Phase 4: async pipeline (Celery+Redis), Postgres job state, local-fs blob storage, API-key auth, Prometheus metrics (#3)
* Phase 4: async pipeline (Celery+Redis), Postgres job state, local-fs blob storage, API-key auth, Prometheus metrics Co-Authored-By: adrian kuman firmansah <adriancuman@gmail.com> * Phase 4: fix sync-mode rollback orphaning blobs + use is_relative_to for path-escape check Devin Review on PR #3 found two real bugs: 1. Sync path mark_failed was rolled back by the request-scoped session. When the pipeline raised an exception in ?sync=true mode, _run_inline modified the FastAPI session and re-raised; get_session caught the exception, called session.rollback(), and wiped both the create() and the mark_failed() writes. The blob was already on disk, so it was permanently orphaned with no DB record. Fix: commit the pending row immediately after create(), and run all subsequent state transitions in independent session_scope blocks (matching the worker task pattern). 2. _resolve used str.startswith for path-escape detection, which lets a sibling directory whose name begins with the storage root pass (e.g. /app/blobs_evil vs /app/blobs). Switched to Path.is_relative_to. Added regression tests for both. Co-Authored-By: adrian kuman firmansah <adriancuman@gmail.com> * Phase 4: honor queue_enabled setting + resolve base_dir for path comparisons Two more bugs found by Devin Review: 3. queue_enabled was declared in config and documented in .env.example but never read by the route. A fresh dev install with QUEUE_ENABLED=false (the default) would still enqueue, then fail with a Redis connection error. Fixed by making the ?sync= query param default to None and resolving to (not queue_enabled) inside the route. Tests now set QUEUE_ENABLED=true so the async flow stays exercised, and a new test verifies the inline fallback when the queue is disabled. 4. LocalFsBlobStorage stored base_dir as-is. _resolve resolved its candidate paths, so the empty-dir cleanup loop in delete() compared a resolved candidate against an unresolved base_dir and broke on the first iteration (no cleanup ever happened). Fixed by resolving base_dir once in __init__ so every path comparison is apples-to-apples. Co-Authored-By: adrian kuman firmansah <adriancuman@gmail.com> * Phase 4: derive ocr_jobs_total from DB so worker writes are visible at /metrics Devin Review correctly noted the Counter-based JOBS_TOTAL would never increment in production because the worker runs in a separate process from the API and the registry is process-local. Replaced JOBS_TOTAL with a custom Collector that issues SELECT status, COUNT(*) FROM jobs GROUP BY status on every /metrics scrape. Result: the metric stays accurate regardless of which process wrote the row. Also corrected the metrics.py docstring (the old comment claimed the counter was 'incremented by the worker', which was the bug). Removed the JOBS_TOTAL.inc() calls from the sync route — the DB collector covers both paths now. JOB_PROCESSING_SECONDS stays as an API-process histogram with an updated docstring noting its scope; cross-process latency belongs to derived dashboards over jobs.created_at/updated_at. Added regression test test_metrics_jobs_total_reflects_worker_writes. Co-Authored-By: adrian kuman firmansah <adriancuman@gmail.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: adrian kuman firmansah <adriancuman@gmail.com>
This commit is contained in:
devin-ai-integration[bot]
committed by
GitHub
GitHub
parent
33b38aacc7
commit
2112023b6e
@@ -1,58 +1,194 @@
|
||||
"""Documents API — Phase 1 synchronous endpoint.
|
||||
"""Documents API.
|
||||
|
||||
POST /documents accepts a single PDF or image upload, runs the synchronous
|
||||
pipeline inline, and returns the structured result. This is suitable for
|
||||
development and low-traffic production; Phase 4 will introduce an async
|
||||
queue and a polling-style API at the same path.
|
||||
Phase 1 shipped a single synchronous endpoint. Phase 4 adds an async
|
||||
flow on top:
|
||||
|
||||
* `POST /documents` — async by default. Saves the upload to blob
|
||||
storage, creates a `pending` job row, and
|
||||
enqueues a Celery task. Returns `202` with
|
||||
the job id.
|
||||
* `POST /documents?sync=true` — runs the pipeline inline (the original
|
||||
Phase 1 behaviour). Useful for tests and
|
||||
small-volume single-tenant deploys without
|
||||
a Celery worker.
|
||||
* `GET /documents/{job_id}` — returns the current job state. Async
|
||||
clients poll this until `status` is in a
|
||||
terminal state (completed / needs_review /
|
||||
failed).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from uuid import uuid4
|
||||
from typing import Annotated
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
from fastapi import APIRouter, File, UploadFile, status
|
||||
from fastapi import APIRouter, Depends, File, HTTPException, Query, Response, UploadFile, status
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from ocr_sprint.api.deps.auth import require_api_key
|
||||
from ocr_sprint.api.deps.db import get_session
|
||||
from ocr_sprint.api.errors import UnsupportedDocumentError
|
||||
from ocr_sprint.api.metrics import JOB_PROCESSING_SECONDS
|
||||
from ocr_sprint.config import get_settings
|
||||
from ocr_sprint.db.base import session_scope
|
||||
from ocr_sprint.db.repositories import JobNotFoundError, JobRepository
|
||||
from ocr_sprint.pipeline.ingest import detect_source_kind
|
||||
from ocr_sprint.pipeline.orchestrator import run_pipeline
|
||||
from ocr_sprint.schemas.document import DocumentResponse
|
||||
from ocr_sprint.schemas.document import DocumentResponse, DocumentStatus
|
||||
from ocr_sprint.schemas.extraction import ExtractionResult
|
||||
from ocr_sprint.storage.blob import get_blob_storage
|
||||
from ocr_sprint.utils.logging import get_logger
|
||||
|
||||
router = APIRouter(prefix="/documents", tags=["documents"])
|
||||
router = APIRouter(
|
||||
prefix="/documents",
|
||||
tags=["documents"],
|
||||
dependencies=[Depends(require_api_key)],
|
||||
)
|
||||
_logger = get_logger(__name__)
|
||||
|
||||
_MAX_UPLOAD_BYTES = 25 * 1024 * 1024 # 25 MB
|
||||
|
||||
# ---------- helpers ----------
|
||||
|
||||
|
||||
@router.post("", status_code=status.HTTP_200_OK, response_model=DocumentResponse)
|
||||
async def create_document(file: UploadFile = File(...)) -> DocumentResponse:
|
||||
"""Run OCR + extraction synchronously on a single upload."""
|
||||
def _enforce_size(content: bytes) -> None:
|
||||
s = get_settings()
|
||||
if not content:
|
||||
raise UnsupportedDocumentError("Uploaded file is empty.")
|
||||
max_bytes = s.blob_max_upload_mb * 1024 * 1024
|
||||
if len(content) > max_bytes:
|
||||
raise UnsupportedDocumentError(f"Uploaded file exceeds {s.blob_max_upload_mb} MB limit.")
|
||||
|
||||
|
||||
def _row_to_response(row: object) -> DocumentResponse:
|
||||
# Local import to avoid a circular import at module load time.
|
||||
from ocr_sprint.db.models import JobRow
|
||||
|
||||
assert isinstance(row, JobRow)
|
||||
status_enum = DocumentStatus(row.status)
|
||||
result_obj: ExtractionResult | None = None
|
||||
if row.result is not None:
|
||||
result_obj = ExtractionResult.model_validate(row.result)
|
||||
return DocumentResponse(
|
||||
job_id=row.job_id,
|
||||
status=status_enum,
|
||||
confidence=row.confidence,
|
||||
data=result_obj,
|
||||
review_flags=list(row.review_flags or []),
|
||||
error=row.error,
|
||||
)
|
||||
|
||||
|
||||
# ---------- POST ----------
|
||||
|
||||
|
||||
@router.post("", response_model=DocumentResponse)
|
||||
async def create_document(
|
||||
file: Annotated[UploadFile, File(...)],
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
response: Response,
|
||||
sync: Annotated[
|
||||
bool | None,
|
||||
Query(description="Run pipeline inline (skip queue). Defaults to !queue_enabled."),
|
||||
] = None,
|
||||
) -> DocumentResponse:
|
||||
# When the queue is disabled (default for local dev), running the async
|
||||
# path would try to dial Redis and fail with a 500. Auto-fall-back to the
|
||||
# inline pipeline unless the caller explicitly asked for async.
|
||||
if sync is None:
|
||||
sync = not get_settings().queue_enabled
|
||||
|
||||
job_id = uuid4()
|
||||
log = _logger.bind(job_id=str(job_id), filename=file.filename or "")
|
||||
|
||||
content = await file.read()
|
||||
if not content:
|
||||
raise UnsupportedDocumentError("Uploaded file is empty.")
|
||||
if len(content) > _MAX_UPLOAD_BYTES:
|
||||
raise UnsupportedDocumentError(
|
||||
f"Uploaded file exceeds {_MAX_UPLOAD_BYTES // (1024 * 1024)} MB limit."
|
||||
)
|
||||
_enforce_size(content)
|
||||
|
||||
log.info("documents.received", size=len(content))
|
||||
storage = get_blob_storage()
|
||||
blob_key = storage.put(content, original_filename=file.filename)
|
||||
source_kind = detect_source_kind(content)
|
||||
JobRepository(session).create(
|
||||
job_id=job_id,
|
||||
filename=file.filename or "",
|
||||
source_kind=source_kind,
|
||||
blob_key=blob_key,
|
||||
)
|
||||
# Commit the `pending` row immediately so it is observable regardless
|
||||
# of what happens next. Both code paths below open their own session
|
||||
# for state transitions; that way an exception in `_run_inline` cannot
|
||||
# roll back the create() (which would orphan the blob on disk).
|
||||
session.commit()
|
||||
log.info("documents.received", size=len(content), blob_key=blob_key, sync=sync)
|
||||
|
||||
if sync:
|
||||
# Status code stays at the default 200; the body's `status` field
|
||||
# tells the client whether the job needs review.
|
||||
return await _run_inline(job_id, content)
|
||||
|
||||
# Async path — enqueue and return 202. The Celery worker will pick up
|
||||
# the row using its own session.
|
||||
from ocr_sprint.worker.tasks import process_document_task
|
||||
|
||||
process_document_task.delay(str(job_id))
|
||||
with session_scope() as poll:
|
||||
row = JobRepository(poll).get_or_raise(job_id)
|
||||
body = _row_to_response(row)
|
||||
response.status_code = status.HTTP_202_ACCEPTED
|
||||
return body
|
||||
|
||||
|
||||
async def _run_inline(job_id: UUID, content: bytes) -> DocumentResponse:
|
||||
"""Synchronous pipeline execution.
|
||||
|
||||
Each state transition opens its own short session so the request-scoped
|
||||
session's rollback-on-exception behaviour cannot wipe out the
|
||||
``mark_failed`` write or strand the blob on disk.
|
||||
"""
|
||||
import time
|
||||
|
||||
with session_scope() as s:
|
||||
JobRepository(s).mark_processing(job_id)
|
||||
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
output = run_pipeline(content)
|
||||
except ValueError as exc:
|
||||
with session_scope() as s:
|
||||
JobRepository(s).mark_failed(job_id, error=str(exc))
|
||||
raise UnsupportedDocumentError(str(exc)) from exc
|
||||
except Exception as exc:
|
||||
with session_scope() as s:
|
||||
JobRepository(s).mark_failed(job_id, error=str(exc))
|
||||
raise
|
||||
|
||||
log.info(
|
||||
"documents.completed",
|
||||
status=output.status.value,
|
||||
confidence=round(output.confidence, 3),
|
||||
flags=[f.value for f in output.result.review_flags],
|
||||
)
|
||||
return DocumentResponse(
|
||||
job_id=job_id,
|
||||
status=output.status,
|
||||
confidence=output.confidence,
|
||||
data=output.result,
|
||||
review_flags=[f.value for f in output.result.review_flags],
|
||||
)
|
||||
flags = [f.value for f in output.result.review_flags]
|
||||
JOB_PROCESSING_SECONDS.observe(time.perf_counter() - started)
|
||||
with session_scope() as s:
|
||||
repo = JobRepository(s)
|
||||
repo.mark_completed(
|
||||
job_id,
|
||||
status=output.status,
|
||||
confidence=output.confidence,
|
||||
result=output.result.model_dump(mode="json"),
|
||||
review_flags=flags,
|
||||
)
|
||||
row = repo.get_or_raise(job_id)
|
||||
return _row_to_response(row)
|
||||
|
||||
|
||||
# ---------- GET ----------
|
||||
|
||||
|
||||
@router.get(
|
||||
"/{job_id}",
|
||||
response_model=DocumentResponse,
|
||||
)
|
||||
async def get_document(
|
||||
job_id: UUID,
|
||||
session: Annotated[Session, Depends(get_session)],
|
||||
) -> DocumentResponse:
|
||||
repo = JobRepository(session)
|
||||
try:
|
||||
row = repo.get_or_raise(job_id)
|
||||
except JobNotFoundError as exc:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(exc)) from exc
|
||||
return _row_to_response(row)
|
||||
|
||||
Reference in New Issue
Block a user