Update worker.py

worker.py

@@ -15,18 +15,93 @@ url: str = os.environ.get("SUPABASE_URL")
 key: str = os.environ.get("SUPABASE_KEY")
 supabase: Client = create_client(url, key)
 
+fire_url = os.environ.get("FIRE_URL")
+fire_bucket = os.environ.get("FIRE_BUCKET")
+cred = credentials.Certificate('firekey.json')
+firebase_admin.initialize_app(cred, {'storageBucket': fire_bucket})
+ref = db.reference('/Tasks',url=fire_url)
+bucket = storage.bucket()
+
+IsSuperbase = False
+if time.localtime().tm_mday > 15:
+    IsSuperbase = True
+
+    
+class Task:
+    def __init__(self,tasks):
+        """Initialize the task with the first task in the list of tasks"""
+        self.key = tasks[0][0]
+        self.text = tasks[0][1]['text']
+        self.status = tasks[0][1]['status']
+        self.result = tasks[0][1]['result']
+
+    def __str__(self):
+        return f"Task: {self.text}, Status: {self.status}, Result: {self.result}, Key: {self.key}"
+    
+    def processing(self):
+        ref.child(self.key).set({
+            "text": self.text,
+            "status": 'processing',
+            'result': ''
+        })
+        print('Processing')
+
+    def complete(self, result):
+        ref.child(self.key).set({
+            "text": self.text,
+            "status": 'complete',
+            'result': result
+        })
+        print('Complete')
+
+    def error(self):
+        ref.child(self.key).set({
+            "text": self.text,
+            "status": 'error',
+            'result': ''
+        })
+        print('Error')
+
+    def save_to_storage(self, binary_file_path, key):
+        try:
+            blob = bucket.blob(key+'.obj')
+            #upload but get download url as well
+            blob.upload_from_filename(binary_file_path)
+            blob.make_public()
+            self.complete(blob.public_url)
+            os.remove('task_'+str(self.key)+'.obj')
+        except Exception as e:
+            print(f"Error saving file to storage: {e}")
+            self.error()
+            return None
+    
 def check_queue():
-    try:
-        tasks = supabase.table("Tasks").select("*").eq("status", "pending").execute()
-        assert len(tasks.data) > 0
-        if len(tasks.data) > 0:
-            return tasks.data[0]
-        else:
+    if IsSuperbase:
+        try:
+            tasks = supabase.table("Tasks").select("*").eq("status", "pending").execute()
+            assert len(tasks.data) > 0
+            if len(tasks.data) > 0:
+                return tasks.data[0]
+            else:
+                return None
+        except Exception as e:
+            print(f"Error checking queue: {e}")
+            return None
+        
+    else:
+        try:
+            response = ref.order_by_child('status').equal_to('pending').limit_to_first(1).get()
+            tasks = list(response.items())  # Convert tasks to a list
+            print(tasks)
+            if len(tasks) > 0:
+                task = Task(tasks)
+                print(task)
+                return task
+            else:
+                return None
+        except Exception as e:
+            print(f"Error checking queue: {e}")
             return None
-    except Exception as e:
-        print(f"Error checking queue: {e}")
-        return None
-
 
 def generate_image(text):
     try:
@@ -65,17 +140,31 @@ def worker():
     while True:
         task = check_queue()
         if task:
-            supabase.table("Tasks").update({"status": "processing"}).eq("id", task['id']).execute()
-            print(f"Processing task {task['id']}")
-            img = generate_image(task["text"])
+            if IsSuperbase:
+                task_id = task['id']
+                supabase.table("Tasks").update({"status": "processing"}).eq("id", task_id).execute()
+                print(f"Processing task {task_id}")
+                img = generate_image(task["text"])
+
+            else:
+                task_id = task.key
+                task.processing()
+                img = generate_image(task.text)
+            
             if img:
-                print(f"Image generated for task {task['id']}")
-                create_obj_file(img,task["id"])
-                send_back_to_supabase(task["id"])
-                print(f"Task {task['id']} completed")
+                print(f"Image generated for task {task_id}")
+                create_obj_file(img,task_id)
+                if IsSuperbase:
+                    send_back_to_supabase(task_id)
+                else:
+                    task.save_to_storage('task_'+str(task_id)+'.obj', task_id)
+                    print(f"Task {task_id} completed")
             else:
-                print(f"Error generating image for task {task['id']}")
-                supabase.table("Tasks").update({"status": "error"}).eq("id", task['id']).execute()
+                print(f"Error generating image for task {task_id}")
+                if IsSuperbase:
+                    supabase.table("Tasks").update({"status": "error"}).eq("id", task_id).execute()
+                else:
+                    task.error()
 
         else:
             print("No pending tasks in the queue")