fix: make get_or_create routes more resilient (#175)

* fix: make get_or_create routes more resilien

* fix: standardize retry

src/crud/collection.py

@@ -1,10 +1,11 @@
 from logging import getLogger
 
 from sqlalchemy import select
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from src import models
-from src.exceptions import ResourceNotFoundException
+from src.exceptions import ConflictException, ResourceNotFoundException
 
 logger = getLogger(__name__)
 
@@ -51,15 +52,27 @@ async def get_or_create_collection(
     workspace_name: str,
     collection_name: str,
     peer_name: str | None = None,
+    *,
+    _retry: bool = False,
 ) -> models.Collection:
     try:
         return await get_collection(db, workspace_name, collection_name, peer_name)
     except ResourceNotFoundException:
-        honcho_collection = models.Collection(
-            workspace_name=workspace_name,
-            peer_name=peer_name,
-            name=collection_name,
-        )
-        db.add(honcho_collection)
-        await db.commit()
-        return honcho_collection
+        try:
+            honcho_collection = models.Collection(
+                workspace_name=workspace_name,
+                peer_name=peer_name,
+                name=collection_name,
+            )
+            db.add(honcho_collection)
+            await db.commit()
+            return honcho_collection
+        except IntegrityError:
+            await db.rollback()
+            if _retry:
+                raise ConflictException(
+                    f"Unable to create or get collection: {collection_name}"
+                ) from None
+            return await get_or_create_collection(
+                db, workspace_name, collection_name, peer_name, _retry=True
+            )

src/crud/peer.py

@@ -2,10 +2,11 @@ from logging import getLogger
 from typing import Any
 
 from sqlalchemy import Select, select
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from src import models, schemas
-from src.exceptions import ResourceNotFoundException
+from src.exceptions import ConflictException, ResourceNotFoundException
 from src.utils.filter import apply_filter
 
 logger = getLogger(__name__)
@@ -15,6 +16,8 @@ async def get_or_create_peers(
     db: AsyncSession,
     workspace_name: str,
     peers: list[schemas.PeerCreate],
+    *,
+    _retry: bool = False,
 ) -> list[models.Peer]:
     """
     Get an existing list of peers or create new peers if they don't exist.
@@ -24,9 +27,13 @@ async def get_or_create_peers(
         db: Database session
         workspace_name: Name of the workspace
         peers: List of peer creation schemas
+        _retry: Whether to retry the operation
 
     Returns:
         List of peers if found or created
+
+    Raises:
+        ConflictException: If we fail to get or create the peers
     """
     peer_names = [p.name for p in peers]
     stmt = (
@@ -65,12 +72,18 @@ async def get_or_create_peers(
         )
         for p in peers_to_create
     ]
-    db.add_all(new_peers)
-
-    await db.commit()
-
-    # Return combined list of existing and new peers
-    return existing_peers + new_peers
+    try:
+        db.add_all(new_peers)
+        await db.commit()
+        # Return combined list of existing and new peers
+        return existing_peers + new_peers
+    except IntegrityError:
+        await db.rollback()
+        if _retry:
+            raise ConflictException(
+                f"Unable to create or get peers: {peer_names}"
+            ) from None
+        return await get_or_create_peers(db, workspace_name, peers, _retry=True)
 
 
 async def get_peer(

src/crud/session.py

@@ -4,12 +4,13 @@ from typing import Any
 from nanoid import generate as generate_nanoid
 from sqlalchemy import Select, case, cast, func, insert, select, update
 from sqlalchemy.dialects.postgresql import insert as pg_insert
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.types import BigInteger
 
 from src import models, schemas
 from src.config import settings
-from src.exceptions import ResourceNotFoundException
+from src.exceptions import ConflictException, ResourceNotFoundException
 from src.utils.filter import apply_filter
 
 from .peer import get_or_create_peers, get_peer
@@ -38,6 +39,8 @@ async def get_or_create_session(
     db: AsyncSession,
     session: schemas.SessionCreate,
     workspace_name: str,
+    *,
+    _retry: bool = False,
 ) -> models.Session:
     """
     Get or create a session in a workspace with specified peers.
@@ -48,12 +51,13 @@ async def get_or_create_session(
         session: Session creation schema
         workspace_name: Name of the workspace
         peer_names: List of peer names to add to the session
-
+        _retry: Whether to retry the operation
     Returns:
         The created session
 
     Raises:
         ResourceNotFoundException: If the session does not exist and create is false
+        ConflictException: If we fail to get or create the session
     """
 
     stmt = (
@@ -89,9 +93,20 @@ async def get_or_create_session(
             h_metadata=session.metadata or {},
             configuration=session.configuration or {},
         )
-        db.add(honcho_session)
-        # Flush to ensure session exists in DB before adding peers
-        await db.flush()
+        try:
+            db.add(honcho_session)
+            # Flush to ensure session exists in DB before adding peers
+            await db.flush()
+        except IntegrityError:
+            await db.rollback()
+            logger.debug(
+                f"Race condition detected for session: {session.name}, retrying get"
+            )
+            if _retry:
+                raise ConflictException(
+                    f"Unable to create or get session: {session.name}"
+                ) from None
+            return await get_or_create_session(db, session, workspace_name, _retry=True)
     else:
         # Update existing session with metadata and feature flags if provided
         if session.metadata is not None:

src/crud/workspace.py

@@ -2,16 +2,21 @@ from logging import getLogger
 from typing import Any
 
 from sqlalchemy import Select, select
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from src import models, schemas
+from src.exceptions import ConflictException
 from src.utils.filter import apply_filter
 
 logger = getLogger(__name__)
 
 
 async def get_or_create_workspace(
-    db: AsyncSession, workspace: schemas.WorkspaceCreate
+    db: AsyncSession,
+    workspace: schemas.WorkspaceCreate,
+    *,
+    _retry: bool = False,
 ) -> models.Workspace:
     """
     Get an existing workspace or create a new one if it doesn't exist.
@@ -24,7 +29,7 @@ async def get_or_create_workspace(
         The workspace if found or created
 
     Raises:
-        ConflictException: If there's an integrity error when creating the workspace
+        ConflictException: If we fail to get or create the workspace
     """
     # Try to get the existing workspace
     stmt = select(models.Workspace).where(models.Workspace.name == workspace.name)
@@ -42,10 +47,18 @@ async def get_or_create_workspace(
         h_metadata=workspace.metadata,
         configuration=workspace.configuration,
     )
-    db.add(honcho_workspace)
-    await db.commit()
-    logger.info(f"Workspace created successfully: {workspace.name}")
-    return honcho_workspace
+    try:
+        db.add(honcho_workspace)
+        await db.commit()
+        logger.info(f"Workspace created successfully: {workspace.name}")
+        return honcho_workspace
+    except IntegrityError:
+        await db.rollback()
+        if _retry:
+            raise ConflictException(
+                f"Unable to create or get workspace: {workspace.name}"
+            ) from None
+        return await get_or_create_workspace(db, workspace, _retry=True)
 
 
 async def get_all_workspaces(