"""Async document and topic storage helpers.""" from __future__ import annotations import sys import uuid from datetime import datetime, timezone from typing import Optional from sqlalchemy import select, delete, text, or_ from sqlalchemy import func as sql_func from sqlalchemy.ext.asyncio import AsyncSession from db.models import Document, DocumentTopic, Topic from storage import get_storage_backend _storage = None def _backend(): """Return the lazily-instantiated StorageBackend singleton.""" global _storage _storage = _storage or get_storage_backend() return _storage def _doc_to_dict(doc: Document, topic_names: list) -> dict: return { "id": str(doc.id), "original_name": doc.filename, "filename": doc.filename, "mime_type": doc.content_type, "size_bytes": doc.size_bytes, "extracted_text": doc.extracted_text or "", "topics": topic_names, "created_at": doc.created_at.isoformat() if doc.created_at else None, "classified_at": doc.updated_at.isoformat() if doc.status == "classified" else None, "status": doc.status, } def _topic_to_dict(topic: Topic) -> dict: return { "id": str(topic.id), "name": topic.name, "description": topic.description, "color": topic.color, } def _topic_namespace_filter(name: str, user_id: Optional[uuid.UUID]): criteria = [sql_func.lower(Topic.name) == name.lower()] criteria.append(Topic.user_id.is_(None) if user_id is None else Topic.user_id == user_id) return criteria async def _load_topic_names(session: AsyncSession, doc_id: uuid.UUID) -> list: q = await session.execute( select(Topic.name) .join(DocumentTopic, DocumentTopic.topic_id == Topic.id) .where(DocumentTopic.document_id == doc_id) ) return [row[0] for row in q] async def save_metadata(session: AsyncSession, meta: dict) -> None: """Update a Document row from the legacy metadata dict shape. Keys consumed: id, extracted_text, topics (list[str]), classified_at. """ try: uid = uuid.UUID(meta["id"]) except (ValueError, KeyError): return doc = await session.get(Document, uid) if doc is None: return doc.extracted_text = meta.get("extracted_text", "") topics_list = meta.get("topics") if topics_list: await update_document_topics(session, meta["id"], topics_list) doc.status = "classified" if meta.get("classified_at") else "pending" await session.commit() async def get_metadata(session: AsyncSession, doc_id: str) -> Optional[dict]: """Return the legacy metadata dict for a document, or None if not found.""" try: uid = uuid.UUID(doc_id) except ValueError: return None doc = await session.get(Document, uid) if doc is None: return None topic_names = await _load_topic_names(session, uid) return _doc_to_dict(doc, topic_names) async def list_metadata( session: AsyncSession, user_id: uuid.UUID, topic: Optional[str] = None ) -> list: """Return metadata dicts for a user's documents, optionally filtered by topic.""" stmt = select(Document).where(Document.user_id == user_id).order_by(Document.created_at.desc()) if topic is not None: stmt = ( stmt.join(DocumentTopic, DocumentTopic.document_id == Document.id) .join(Topic, Topic.id == DocumentTopic.topic_id) .where(Topic.name == topic) ) result = await session.execute(stmt) docs = result.scalars().all() rows = [] for doc in docs: topic_names = await _load_topic_names(session, doc.id) rows.append(_doc_to_dict(doc, topic_names)) return rows async def delete_document( session: AsyncSession, doc_id: str, skip_quota: bool = False, auto_commit: bool = True, ) -> bool: """Delete a document's MinIO object and its PostgreSQL row. Returns False if the document is not found; True on success. MinIO deletion failures are logged to stderr but do not prevent the DB row deletion (the bytes may already be gone). skip_quota=True skips the quota decrement — used for cloud-stored documents that were never charged against the user's MinIO quota (T-06.2-03-01). auto_commit=False defers the session.commit() to the caller, allowing the caller to write an audit log entry in the same transaction before committing (avoids the split-transaction gap where a failed audit write loses the record while the document row is already gone). """ try: uid = uuid.UUID(doc_id) except ValueError: return False doc = await session.get(Document, uid) if doc is None: return False try: await _backend().delete_object(doc.object_key) except Exception as exc: print(f"[storage] WARNING: MinIO delete_object failed for {doc.object_key!r}: {exc}", file=sys.stderr) if not skip_quota: await session.execute( text( "UPDATE quotas " "SET used_bytes = CASE WHEN used_bytes > :delta THEN used_bytes - :delta ELSE 0 END " "WHERE user_id = :uid" ), {"delta": doc.size_bytes, "uid": doc.user_id.hex}, ) await session.delete(doc) if auto_commit: await session.commit() return True async def update_document_topics( session: AsyncSession, doc_id: str, topics: list ) -> Optional[dict]: """Replace all topic associations for a document. Auto-creates topics that don't yet exist. Returns the refreshed metadata dict, or None if the document is not found. """ try: uid = uuid.UUID(doc_id) except ValueError: return None doc = await session.get(Document, uid) if doc is None: return None await session.execute( delete(DocumentTopic).where(DocumentTopic.document_id == uid) ) seen: set = set() for name in topics: if name in seen: continue seen.add(name) topic_dict = await create_topic(session, name) session.add( DocumentTopic( document_id=uid, topic_id=uuid.UUID(topic_dict["id"]), ) ) doc.status = "classified" await session.commit() return await get_metadata(session, doc_id) async def remove_topic_from_all_documents( session: AsyncSession, topic_name: str ) -> int: """Delete all DocumentTopic rows for the named topic. Returns the number of rows deleted. """ q = await session.execute( select(Topic).where(sql_func.lower(Topic.name) == topic_name.lower()) ) topic = q.scalars().first() if topic is None: return 0 result = await session.execute( delete(DocumentTopic).where(DocumentTopic.topic_id == topic.id) ) await session.commit() return result.rowcount async def load_topics(session: AsyncSession) -> list: """Return all topics ordered by name.""" q = await session.execute(select(Topic).order_by(Topic.name)) return [_topic_to_dict(t) for t in q.scalars()] async def load_topics_for_user(session: AsyncSession, user_id: uuid.UUID) -> list: """Return system topics (user_id IS NULL) + the user's own topics, ordered by name. D-08 + D-17 + DOC-04: layered topic namespace. System topics are visible to all users; per-user topics are visible only to their owner. A user's topic list is the union of both sets. """ q = await session.execute( select(Topic).where( or_(Topic.user_id == user_id, Topic.user_id.is_(None)) ).order_by(Topic.name) ) return [_topic_to_dict(t) for t in q.scalars()] async def save_topics(session: AsyncSession, topics: list) -> None: """Idempotent bulk replace; kept for compatibility with older callers.""" await session.execute(delete(Topic)) for t in topics: session.add( Topic( id=uuid.UUID(t["id"]) if t.get("id") else uuid.uuid4(), name=t["name"], description=t.get("description", ""), color=t.get("color", "#6366f1"), ) ) await session.commit() async def get_topic(session: AsyncSession, topic_id: str) -> Optional[dict]: """Return a topic dict by UUID string, or None if not found.""" try: uid = uuid.UUID(topic_id) except ValueError: return None t = await session.get(Topic, uid) if t is None: return None return _topic_to_dict(t) async def create_topic( session: AsyncSession, name: str, description: str = "", color: str = "#6366f1", user_id: Optional[uuid.UUID] = None, ) -> dict: """Create a topic, or return an existing case-insensitive namespace match.""" q = await session.execute(select(Topic).where(*_topic_namespace_filter(name, user_id))) existing = q.scalars().first() if existing is not None: return _topic_to_dict(existing) topic = Topic(name=name, description=description, color=color, user_id=user_id) session.add(topic) await session.commit() return _topic_to_dict(topic) async def update_topic( session: AsyncSession, topic_id: str, name: Optional[str] = None, description: Optional[str] = None, color: Optional[str] = None, ) -> Optional[dict]: """Update non-None fields on a topic. Returns updated dict or None.""" try: uid = uuid.UUID(topic_id) except ValueError: return None t = await session.get(Topic, uid) if t is None: return None if name is not None: t.name = name if description is not None: t.description = description if color is not None: t.color = color await session.commit() return _topic_to_dict(t) async def delete_topic(session: AsyncSession, topic_id: str) -> Optional[str]: """Delete a topic and cascade-remove its DocumentTopic rows. Returns the deleted topic name, or None if not found. """ try: uid = uuid.UUID(topic_id) except ValueError: return None t = await session.get(Topic, uid) if t is None: return None name = t.name await session.delete(t) await session.commit() return name async def topic_doc_counts( session: AsyncSession, user_id: Optional[uuid.UUID] = None ) -> dict: """Return a mapping of topic name -> document count. If user_id is provided, counts only documents belonging to that user. This ensures a user sees the count of their own documents for each topic, not the global count across all users. """ stmt = ( select(Topic.name, sql_func.count(DocumentTopic.document_id)) .join(DocumentTopic, DocumentTopic.topic_id == Topic.id, isouter=True) ) if user_id is not None: stmt = stmt.join( Document, Document.id == DocumentTopic.document_id, isouter=True ).where( or_(Document.user_id == user_id, Document.user_id.is_(None)) ) stmt = stmt.group_by(Topic.name) q = await session.execute(stmt) return {name: count for name, count in q} __all__ = [ "save_metadata", "get_metadata", "list_metadata", "delete_document", "update_document_topics", "remove_topic_from_all_documents", "load_topics", "load_topics_for_user", "save_topics", "get_topic", "create_topic", "update_topic", "delete_topic", "topic_doc_counts", ]