refact(db): make cursor arg optional for finish_run

bento-platform · Apr 11, 2024 · d15d755 · d15d755
1 parent 15222fb
commit d15d755
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/bento_wes/backends/_wes_backend.py b/bento_wes/backends/_wes_backend.py
@@ -283,7 +283,7 @@ def _finish_run_and_clean_up(self, run: Run, state: str) -> None:
 
         # Finish run ----------------------------------------------------------
 
-        self.db.finish_run(self.db.cursor(), self.event_bus, run, state)
+        self.db.finish_run(self.event_bus, run, state)
 
         # Clean up ------------------------------------------------------------
 

diff --git a/bento_wes/db.py b/bento_wes/db.py
@@ -104,10 +104,10 @@ def init(self):
 
     def finish_run(
         self,
-        c: sqlite3.Cursor,
         event_bus: EventBus,
         run: Run,
         state: str,
+        cursor: sqlite3.Cursor | None = None,
         logger: logging.Logger | None = None,
     ) -> None:
         """
@@ -121,6 +121,8 @@ def finish_run(
         :return:
         """
 
+        c: sqlite3.Cursor = cursor or self.cursor()
+
         run_id = run.run_id
         end_time = iso_now()
 
@@ -176,7 +178,7 @@ def update_stuck_runs(self):
 
             logger.info(
                 f"Found stuck run: {run.run_id} at state {run.state}. Setting state to {states.STATE_SYSTEM_ERROR}")
-            self.finish_run(c, event_bus, run, states.STATE_SYSTEM_ERROR)
+            self.finish_run(event_bus, run, states.STATE_SYSTEM_ERROR, cursor=c)
 
         self.commit()
 

diff --git a/bento_wes/runner.py b/bento_wes/runner.py
@@ -79,7 +79,7 @@ def run_workflow(self, run_id: uuid.UUID):
     except Exception as e:
         # Intercept any uncaught exceptions and finish with an error state
         logger.error(f"Uncaught exception while obtaining access token: {type(e).__name__} {e}")
-        db.finish_run(c, event_bus, run, states.STATE_SYSTEM_ERROR, logger=logger)
+        db.finish_run(event_bus, run, states.STATE_SYSTEM_ERROR, cursor=c, logger=logger)
         raise e
 
     # Perform the run
@@ -89,5 +89,5 @@ def run_workflow(self, run_id: uuid.UUID):
     except Exception as e:
         # Intercept any uncaught exceptions and finish with an error state
         logger.error(f"Uncaught exception while performing run: {type(e).__name__} {e}")
-        db.finish_run(c, event_bus, run, states.STATE_SYSTEM_ERROR, logger=logger)
+        db.finish_run(event_bus, run, states.STATE_SYSTEM_ERROR, cursor=c, logger=logger)
         raise e