document id validation with dep injection too

Author: raphaellaude

backend/app/core/dependencies.py

@@ -1,5 +1,6 @@
 from fastapi import Depends, HTTPException, status
 from sqlmodel import select, Session, literal
+from app.core.models import DocumentID
 from app.models import Document, DocumentPublic, DistrictrMap
 from app.save_share.models import (
     DocumentShareStatus,
@@ -17,10 +18,20 @@
 logging.basicConfig(level=logging.INFO)
 
 
-def get_document(document_id: str, session: Session = Depends(get_session)) -> Document:
+def parse_document_id(document_id: str) -> DocumentID:
+    try:
+        return DocumentID(document_id=document_id)
+    except ValueError:
+        raise HTTPException(status_code=400, detail="Invalid document ID")
+
+
+def get_document(
+    document_id: DocumentID = Depends(parse_document_id),
+    session: Session = Depends(get_session),
+) -> Document:
     try:
         document = session.exec(
-            select(Document).where(Document.document_id == document_id)
+            select(Document).where(Document.document_id == document_id.value)
         ).one()
     except NoResultFound:
         raise HTTPException(status_code=404, detail="Document not found")
@@ -32,7 +43,8 @@ def get_document(document_id: str, session: Session = Depends(get_session)) -> D
 
 
 def get_protected_document(
-    document_id: str | int, session: Session = Depends(get_session)
+    document_id: DocumentID = Depends(parse_document_id),
+    session: Session = Depends(get_session),
 ) -> Document:
     """
     Always returns a document even if the public_id was used instead of the document_id. This function
@@ -44,21 +56,14 @@ def get_protected_document(
     - UUID document IDs
     - Public IDs (numeric, for public sharing)
     """
-    id_is_public = isinstance(document_id, int) or (
-        isinstance(document_id, str) and document_id.isdigit()
-    )
-
-    if id_is_public:
-        document_id = int(document_id)
-
     stmt = select(Document)
 
-    if id_is_public:
-        stmt = stmt.where(Document.public_id == document_id).where(
+    if document_id.is_public:
+        stmt = stmt.where(Document.public_id == document_id.value).where(
             text("map_metadata->>'draft_status' = 'ready_to_share'")
         )
     else:
-        stmt = stmt.where(Document.document_id == document_id)
+        stmt = stmt.where(Document.document_id == document_id.value)
 
     try:
         document = session.exec(stmt).one()
@@ -73,33 +78,21 @@ def get_protected_document(
 
 def get_document_public(
     session: Session,
-    document_id: str | int,
+    document_id: DocumentID = Depends(parse_document_id),
     user_id: str | None = None,
     shared: bool = False,
     lock_status: DocumentEditStatus | None = None,
 ) -> DocumentPublic:
-    id_is_public = isinstance(document_id, int) or (
-        isinstance(document_id, str) and document_id.isdigit()
-    )
-
-    if id_is_public:
-        document_id = int(document_id)
-
-    if not document_id:
-        raise HTTPException(status_code=404, detail="Document not found")
-    document = get_protected_document(document_id=document_id, session=session)
-    # TODO: Rather than being a separate query, this should be part of the main query
-
     access_type = DocumentShareStatus.read
     # Store if lock_status was explicitly provided
     lock_status_provided = lock_status is not None
 
-    if document.document_id == document_id and not id_is_public:
+    if not document_id.is_public:
         access_type = DocumentShareStatus.edit
         # Only check map lock if no lock_status was explicitly provided
         if not lock_status_provided:
             lock_status = check_map_lock(
-                document.document_id, user_id=user_id, session=session
+                document_id.document_id, user_id=user_id, session=session
             )
 
     # Set default lock_status if not provided and not already set
@@ -108,7 +101,9 @@ def get_document_public(
 
     stmt = select(
         # Obsured document ID
-        literal("anonymous" if id_is_public else document_id).label("document_id"),
+        literal(
+            "anonymous" if document_id.is_public else document_id.document_id
+        ).label("document_id"),
         Document.created_at,
         Document.districtr_map_slug,
         Document.gerrydb_table,
@@ -139,20 +134,20 @@ def get_document_public(
         isouter=True,
     )
 
-    if id_is_public:
-        stmt = stmt.where(Document.public_id == document_id).where(
+    if document_id.is_public:
+        stmt = stmt.where(Document.public_id == document_id.value).where(
             text("map_metadata->>'draft_status' = 'ready_to_share'")
         )
     else:
-        stmt = stmt.where(Document.document_id == document_id)
+        stmt = stmt.where(Document.document_id == document_id.value)
 
     result = session.exec(stmt)
 
     return result.one()
 
 
 def get_districtr_map(
-    document_id: str,
+    document_id: DocumentID = Depends(parse_document_id),
     session: Session = Depends(get_session),
 ) -> DistrictrMap:
     stmt = (
@@ -169,8 +164,8 @@ def get_districtr_map(
         )
         .filter(
             or_(
-                Document.document_id == document_id,
-                MapDocumentToken.document_id == document_id,
+                Document.document_id == document_id.document_id,
+                MapDocumentToken.document_id == document_id.document_id,
             )
         )  # pyright: ignore
     )

backend/app/core/models.py

@@ -1,14 +1,16 @@
 from datetime import datetime
+from pydantic import BaseModel, computed_field, model_validator
+from uuid import UUID
 from sqlmodel import (
     Field,
     SQLModel,
-    UUID,
     TIMESTAMP,
+    UUID as UUIDField,
     text,
 )
 
 
-class UUIDType(UUID):
+class UUIDType(UUIDField):
     def __init__(self, *args, **kwargs):
         kwargs["as_uuid"] = False
         super().__init__(*args, **kwargs)
@@ -32,3 +34,44 @@ class TimeStampMixin(SQLModel):
         nullable=False,
         default=None,
     )
+
+
+class DocumentID(BaseModel):
+    """
+    Represents a document identifier.
+
+    Attributes:
+        document_id (str): The unique identifier of the document.
+        is_public (bool): Indicates whether the document is public.
+        value (int | str): The value of the document identifier.
+
+    Raises:
+        ValueError: If the private document_id is not a valid UUID.
+    """
+
+    document_id: str
+
+    @computed_field
+    @property
+    def is_public(self) -> bool:
+        return isinstance(self.document_id, int) or (
+            isinstance(self.document_id, str) and self.document_id.isdigit()
+        )
+
+    @computed_field
+    @property
+    def value(self) -> int | str:
+        if self.is_public:
+            return int(self.document_id)
+        return self.document_id
+
+    @model_validator(mode="after")
+    def validate_private_uuid(self) -> "DocumentID":
+        if not self.is_public:
+            try:
+                UUID(self.document_id)
+            except ValueError:
+                raise ValueError(
+                    f"Private document_id must be a valid UUID, got: {self.document_id}"
+                )
+        return self

backend/app/main.py

@@ -31,7 +31,9 @@
     get_document_public,
     get_protected_document,
     get_districtr_map,
+    parse_document_id,
 )
+from app.core.models import DocumentID
 from app.core.config import settings
 import app.contiguity.main as contiguity
 import app.cms.main as cms
@@ -155,7 +157,9 @@ async def db_is_alive(session: Session = Depends(get_session)):
 
 @app.post("/api/document/{document_id}/unload", status_code=status.HTTP_200_OK)
 async def unlock_map(
-    document_id: str, user_id: str = Form(...), session: Session = Depends(get_session)
+    document_id: DocumentID = Depends(parse_document_id),
+    user_id: str = Form(...),
+    session: Session = Depends(get_session),
 ):
     """
     unlock map when tab is unloaded
@@ -169,7 +173,7 @@ async def unlock_map(
                 bindparam(key="document_id", type_=UUIDType),
                 bindparam(key="user_id", type_=String),
             ),
-            {"document_id": document_id, "user_id": user_id},
+            {"document_id": document_id.value, "user_id": user_id},
         )
         session.commit()
         return {"status": DocumentEditStatus.unlocked}
@@ -461,7 +465,9 @@ async def reset_map(
     response_model=ColorsSetResult,
 )
 async def update_colors(
-    document_id: str, colors: list[str], session: Session = Depends(get_session)
+    colors: list[str],
+    document_id: DocumentID = Depends(parse_document_id),
+    session: Session = Depends(get_session),
 ):
     districtr_map = session.exec(
         select(DistrictrMap)
@@ -470,7 +476,7 @@ async def update_colors(
             Document.districtr_map_slug == DistrictrMap.districtr_map_slug,  # pyright: ignore
             isouter=True,
         )
-        .where(Document.document_id == document_id)
+        .where(Document.document_id == document_id.value)
     ).one()
 
     if districtr_map.num_districts != len(colors):
@@ -487,7 +493,7 @@ async def update_colors(
         bindparam(key="document_id", type_=UUIDType),
         bindparam(key="colors", type_=ARRAY(String)),
     )
-    session.execute(stmt, {"document_id": document_id, "colors": colors})
+    session.execute(stmt, {"document_id": document_id.value, "colors": colors})
     session.commit()
     return ColorsSetResult(colors=colors)
 
@@ -521,7 +527,7 @@ async def get_assignments(
 
 @app.get("/api/document/{document_id}", response_model=DocumentPublic)
 async def get_document_object(
-    document_id: str | int,
+    document_id: DocumentID = Depends(parse_document_id),
     user_id: str | None = None,
     session: Session = Depends(get_session),
 ):
@@ -611,18 +617,14 @@ async def _get_graph(gerrydb_name: str) -> Graph:
 
 @app.get("/api/document/{document_id}/contiguity")
 async def check_document_contiguity(
-    document_id: str,
     document: Annotated[Document, Depends(get_protected_document)],
+    districtr_map: Annotated[DistrictrMap, Depends(get_districtr_map)],
     zone: list[int] = Query(default=[]),
     session: Session = Depends(get_session),
 ):
-    assert document.document_id is not None
-
-    districtr_map = get_districtr_map(session=session, document_id=document.document_id)
-
     if districtr_map.child_layer is not None:
         logger.info(
-            f"Using child layer {districtr_map.child_layer} for document {document_id}"
+            f"Using child layer {districtr_map.child_layer} for document {document.document_id}"
         )
         gerrydb_name = districtr_map.child_layer
         kwargs = {"zones": zone} if len(zone) > 0 else {}
@@ -632,7 +634,7 @@ async def check_document_contiguity(
     else:
         gerrydb_name = districtr_map.parent_layer
         logger.info(
-            f"No child layer configured for document. Defauling to parent layer {gerrydb_name} for document {document_id}"
+            f"No child layer configured for document. Defauling to parent layer {gerrydb_name} for document {document.document_id}"
         )
         sql = text(
             """
@@ -666,18 +668,14 @@ async def check_document_contiguity(
 
 @app.get("/api/document/{document_id}/contiguity/{zone}/connected_component_bboxes")
 async def get_connected_component_bboxes(
-    document_id: str,
     zone: int,
     document: Annotated[Document, Depends(get_protected_document)],
+    districtr_map: Annotated[DistrictrMap, Depends(get_districtr_map)],
     session: Session = Depends(get_session),
 ):
-    assert document.document_id is not None
-
-    districtr_map = get_districtr_map(session=session, document_id=document.document_id)
-
     if districtr_map.child_layer is not None:
         logger.info(
-            f"Using child layer {districtr_map.child_layer} for document {document_id}"
+            f"Using child layer {districtr_map.child_layer} for document {document.document_id}"
         )
         gerrydb_name = districtr_map.child_layer
         zone_assignments = contiguity.get_block_assignments_bboxes(
@@ -691,7 +689,7 @@ async def get_connected_component_bboxes(
     else:
         gerrydb_name = districtr_map.parent_layer
         logger.info(
-            f"No child layer configured for document. Defauling to parent layer {gerrydb_name} for document {document_id}"
+            f"No child layer configured for document. Defauling to parent layer {gerrydb_name} for document {document.document_id}"
         )
         sql = text(
             f"""
@@ -820,7 +818,6 @@ async def update_districtrmap_metadata(
     #  response_model=list[DistrictrMapPublic]
 )
 async def get_projects(
-    *,
     session: Session = Depends(get_session),
     group: str = Query(default="states"),
     offset: int = Query(default=0, ge=0),