Fix connection with non sqlite db

src/plombery/database/base.py

@@ -10,11 +10,14 @@
 def json_serializer(*args, **kwargs) -> str:
     return json.dumps(*args, default=jsonable_encoder, **kwargs)
 
+connect_args = {}
+if settings.database_url[:6] == "sqlite":
+    connect_args["check_same_thread"] = False
 
 engine = create_engine(
     settings.database_url,
     json_serializer=json_serializer,
-    connect_args={"check_same_thread": False},
+    connect_args=connect_args,
 )
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 

src/plombery/database/models.py

@@ -81,17 +81,16 @@ class PipelineRun(Base):
 
 
 def _mark_cancelled_runs():
-    db = SessionLocal()
-
-    db.query(PipelineRun).filter(
-        PipelineRun.status == PipelineRunStatus.RUNNING.value
-    ).update(
-        dict(
-            status=PipelineRunStatus.CANCELLED.value,
+    with SessionLocal() as db:
+        db.query(PipelineRun).filter(
+            PipelineRun.status == PipelineRunStatus.RUNNING.value
+        ).update(
+            dict(
+                status=PipelineRunStatus.CANCELLED.value,
+            )
         )
-    )
 
-    db.commit()
+        db.commit()
 
 
 _mark_cancelled_runs()

src/plombery/database/repository.py

@@ -9,82 +9,79 @@
 
 
 def create_pipeline_run(data: PipelineRunCreate):
-    db = SessionLocal()
-    db.expire_on_commit = False
+    with SessionLocal() as db:
+        db.expire_on_commit = False
 
-    created_model = models.PipelineRun(**data.model_dump())
-    db.add(created_model)
-    db.commit()
-    db.refresh(created_model)
+        created_model = models.PipelineRun(**data.model_dump())
+        db.add(created_model)
+        db.commit()
+        db.refresh(created_model)
     return created_model
 
 
 def update_pipeline_run(
     pipeline_run: models.PipelineRun, end_time: datetime, status: PipelineRunStatus
 ):
-    db = SessionLocal()
-
     pipeline_run.duration = (end_time - pipeline_run.start_time).total_seconds() * 1000
     pipeline_run.status = status.value
-
-    db.query(models.PipelineRun).filter(
-        models.PipelineRun.id == pipeline_run.id
-    ).update(
-        dict(
-            duration=pipeline_run.duration,
-            status=pipeline_run.status,
-            tasks_run=pipeline_run.tasks_run,
+    with SessionLocal() as db:
+
+        db.query(models.PipelineRun).filter(
+            models.PipelineRun.id == pipeline_run.id
+        ).update(
+            dict(
+                duration=pipeline_run.duration,
+                status=pipeline_run.status,
+                tasks_run=pipeline_run.tasks_run,
+            )
         )
-    )
-
-    db.commit()
+        db.commit()
 
 
 def list_pipeline_runs(
     pipeline_id: Optional[str] = None, trigger_id: Optional[str] = None
 ):
-    db = SessionLocal()
-    db.expire_on_commit = False
-
     filters = []
-
     if pipeline_id:
         filters.append(models.PipelineRun.pipeline_id == pipeline_id)
     if trigger_id:
         filters.append(models.PipelineRun.trigger_id == trigger_id)
 
-    pipeline_runs: List[models.PipelineRun] = (
-        db.query(models.PipelineRun)
-        .filter(*filters)
-        .order_by(models.PipelineRun.id.desc())
-        .limit(30)
-        .all()
-    )
+    with SessionLocal() as db:
+        db.expire_on_commit = False
+
+        pipeline_runs: List[models.PipelineRun] = (
+            db.query(models.PipelineRun)
+            .filter(*filters)
+            .order_by(models.PipelineRun.id.desc())
+            .limit(30)
+            .all()
+        )
 
     return pipeline_runs
 
 
 def get_pipeline_run(pipeline_run_id: int) -> Optional[models.PipelineRun]:
-    db = SessionLocal()
+    with SessionLocal() as db:
 
-    pipeline_run: Optional[models.PipelineRun] = db.query(models.PipelineRun).get(
-        pipeline_run_id
-    )
+        pipeline_run: Optional[models.PipelineRun] = db.query(models.PipelineRun).get(
+            pipeline_run_id
+        )
 
     return pipeline_run
 
 
 def get_latest_pipeline_run(pipeline_id, trigger_id):
-    db = SessionLocal()
-
-    pipeline_run: models.PipelineRun = (
-        db.query(models.PipelineRun)
-        .filter(
-            models.PipelineRun.pipeline_id == pipeline_id,
-            models.PipelineRun.trigger_id == trigger_id,
+    with SessionLocal() as db:
+
+        pipeline_run: models.PipelineRun = (
+            db.query(models.PipelineRun)
+            .filter(
+                models.PipelineRun.pipeline_id == pipeline_id,
+                models.PipelineRun.trigger_id == trigger_id,
+            )
+            .order_by(models.PipelineRun.id.desc())
+            .first()
         )
-        .order_by(models.PipelineRun.id.desc())
-        .first()
-    )
 
     return pipeline_run