Proactively look for duplicate blocks instead of relying on `Integrit…

…yError`. (#7140)
PrefectHQ · Oct 12, 2022 · df540ca · df540ca
1 parent ab53108
commit df540ca
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 21 deletions.
diff --git a/src/prefect/orion/api/block_documents.py b/src/prefect/orion/api/block_documents.py
@@ -4,7 +4,6 @@
 from typing import List, Optional
 from uuid import UUID
 
-import sqlalchemy as sa
 from fastapi import Body, Depends, HTTPException, Path, Query, status
 
 from prefect.orion import models, schemas
@@ -25,16 +24,23 @@ async def create_block_document(
     Create a new block document.
     """
     async with db.session_context(begin_transaction=True) as session:
-        try:
-            new_block_document = await models.block_documents.create_block_document(
-                session=session, block_document=block_document
+        if block_document.name is not None:
+            exists = (
+                await models.block_documents.block_document_with_unique_values_exists(
+                    session=session,
+                    block_type_id=block_document.block_type_id,
+                    name=block_document.name,
+                )
             )
-        except sa.exc.IntegrityError:
-            raise HTTPException(
-                status.HTTP_409_CONFLICT,
-                detail="Block already exists",
-            )
-        return new_block_document
+            if exists:
+                raise HTTPException(
+                    status.HTTP_409_CONFLICT,
+                    detail="Block already exists",
+                )
+
+        return await models.block_documents.create_block_document(
+            session=session, block_document=block_document
+        )
 
 
 @router.post("/filter")

diff --git a/src/prefect/orion/models/block_documents.py b/src/prefect/orion/models/block_documents.py
@@ -7,6 +7,7 @@
 from uuid import UUID, uuid4
 
 import sqlalchemy as sa
+from sqlalchemy.ext.asyncio import AsyncSession
 
 import prefect.orion.models as models
 from prefect.orion import schemas
@@ -23,7 +24,7 @@
 
 @inject_db
 async def create_block_document(
-    session: sa.orm.Session,
+    session: AsyncSession,
     block_document: schemas.actions.BlockDocumentCreate,
     db: OrionDBInterface,
 ):
@@ -63,14 +64,28 @@ async def create_block_document(
             ),
         )
 
-    # reload the block document in order to load the associated block schema relationship
+    # reload the block document in order to load the associated block schema
+    # relationship
     return await read_block_document_by_id(
         session=session,
         block_document_id=orm_block.id,
         include_secrets=False,
     )
 
 
+@inject_db
+async def block_document_with_unique_values_exists(
+    session: AsyncSession, block_type_id: UUID, name: str, db: OrionDBInterface
+) -> bool:
+    result = await session.execute(
+        sa.select(sa.exists(db.BlockDocument)).where(
+            db.BlockDocument.block_type_id == block_type_id,
+            db.BlockDocument.name == name,
+        )
+    )
+    return bool(result.scalar_one_or_none())
+
+
 def _separate_block_references_from_data(
     block_document_data: Dict,
 ) -> Tuple[Dict, List[Tuple[str, UUID]]]:
@@ -109,7 +124,7 @@ def _separate_block_references_from_data(
 
 @inject_db
 async def read_block_document_by_id(
-    session: sa.orm.Session,
+    session: AsyncSession,
     block_document_id: UUID,
     db: OrionDBInterface,
     include_secrets: bool = False,
@@ -130,7 +145,7 @@ async def read_block_document_by_id(
 
 
 async def _construct_full_block_document(
-    session: sa.orm.Session,
+    session: AsyncSession,
     block_documents_with_references: List[
         Tuple[ORMBlockDocument, Optional[str], Optional[UUID]]
     ],
@@ -176,7 +191,9 @@ async def _construct_full_block_document(
                     "name": block_document.name,
                     "block_type": block_document.block_type,
                     "is_anonymous": block_document.is_anonymous,
-                    "block_document_references": full_child_block_document.block_document_references,
+                    "block_document_references": (
+                        full_child_block_document.block_document_references
+                    ),
                 }
             }
 
@@ -211,7 +228,7 @@ async def _find_parent_block_document(
 
 @inject_db
 async def read_block_document_by_name(
-    session: sa.orm.Session,
+    session: AsyncSession,
     name: str,
     block_type_slug: str,
     db: OrionDBInterface,
@@ -239,7 +256,7 @@ async def read_block_document_by_name(
 
 @inject_db
 async def read_block_documents(
-    session: sa.orm.Session,
+    session: AsyncSession,
     db: OrionDBInterface,
     block_document_filter: Optional[schemas.filters.BlockDocumentFilter] = None,
     block_type_filter: Optional[schemas.filters.BlockTypeFilter] = None,
@@ -390,7 +407,7 @@ async def read_block_documents(
 
 @inject_db
 async def delete_block_document(
-    session: sa.orm.Session,
+    session: AsyncSession,
     block_document_id: UUID,
     db: OrionDBInterface,
 ) -> bool:
@@ -402,7 +419,7 @@ async def delete_block_document(
 
 @inject_db
 async def update_block_document(
-    session: sa.orm.Session,
+    session: AsyncSession,
     block_document_id: UUID,
     block_document: schemas.actions.BlockDocumentUpdate,
     db: OrionDBInterface,
@@ -505,7 +522,7 @@ def _find_block_document_reference(
 
 @inject_db
 async def create_block_document_reference(
-    session: sa.orm.Session,
+    session: AsyncSession,
     block_document_reference: schemas.actions.BlockDocumentReferenceCreate,
     db: OrionDBInterface,
 ):
@@ -531,7 +548,7 @@ async def create_block_document_reference(
 
 @inject_db
 async def delete_block_document_reference(
-    session: sa.orm.Session,
+    session: AsyncSession,
     block_document_reference_id: UUID,
     db: OrionDBInterface,
 ):