session scope

src/transcription/services/base.py

@@ -1,5 +1,6 @@
 import asyncio
 from abc import ABC
+from contextlib import asynccontextmanager
 
 from sqlalchemy.ext.asyncio import async_sessionmaker
 from sqlmodel.ext.asyncio.session import AsyncSession
@@ -24,3 +25,14 @@ class ServiceBase(ABC):
         self.settings = get_settings()
         self.session_factory = session_factory or get_session_factory()
         self.queue = queue or asyncio.Queue()
+
+    @asynccontextmanager
+    async def _session_scope(self, session: AsyncSession | None):
+        """Provide a transactional scope around a series of operations."""
+        if session is not None:
+            # Reuse the provided session if one is passed in
+            yield session
+        else:
+            # Otherwise, create a new session for this scope
+            async with self.session_factory() as new_session:
+                yield new_session

src/transcription/services/jobs.py

@@ -3,6 +3,7 @@ from uuid import UUID
 
 from sqlalchemy.orm import selectinload
 from sqlmodel import select
+from sqlmodel.ext.asyncio.session import AsyncSession
 
 from ..models import Job
 from ..models import JobStatus
@@ -12,22 +13,22 @@ from .base import ServiceBase
 class JobService(ServiceBase):
     """Thin service class for managing jobs in the database."""
 
-    async def create_job(self, job: Job) -> Job:
+    async def create_job(self, job: Job, session: AsyncSession | None = None) -> Job:
         """Create a new job in the database."""
-        async with self.session_factory() as session:
+        async with self._session_scope(session) as _session:
-            session.add(job)
+            _session.add(job)
-            await session.commit()
+            await _session.commit()
-            await session.refresh(job)
+            await _session.refresh(job)
             return job
 
-    async def read_job(self, job_id: UUID) -> Job:
+    async def read_job(self, job_id: UUID, session: AsyncSession | None = None) -> Job:
         """Read an existing job from the database.
 
         The selectinload option is used to eagerly load the related document for the job, which makes the full Document
         model object available in the return Job object.
         """
-        async with self.session_factory() as session:
+        async with self._session_scope(session) as _session:
-            job = await session.get(
+            job = await _session.get(
                 Job,
                 job_id,
                 # Makes the full Document model object available in the return Job object
@@ -37,19 +38,19 @@ class JobService(ServiceBase):
             raise ValueError(f"Job with id {job_id} not found")
         return job
 
-    async def update_job(self, job: Job) -> Job:
+    async def update_job(self, job: Job, session: AsyncSession | None = None) -> Job:
         """Update an existing job in the database."""
-        async with self.session_factory() as session:
+        async with self._session_scope(session) as _session:
-            await session.merge(job)
+            await _session.merge(job)
-            await session.commit()
+            await _session.commit()
-            await session.refresh(job)
+            await _session.refresh(job)
-            return job
+            return Job.model_copy(job)
 
-    async def delete_job(self, job: Job) -> None:
+    async def delete_job(self, job: Job, session: AsyncSession | None = None) -> None:
         """Delete a job from the database."""
-        async with self.session_factory() as session:
+        async with self._session_scope(session) as _session:
-            await session.delete(job)
+            await _session.delete(job)
-            await session.commit()
+            await _session.commit()
 
     async def query_jobs(
         self,
@@ -58,7 +59,7 @@ class JobService(ServiceBase):
         filename: str | None = None,
     ) -> Sequence[Job]:
         """Query jobs from the database based on provided filters."""
-        async with self.session_factory() as session:
+        async with self._session_scope(None) as session:
             query = select(Job)
             if status is not None:
                 query = query.where(Job.status == status)
@@ -66,18 +67,23 @@ class JobService(ServiceBase):
                 query = query.where(Job.document.filename == filename)
             return (await session.exec(query)).all()
 
-    async def list_jobs(self) -> Sequence[Job]:
+    async def list_jobs(self, session: AsyncSession | None = None) -> Sequence[Job]:
         """List all jobs in the database."""
-        async with self.session_factory() as session:
+        async with self._session_scope(session) as _session:
-            return (await session.exec(select(Job))).all()
+            return (await _session.exec(select(Job))).all()
 
-    async def mark_job_status(self, job_id: UUID, status: JobStatus) -> Job:
+    async def mark_job_status(
+        self,
+        job_id: UUID,
+        status: JobStatus,
+        session: AsyncSession | None = None,
+    ) -> Job:
         """Mark a job with a new status."""
-        async with self.session_factory() as session:
+        async with self._session_scope(session) as _session:
-            job = await session.get(Job, job_id)
+            job = await _session.get(Job, job_id)
             if job is None:
                 raise ValueError(f"Job with id {job_id} not found")
             job.status = status
-            await session.commit()
+            await _session.commit()
-            await session.refresh(job)
+            await _session.refresh(job)
             return job