fix(documents): serialize process stages with atomic update (#11)

app/main.py

@@ -4,13 +4,10 @@ from fastapi import FastAPI, Request, Depends
 from fastapi.exceptions import RequestValidationError
 from starlette.exceptions import HTTPException as StarletteHTTPException
 from app.errors import AppError, error_response
-from app.routers import auth,hubs
+from app.routers import auth, hubs, documents
 from app.deps import get_current_user
 from app.models import User
 from app.services.gemini import AiUnavailable
-from app.routers import auth, hubs, documents
-
-
 
 logger = logging.getLogger("uvicorn.error")
 
@@ -26,7 +23,6 @@ def create_app() -> FastAPI:
     app.include_router(hubs.router, prefix="/api/v1")
     app.include_router(documents.router, prefix="/api/v1")
 
-
     @app.middleware("http")
     async def catch_unhandled(request: Request, call_next):
         try:
@@ -66,7 +62,7 @@ def create_app() -> FastAPI:
         if exc.status_code == 405:
             return error_response(request, 405, "METHOD_NOT_ALLOWED", "Method not allowed for this route.")
         return error_response(request, exc.status_code, "HTTP_ERROR", str(exc.detail))
-    
+
     AI_ERRORS = {
         "rate_limited": (429, "AI_RATE_LIMITED", "The AI provider's free quota is used up. Try again later."),
         "timeout": (504, "AI_TIMEOUT", "The AI provider took too long to answer."),
@@ -78,6 +74,7 @@ def create_app() -> FastAPI:
         status, code, message = AI_ERRORS[exc.kind]
         headers = {"Retry-After": "30"} if exc.kind == "rate_limited" else None
         return error_response(request, status, code, message, headers=headers)
+
     @app.get("/healthz", tags=["system"])
     def healthz() -> dict[str, str]:
         return {"status": "ok"}

app/routers/documents.py

@@ -27,40 +27,31 @@ async def upload_document(
     current_user: User = Depends(get_current_user),
     db: Session = Depends(get_db)
 ):
-    # 1. Verify hub ownership
     hub = db.query(Hub).filter(Hub.id == hub_id, Hub.owner_id == current_user.id).first()
     if not hub:
         raise AppError(404, "HUB_NOT_FOUND", "Hub not found.")
 
-    # 2. Validate MIME type and size
     allowed_mimes = ["application/pdf", "text/plain", "text/markdown"]
     if file.content_type not in allowed_mimes:
         raise AppError(415, "UNSUPPORTED_MIME_TYPE", "Only PDF, TXT, and MD are allowed.")
     
-    # Read bytes (limit to 10MB)
     contents = await file.read()
     if len(contents) > 10 * 1024 * 1024:
         raise AppError(413, "DOCUMENT_TOO_LARGE", "File exceeds 10MB.")
 
-    # 3. Calculate checksum
     checksum = hashlib.sha256(contents).hexdigest()
-
-    # 4. Idempotency hash (hub_id + filename + size + checksum)
     hash_input = f"{hub_id}{file.filename}{len(contents)}{checksum}"
     req_hash = hashlib.sha256(hash_input.encode()).hexdigest()
 
-    # 5. Claim key
     purge_stale(db, current_user.id)
     if not claim_key(db, current_user.id, key, req_hash, f"/api/v1/hubs/{hub_id}/documents"):
         return replay_or_conflict(db, current_user.id, key, req_hash)
 
-    # 6. Check for duplicate upload
     existing = db.query(Document).filter(Document.hub_id == hub_id, Document.checksum_sha256 == checksum).first()
     if existing:
         release_claim(db, current_user.id, key)
         raise AppError(409, "DOCUMENT_ALREADY_EXISTS", "This file has already been uploaded to this hub.")
 
-    # 7. Save document
     new_doc = Document(
         hub_id=hub_id, owner_id=current_user.id, origin="upload",
         filename=file.filename, mime_type=file.content_type, size_bytes=len(contents),
@@ -70,7 +61,6 @@ async def upload_document(
     db.commit()
     db.refresh(new_doc)
 
-    # 8. Build response and complete key
     payload = DocumentResponse.model_validate(new_doc).model_dump(mode="json")
     complete_key(db, current_user.id, key, 202, payload)
     return payload
@@ -83,7 +73,6 @@ def process_document(
     current_user: User = Depends(get_current_user),
     db: Session = Depends(get_db)
 ):
-    # 1. Fetch document with tenant isolation
     doc = db.query(Document).filter(
         Document.id == doc_id, Document.hub_id == hub_id, Document.owner_id == current_user.id
     ).first()
@@ -93,21 +82,31 @@ def process_document(
     if doc.status in ("ready", "failed"):
         raise AppError(409, "DOCUMENT_NOT_PROCESSABLE", "Document is already complete or failed.")
 
-    # 2. Idempotency check: If the same key is used, return the current state without advancing
     if key and doc.process_claim and str(doc.process_claim) == str(key):
         return DocumentResponse.model_validate(doc)
 
-    # 3. Claim the stage
     if key:
         doc.process_claim = key
-        db.commit() # Commit the claim immediately so concurrent requests see it
+        db.commit()
 
-    # 4. Advance state machine
     if doc.status == "pending":
-        doc.status = "chunking"
+        rows = db.query(Document).filter(
+            Document.id == doc.id, Document.status == "pending"
+        ).update({"status": "chunking"}, synchronize_session=False)
         db.commit()
+        db.refresh(doc)
+        return DocumentResponse.model_validate(doc)
         
     elif doc.status == "chunking":
+        rows = db.query(Document).filter(
+            Document.id == doc.id, Document.status == "chunking"
+        ).update({"status": "embedding"}, synchronize_session=False)
+        db.commit()
+        
+        if rows == 0:
+            db.refresh(doc)
+            return DocumentResponse.model_validate(doc)
+
         try:
             if doc.mime_type == "application/pdf":
                 reader = PdfReader(io.BytesIO(doc.file_bytes))
@@ -123,11 +122,10 @@ def process_document(
                 db.add(Chunk(
                     document_id=doc.id, hub_id=hub_id, owner_id=current_user.id,
                     chunk_index=i, text=chunk_text_str, token_count=tokens,
-                    embedding=[0.0] * 768  # Placeholder
+                    embedding=[0.0] * 768
                 ))
                 
-            doc.file_bytes = None  # Delete raw bytes!
-            doc.status = "embedding"
+            doc.file_bytes = None
             db.commit()
         except Exception as e:
             doc.status = "failed"
@@ -135,6 +133,15 @@ def process_document(
             db.commit()
             
     elif doc.status == "embedding":
+        rows = db.query(Document).filter(
+            Document.id == doc.id, Document.status == "embedding"
+        ).update({"status": "ready"}, synchronize_session=False)
+        db.commit()
+        
+        if rows == 0:
+            db.refresh(doc)
+            return DocumentResponse.model_validate(doc)
+
         chunks = db.query(Chunk).filter(Chunk.document_id == doc.id).order_by(Chunk.chunk_index).all()
         if not chunks:
             doc.status = "failed"
@@ -148,8 +155,6 @@ def process_document(
                 for chunk, emb in zip(chunks, embeddings):
                     chunk.embedding = emb
                 
-                doc.status = "ready"
-                doc.process_claim = None # Clear claim on success
                 db.commit()
             except Exception as e:
                 doc.status = "failed"

app/schemas.py

@@ -45,7 +45,12 @@ class ChangePasswordRequest(BaseModel):
         if len(v.encode("utf-8")) > 72:
             raise ValueError("must be at most 72 bytes when encoded as UTF-8")
         return v
-# ... (keep existing schemas) ...
+
+class HubCreate(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    name: str = Field(min_length=1, max_length=120)
+    description: str | None = Field(default=None, max_length=500)
+    content: str = Field(default="", max_length=20000)
 
 class DocumentSummary(BaseModel):
     model_config = ConfigDict(from_attributes=True)
@@ -68,14 +73,7 @@ class HubResponse(BaseModel):
     updated_at: datetime
 
 class HubDetail(HubResponse):
-    documents: list[DocumentSummary] = []  
-
-class HubCreate(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-    name: str = Field(min_length=1, max_length=120)
-    description: str | None = Field(default=None, max_length=500)
-    content: str = Field(default="", max_length=20000)      
-    documents: list[DocumentSummary] = []    
+    documents: list[DocumentSummary] = []
 
 class DocumentResponse(BaseModel):
     model_config = ConfigDict(from_attributes=True)
@@ -89,4 +87,4 @@ class DocumentResponse(BaseModel):
     error_message: str | None
     chunk_count: int = 0
     created_at: datetime
-    updated_at: datetime    
+    updated_at: datetime

tests/test_concurrency_ingestion.py

@@ -0,0 +1,43 @@
+import pytest
+import time
+from uuid import uuid4
+from unittest.mock import patch
+from concurrent.futures import ThreadPoolExecutor
+from fastapi.testclient import TestClient
+from app.main import app
+from app.models import Chunk
+from app.chunking import chunk_text as original_chunk_text
+
+def test_concurrent_process_locks(client, db_session):
+    client.post("/api/v1/auth/signup", json={"email": "race_doc@test.com", "password": "Password123"})
+    token = client.post("/api/v1/auth/login", json={"email": "race_doc@test.com", "password": "Password123"}).json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+    
+    hub_res = client.post("/api/v1/hubs", headers={**headers, "X-Idempotency-Key": str(uuid4())}, json={"name": "Race Hub"})
+    hub_id = hub_res.json()["id"]
+
+    files = {"file": ("test.txt", b"Test text for race condition.", "text/plain")}
+    upload_res = client.post(f"/api/v1/hubs/{hub_id}/documents", headers={**headers, "X-Idempotency-Key": str(uuid4())}, files=files)
+    doc_id = upload_res.json()["id"]
+
+    client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers=headers)
+
+    def send_process(_):
+        with TestClient(app) as c:
+            return c.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers=headers)
+
+    def slow_chunk_text(*args, **kwargs):
+        time.sleep(0.5)
+        return original_chunk_text(*args, **kwargs)
+
+    with patch("app.routers.documents.chunk_text", side_effect=slow_chunk_text) as mock_chunk:
+        with ThreadPoolExecutor(max_workers=2) as executor:
+            future1 = executor.submit(send_process, 1)
+            future2 = executor.submit(send_process, 2)
+            r1 = future1.result()
+            r2 = future2.result()
+
+    assert r1.status_code == 200
+    assert r2.status_code == 200
+    assert mock_chunk.call_count == 1
+    assert db_session.query(Chunk).count() > 0

tests/test_ingestion.py

@@ -4,17 +4,14 @@ from unittest.mock import patch
 from app.models import Document, Chunk
 
 def test_upload_and_process_pdf(client, db_session):
-    # Setup user and hub
     client.post("/api/v1/auth/signup", json={"email": "doc@test.com", "password": "Password123"})
     token = client.post("/api/v1/auth/login", json={"email": "doc@test.com", "password": "Password123"}).json()["access_token"]
     headers = {"Authorization": f"Bearer {token}"}
     
-    # Create a hub
     key = str(uuid4())
     hub_res = client.post("/api/v1/hubs", headers={**headers, "X-Idempotency-Key": key}, json={"name": "PDF Hub"})
     hub_id = hub_res.json()["id"]
 
-    # Upload a fake text file
     file_content = b"This is a test document. It has multiple sentences to chunk."
     upload_key = str(uuid4())
     files = {"file": ("test.txt", file_content, "text/plain")}
@@ -24,24 +21,19 @@ def test_upload_and_process_pdf(client, db_session):
     doc_id = upload_res.json()["id"]
     assert upload_res.json()["status"] == "pending"
 
-    # Process Stage 1 (chunking)
     with patch("app.routers.documents.embed_texts", return_value=[[0.1]*768]):
         res1 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers=headers)
         assert res1.json()["status"] == "chunking"
 
-        # Process Stage 2 (embedding)
         res2 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers=headers)
         assert res2.json()["status"] == "embedding"
         
-        # Process Stage 3 (ready)
         res3 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers=headers)
         assert res3.json()["status"] == "ready"
 
-    # Verify chunks exist in DB
     assert db_session.query(Chunk).count() > 0
 
 def test_process_is_idempotent(client, db_session):
-    # Setup user and hub
     client.post("/api/v1/auth/signup", json={"email": "idem_doc@test.com", "password": "Password123"})
     token = client.post("/api/v1/auth/login", json={"email": "idem_doc@test.com", "password": "Password123"}).json()["access_token"]
     headers = {"Authorization": f"Bearer {token}"}
@@ -56,16 +48,12 @@ def test_process_is_idempotent(client, db_session):
     key1 = str(uuid4())
     
     with patch("app.routers.documents.embed_texts", return_value=[[0.1]*768]):
-        # 1. Advance to chunking
         res1 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers={**headers, "X-Idempotency-Key": key1})
         assert res1.json()["status"] == "chunking"
 
-        # 2. Retry the exact same request with key1
         res1_retry = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers={**headers, "X-Idempotency-Key": key1})
-        # It should return "chunking", NOT "embedding"!
         assert res1_retry.json()["status"] == "chunking"
 
-        # 3. Advance to embedding with a NEW key
         key2 = str(uuid4())
         res2 = client.post(f"/api/v1/hubs/{hub_id}/documents/{doc_id}/process", headers={**headers, "X-Idempotency-Key": key2})
-        assert res2.json()["status"] == "embedding"   
+        assert res2.json()["status"] == "embedding"