Update main.py

main.py

@@ -45,10 +45,6 @@ OUTPUT_DIR.mkdir(exist_ok=True)
 # Initialize redactor
 redactor = PDFRedactor()
 
-# ---------------- In-Memory Redaction Status Tracker ----------------
-# request_id -> status (pending | processing | completed | failed)
-redaction_status: Dict[str, str] = {}
-
 # ---------------- Response Models ----------------
 class RedactionEntity(BaseModel):
     entity_type: str
@@ -69,15 +65,35 @@ class RedactionStatusResponse(BaseModel):
     files: List[str]
     message: str
 
-
 class HealthResponse(BaseModel):
     status: str
     version: str
     model_loaded: bool
 
+# ---------------- DB Status Helpers ----------------
+def set_request_status(request_id: str, status: str):
+    """Update the status column in document_requests for the given request_id."""
+    supabase.from_("document_requests").update({"status": status}).eq("id", request_id).execute()
+    logger.info(f"Request {request_id} status -> {status}")
+
+def get_request_status(request_id: str) -> str:
+    """Fetch current status from document_requests."""
+    response = (
+        supabase
+        .from_("document_requests")
+        .select("status")
+        .eq("id", request_id)
+        .maybe_single()
+        .execute()
+    )
+    if response.data:
+        return response.data["status"]
+    return "not_found"
+
 # ---------------- Helper Functions ----------------
 def get_public_url(bucket: str, storage_path: str) -> str:
     return f"{os.getenv('SUPABASE_URL')}/storage/v1/object/public/{bucket}/{storage_path}"
+
 def cleanup_files(job_id: str):
     """Clean up temporary files after a delay"""
     try:
@@ -103,10 +119,8 @@ def download_file_from_supabase(bucket: str, storage_path: str, local_path: Path
 
 def upload_file_to_supabase(bucket: str, storage_path: str, local_path: Path):
     logger.info(f"Uploading {local_path} to {storage_path}")
-
     with local_path.open("rb") as f:
         content = f.read()
-
     supabase.storage.from_(bucket).upload(
         path=storage_path,
         file=content,
@@ -117,11 +131,15 @@ def upload_file_to_supabase(bucket: str, storage_path: str, local_path: Path):
     )
 
 def redact_request(request_id: str, bucket: str = "doc_storage"):
-    """Background task: redact all files for a given request_id"""
+    """
+    Background task: redact all files for a given request_id.
+    DB writes: 2 total — one at start (redacting), one at end (redacted | failed).
+    The 'pending' write is done by the endpoint before this task is dispatched.
+    """
     try:
-        redaction_status[request_id] = "processing"
+        # Write 1: mark as redacting
+        set_request_status(request_id, "redacting")
 
-        # Fetch all files for this request_id
         response = (
             supabase
             .from_("request_files")
@@ -131,9 +149,6 @@ def redact_request(request_id: str, bucket: str = "doc_storage"):
         )
 
         files = response.data
-
-        if not files:
-            raise Exception(f"No files found for request {request_id}")
         if not files:
             raise Exception(f"No files found for request {request_id}")
 
@@ -142,23 +157,18 @@ def redact_request(request_id: str, bucket: str = "doc_storage"):
             local_upload = UPLOAD_DIR / f"{uuid.uuid4()}.pdf"
             local_output = OUTPUT_DIR / f"{uuid.uuid4()}_redacted.pdf"
 
-            # Download from Supabase
             download_file_from_supabase(bucket, storage_path, local_upload)
-
-            # Redact
             redactor.redact_document(pdf_path=str(local_upload), output_path=str(local_output))
-
-            # Upload redacted back to same path
             upload_file_to_supabase(bucket, storage_path, local_output)
-
-            # Cleanup local files
             cleanup_temp_files([local_upload, local_output])
 
-        redaction_status[request_id] = "completed"
+        # Write 2: mark as redacted
+        set_request_status(request_id, "redacted")
 
     except Exception as e:
         logger.error(f"Redaction failed for {request_id}: {str(e)}")
-        redaction_status[request_id] = "failed"
+        # Write 2 (error path): mark as failed
+        set_request_status(request_id, "failed")
 
 # ----------------- Existing Endpoints -----------------
 @app.get("/", response_model=HealthResponse)
@@ -266,15 +276,21 @@ async def get_stats():
 # ----------------- NEW Endpoints -----------------
 @app.post("/redact_by_request/{request_id}", response_model=RedactionStatusResponse)
 async def redact_by_request(request_id: str, background_tasks: BackgroundTasks):
-    if redaction_status.get(request_id) == "processing":
+    # Check current DB status to avoid re-triggering an in-progress job
+    current_status = get_request_status(request_id)
+
+    if current_status == "redacting":
         return RedactionStatusResponse(
             request_id=request_id,
-            status="processing",
+            status="redacting",
             files=[],
             message="Redaction already in progress"
         )
-    redaction_status[request_id] = "pending"
+
+    # Write 1: set pending before dispatching background task
+    set_request_status(request_id, "pending")
     background_tasks.add_task(redact_request, request_id)
+
     return RedactionStatusResponse(
         request_id=request_id,
         status="pending",
@@ -284,13 +300,11 @@ async def redact_by_request(request_id: str, background_tasks: BackgroundTasks):
 
 @app.get("/redaction_status/{request_id}", response_model=RedactionStatusResponse)
 async def get_redaction_status(request_id: str):
-    status = redaction_status.get(request_id, "not_found")
+    status = get_request_status(request_id)
 
-    # Default empty response
     files: List[str] = []
 
-    if status == "completed":
-        # Fetch file paths from DB
+    if status == "redacted":
         response = (
             supabase
             .from_("request_files")
@@ -298,35 +312,23 @@ async def get_redaction_status(request_id: str):
             .eq("request_id", request_id)
             .execute()
         )
-        
         if response.data:
             files = [
                 get_public_url("doc_storage", row["storage_path"])
                 for row in response.data
             ]
 
-    message = (
-        "Redaction completed"
-        if status == "completed"
-        else "Redaction pending"
-        if status == "pending"
-        else "Redaction failed"
-        if status == "failed"
-        else "Request not found"
-    )
+    message = {
+        "redacted": "Redaction completed",
+        "pending": "Redaction pending",
+        "redacting": "Redaction in progress",
+        "failed": "Redaction failed",
+        "not_found": "Request not found",
+    }.get(status, status)
 
     return RedactionStatusResponse(
         request_id=request_id,
         status=status,
         files=files,
         message=message
-    )
-
-# ----------------- Run Server -----------------
-# if __name__ == "__main__":
-#     uvicorn.run(
-#         "main:app",
-#         host="localhost",
-#         port=2700,
-#         reload=False
-#     )
+    )