app/db.py

"""MongoDB metadata store for uploaded documents.

If MongoDB is unreachable (e.g. Atlas IP allow-list not configured yet), the
store gracefully falls back to an in-process dictionary so the rest of the app
keeps working. Fallback data is not persisted across restarts.
"""
from __future__ import annotations

from datetime import datetime, timezone
from functools import lru_cache
from typing import List, Optional

from pymongo import MongoClient, DESCENDING

from .config import get_settings

# In-memory fallback store (used only when Mongo can't be reached).
_MEM: dict[str, dict] = {}
_USING_FALLBACK = False


@lru_cache
def _client() -> MongoClient:
    settings = get_settings()
    return MongoClient(
        settings.mongodb_uri,
        appname="doc-intelligence-rag",
        serverSelectionTimeoutMS=4000,
    )


def _collection():
    settings = get_settings()
    return _client()[settings.mongodb_db]["documents"]


def _fallback(reason: str = "") -> None:
    global _USING_FALLBACK
    if not _USING_FALLBACK:
        _USING_FALLBACK = True


def using_fallback() -> bool:
    return _USING_FALLBACK


def save_document(
    document_id: str,
    filename: str,
    content_type: str,
    num_chunks: int,
    num_chars: int,
    ocr_used: bool,
) -> dict:
    doc = {
        "_id": document_id,
        "filename": filename,
        "content_type": content_type or "application/octet-stream",
        "num_chunks": num_chunks,
        "num_chars": num_chars,
        "ocr_used": ocr_used,
        "created_at": datetime.now(timezone.utc).isoformat(),
    }
    try:
        _collection().insert_one(dict(doc))
    except Exception as exc:  # noqa: BLE001
        _fallback(str(exc))
        _MEM[document_id] = doc
    return doc


def list_documents() -> List[dict]:
    try:
        return list(_collection().find().sort("created_at", DESCENDING))
    except Exception as exc:  # noqa: BLE001
        _fallback(str(exc))
        return sorted(_MEM.values(), key=lambda d: d["created_at"], reverse=True)


def get_document(document_id: str) -> Optional[dict]:
    try:
        return _collection().find_one({"_id": document_id})
    except Exception as exc:  # noqa: BLE001
        _fallback(str(exc))
        return _MEM.get(document_id)


def delete_document(document_id: str) -> int:
    try:
        result = _collection().delete_one({"_id": document_id})
        return result.deleted_count
    except Exception as exc:  # noqa: BLE001
        _fallback(str(exc))
        return 1 if _MEM.pop(document_id, None) else 0


def ping() -> bool:
    _client().admin.command("ping")
    return True