diff --git a/src/transcription/services/base.py b/src/transcription/services/base.py index 1d6522a..307dff4 100644 --- a/src/transcription/services/base.py +++ b/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 diff --git a/src/transcription/services/jobs.py b/src/transcription/services/jobs.py index 5963cd5..aeb2b31 100644 --- a/src/transcription/services/jobs.py +++ b/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: - session.add(job) - await session.commit() - await session.refresh(job) + async with self._session_scope(session) as _session: + _session.add(job) + await _session.commit() + 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: - job = await session.get( + async with self._session_scope(session) as _session: + 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: - await session.merge(job) - await session.commit() - await session.refresh(job) - return job + async with self._session_scope(session) as _session: + await _session.merge(job) + await _session.commit() + await _session.refresh(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: - await session.delete(job) - await session.commit() + async with self._session_scope(session) as _session: + await _session.delete(job) + 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: - return (await session.exec(select(Job))).all() + async with self._session_scope(session) as _session: + 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: - job = await session.get(Job, job_id) + async with self._session_scope(session) as _session: + 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.refresh(job) + await _session.commit() + await _session.refresh(job) return job