diff --git a/src/crud/collection.py b/src/crud/collection.py index e44116c3..49fe9d88 100644 --- a/src/crud/collection.py +++ b/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 + ) diff --git a/src/crud/peer.py b/src/crud/peer.py index 6a609ff7..1e14aed7 100644 --- a/src/crud/peer.py +++ b/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( diff --git a/src/crud/session.py b/src/crud/session.py index 327d74bd..05830c6b 100644 --- a/src/crud/session.py +++ b/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: diff --git a/src/crud/workspace.py b/src/crud/workspace.py index 9ca53380..f4f4287e 100644 --- a/src/crud/workspace.py +++ b/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(