Add RenderClient object to DB for client tracking

lib/render_queue.py

@@ -4,7 +4,7 @@ from datetime import datetime
 
 import psutil
 import requests
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, Column, String, Integer
 from sqlalchemy.orm import sessionmaker
 
 from .render_workers.render_worker import RenderStatus
@@ -19,6 +19,27 @@ class JobNotFoundError(Exception):
         self.job_id = job_id
 
 
+class RenderClient(Base):
+    __tablename__ = 'render_clients'
+    id = Column(Integer, primary_key=True)
+    hostname = Column(String)
+
+    def __init__(self, hostname):
+        self.hostname = hostname
+
+    def is_available(self, timeout=3):
+        try:
+            response = requests.get(f"http://{self.hostname}:8080/api/status", timeout=timeout)
+            if response.ok:
+                return True
+        except requests.ConnectionError as e:
+            pass
+        return False
+
+    def __repr__(self):
+        return "client stuff"
+
+
 class RenderQueue:
     engine = create_engine('sqlite:///database.db')
     Base.metadata.create_all(engine)
@@ -26,9 +47,8 @@ class RenderQueue:
     session = Session()
     ScheduledJob.register_user_events()
     job_queue = []
-    render_clients = []
     maximum_renderer_instances = {'blender': 2, 'aerender': 1, 'ffmpeg': 4}
-    host_name = None
+    hostname = None
     port = 8080
     client_mode = False
     server_hostname = None
@@ -41,9 +61,9 @@ class RenderQueue:
     @classmethod
     def add_to_render_queue(cls, render_job, force_start=False, client=None):
 
-        if not client or render_job.client == cls.host_name:
+        if not client or render_job.client == cls.hostname:
             logger.debug('Adding priority {} job to render queue: {}'.format(render_job.priority, render_job.worker_object))
-            render_job.client = cls.host_name
+            render_job.client = cls.hostname
             cls.job_queue.append(render_job)
             if force_start:
                 cls.start_job(render_job)
@@ -165,44 +185,56 @@ class RenderQueue:
                 "memory_available": psutil.virtual_memory().available,
                 "memory_percent": psutil.virtual_memory().percent,
                 "job_counts": cls.job_counts(),
-                "host_name": cls.host_name
+                "host_name": cls.hostname
                 }
 
     @classmethod
-    def register_client(cls, hostname):
+    def render_clients(cls):
+        all_clients = cls.session.query(RenderClient).all()
+        if not all_clients:
+            cls.session.add(RenderClient(hostname=cls.hostname))
+            cls.save_state()
+            all_clients = cls.session.query(RenderClient).all()
+        return all_clients
 
+    @classmethod
+    def client_with_hostname(cls, hostname):
+        return cls.session.query(RenderClient).filter(RenderClient.hostname == hostname).first()
+
+    @classmethod
+    def register_client(cls, hostname):
+        new_client = None
         err_msg = None
 
-        if hostname == cls.host_name:
+        if hostname == cls.hostname:
             err_msg = "Cannot register same hostname as server"
-        elif hostname in cls.render_clients:
+        elif cls.client_with_hostname(hostname):
             err_msg = f"Client '{hostname}' already registered"
         else:
-            try:
-                response = requests.get(f"http://{hostname}:8080/api/status", timeout=3)
-                if response.ok:
-                    cls.render_clients.append(hostname)
-                    logger.info(f"Client '{hostname}' successfully registered")
-                    cls.save_state()
-                else:
-                    err_msg = f'Response from server not ok: {response.text}'
-            except requests.ConnectionError as e:
+            new_client = RenderClient(hostname=hostname)
+            if not new_client.is_available():
+                cls.session.add(new_client)
+                logger.info(f"Client '{hostname}' successfully registered")
+                cls.save_state()
+            else:
                 err_msg = f"Cannot connect to client at hostname: {hostname}"
 
         if err_msg:
             logger.warning(err_msg)
             return err_msg, 400
         else:
-            return 'success'
+            return new_client.hostname
 
     @classmethod
     def unregister_client(cls, hostname):
         success = False
-        if hostname in cls.render_clients and hostname != cls.host_name:
-            cls.render_clients.remove(hostname)
+        client = cls.client_with_hostname(hostname)
+        if client and hostname != cls.hostname:
+            cls.session.delete(client)
+            cls.save_state()
             logger.info(f"Client '{hostname}' successfully unregistered")
             success = True
-        return success
+        return str(success)
 
     @staticmethod
     def is_client_available(client_hostname, timeout=3):