chore: something changed, commit before reorg

cpv3/modules/project_workspaces/repository.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import uuid
+
+from sqlalchemy import select, update
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from cpv3.db.base import utcnow
+from cpv3.modules.project_workspaces.models import ProjectWorkspace
+
+
+class WorkspaceRevisionConflictError(RuntimeError):
+    """Raised when the optimistic workspace revision check fails."""
+
+
+class ProjectWorkspaceRepository:
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def get_by_project_id(self, project_id: uuid.UUID) -> ProjectWorkspace | None:
+        result = await self._session.execute(
+            select(ProjectWorkspace).where(ProjectWorkspace.project_id == project_id)
+        )
+        return result.scalar_one_or_none()
+
+    async def create(self, *, project_id: uuid.UUID, state: dict) -> ProjectWorkspace:
+        workspace = ProjectWorkspace(project_id=project_id, revision=0, state=state)
+        self._session.add(workspace)
+        await self._session.commit()
+        await self._session.refresh(workspace)
+        return workspace
+
+    async def get_or_create(self, *, project_id: uuid.UUID, state: dict) -> ProjectWorkspace:
+        workspace = await self.get_by_project_id(project_id)
+        if workspace is not None:
+            return workspace
+
+        try:
+            return await self.create(project_id=project_id, state=state)
+        except IntegrityError:
+            await self._session.rollback()
+            workspace = await self.get_by_project_id(project_id)
+            if workspace is None:
+                raise
+            return workspace
+
+    async def update_state(
+        self,
+        *,
+        project_id: uuid.UUID,
+        expected_revision: int,
+        state: dict,
+    ) -> ProjectWorkspace:
+        stmt = (
+            update(ProjectWorkspace)
+            .where(ProjectWorkspace.project_id == project_id)
+            .where(ProjectWorkspace.revision == expected_revision)
+            .values(
+                state=state,
+                revision=expected_revision + 1,
+                updated_at=utcnow(),
+            )
+        )
+        result = await self._session.execute(stmt)
+        if result.rowcount != 1:
+            await self._session.rollback()
+            raise WorkspaceRevisionConflictError
+
+        await self._session.commit()
+        workspace = await self.get_by_project_id(project_id)
+        if workspace is None:
+            raise RuntimeError("Workspace disappeared after update")
+        return workspace