h

API/folder.py

@@ -2028,39 +2028,40 @@ class API_folder_store:
         pipe: Optional[str] = None,
     ) -> int:
         """Insert a new worker entry into the database."""
-        try:
-            cursor = self.connection.cursor()
-            cursor.execute(
-                """
-                INSERT INTO worker (worker_id, worker_type, pipe, status, title, description, total_steps)
-                VALUES (?, ?, ?, ?, ?, ?, ?)
-            """,
-                (
-                    worker_id,
-                    worker_type,
-                    pipe,
-                    "running",
-                    title,
-                    description,
-                    total_steps
-                ),
-            )
-            worker_rowid = cursor.lastrowid or 0
-            
-            # Prune occasionally (1 in 50 chance) or just run it to keep it clean
-            # Running it every time might be overkill, but let's do a light version
-            cursor.execute(
-                "DELETE FROM worker WHERE status != 'running' AND id < (SELECT MAX(id) - ? FROM worker)",
-                (MAX_FINISHED_WORKERS * 2,)
-            )
+        with self._db_lock:
+            try:
+                cursor = self.connection.cursor()
+                cursor.execute(
+                    """
+                    INSERT INTO worker (worker_id, worker_type, pipe, status, title, description, total_steps)
+                    VALUES (?, ?, ?, ?, ?, ?, ?)
+                """,
+                    (
+                        worker_id,
+                        worker_type,
+                        pipe,
+                        "running",
+                        title,
+                        description,
+                        total_steps
+                    ),
+                )
+                worker_rowid = cursor.lastrowid or 0
+                
+                # Prune occasionally (1 in 50 chance) or just run it to keep it clean
+                # Running it every time might be overkill, but let's do a light version
+                cursor.execute(
+                    "DELETE FROM worker WHERE status != 'running' AND id < (SELECT MAX(id) - ? FROM worker)",
+                    (MAX_FINISHED_WORKERS * 2,)
+                )
 
-            self.connection.commit()
-            return worker_rowid
-        except sqlite3.IntegrityError:
-            return self.update_worker_status(worker_id, "running")
-        except Exception as e:
-            logger.error(f"Error inserting worker: {e}", exc_info=True)
-            return 0
+                self.connection.commit()
+                return worker_rowid
+            except sqlite3.IntegrityError:
+                return self.update_worker_status(worker_id, "running")
+            except Exception as e:
+                logger.error(f"Error inserting worker: {e}", exc_info=True)
+                return 0
 
     def update_worker(self, worker_id: str, **kwargs) -> bool:
         """Update worker entry with given fields."""
@@ -2107,40 +2108,41 @@ class API_folder_store:
 
     def update_worker_status(self, worker_id: str, status: str) -> int:
         """Update worker status and return its database ID."""
-        try:
-            cursor = self.connection.cursor()
+        with self._db_lock:
+            try:
+                cursor = self.connection.cursor()
 
-            if status in ("completed", "error"):
-                cursor.execute(
-                    """
-                    UPDATE worker
-                    SET status = ?, completed_at = CURRENT_TIMESTAMP, last_updated = CURRENT_TIMESTAMP
-                    WHERE worker_id = ?
-                """,
-                    (status,
-                     worker_id),
-                )
-            else:
-                cursor.execute(
-                    """
-                    UPDATE worker
-                    SET status = ?, last_updated = CURRENT_TIMESTAMP
-                    WHERE worker_id = ?
-                """,
-                    (status,
-                     worker_id),
-                )
+                if status in ("completed", "error"):
+                    cursor.execute(
+                        """
+                        UPDATE worker
+                        SET status = ?, completed_at = CURRENT_TIMESTAMP, last_updated = CURRENT_TIMESTAMP
+                        WHERE worker_id = ?
+                    """,
+                        (status,
+                         worker_id),
+                    )
+                else:
+                    cursor.execute(
+                        """
+                        UPDATE worker
+                        SET status = ?, last_updated = CURRENT_TIMESTAMP
+                        WHERE worker_id = ?
+                    """,
+                        (status,
+                         worker_id),
+                    )
 
-            self.connection.commit()
+                self.connection.commit()
 
-            cursor.execute("SELECT id FROM worker WHERE worker_id = ?",
-                           (worker_id,
-                            ))
-            row = cursor.fetchone()
-            return row[0] if row else 0
-        except Exception as e:
-            logger.error(f"Error updating worker status: {e}", exc_info=True)
-            return 0
+                cursor.execute("SELECT id FROM worker WHERE worker_id = ?",
+                               (worker_id,
+                                ))
+                row = cursor.fetchone()
+                return row[0] if row else 0
+            except Exception as e:
+                logger.error(f"Error updating worker status: {e}", exc_info=True)
+                return 0
 
     def get_worker(self, worker_id: str) -> Optional[Dict[str, Any]]:
         """Retrieve a worker entry by ID."""
@@ -2484,10 +2486,15 @@ class API_folder_store:
             logger.error(f"Error closing database: {e}", exc_info=True)
 
     def __enter__(self):
+        if not self.connection:
+            self._init_db()
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
+        try:
+            self.close()
+        finally:
+            pass
 
 
 # ============================================================================