from __future__ import annotations import json import logging import os import sqlite3 from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path from typing import Any logger = logging.getLogger("uvicorn.error") @dataclass class FileRecord: """SQLite-backed record for a document file.""" id: str sha256: str original_key: str current_key: str status: str = "unknown" template_name: str | None = None period_format: str | None = None document_type: str | None = None emitting_company: str | None = None period: str | None = None emission_date: str | None = None confidence: float | None = None optional_fields: dict[str, Any] = field(default_factory=dict) extracted_text: str | None = None created_at: str = "" updated_at: str = "" @classmethod def from_metadata(cls, metadata: dict[str, Any]) -> "FileRecord": record_id = str(metadata.get("sha256") or metadata.get("id") or metadata.get("current_key") or "unknown") created_at = metadata.get("created_at") or metadata.get("created") or _utc_now() updated_at = metadata.get("updated_at") or created_at return cls( id=record_id, sha256=str(metadata.get("sha256") or record_id), original_key=str(metadata.get("original_key") or metadata.get("source_key") or ""), current_key=str(metadata.get("current_key") or metadata.get("key") or ""), status=str(metadata.get("status") or "unknown"), template_name=metadata.get("template_name"), period_format=metadata.get("period_format"), document_type=metadata.get("document_type"), emitting_company=metadata.get("emitting_company"), period=metadata.get("period"), emission_date=metadata.get("emission_date"), confidence=metadata.get("confidence"), optional_fields=dict(metadata.get("optional_fields") or {}), extracted_text=metadata.get("extracted_text"), created_at=created_at, updated_at=updated_at, ) def to_row(self) -> tuple[Any, ...]: return ( self.id, self.sha256, self.original_key, self.current_key, self.status, self.template_name, self.period_format, self.document_type, self.emitting_company, self.period, self.emission_date, self.confidence, json.dumps(self.optional_fields, ensure_ascii=False), self.extracted_text, self.created_at or _utc_now(), self.updated_at or self.created_at or _utc_now(), ) @classmethod def from_row(cls, row: tuple[Any, ...]) -> "FileRecord": ( id_, sha256, original_key, current_key, status, template_name, period_format, document_type, emitting_company, period, emission_date, confidence, optional_fields, extracted_text, created_at, updated_at, ) = row try: optional_value = json.loads(optional_fields or "{}") if optional_fields else {} except json.JSONDecodeError: optional_value = {} return cls( id=str(id_), sha256=str(sha256), original_key=str(original_key), current_key=str(current_key), status=str(status), template_name=template_name, period_format=period_format, document_type=document_type, emitting_company=emitting_company, period=period, emission_date=emission_date, confidence=float(confidence) if confidence is not None else None, optional_fields=dict(optional_value), extracted_text=extracted_text, created_at=str(created_at), updated_at=str(updated_at), ) def _utc_now() -> str: return datetime.now(timezone.utc).isoformat() class SQLiteFileRepository: """Repository for storing and querying file metadata in SQLite.""" def __init__(self, db_path: str | None = None): default_path = db_path or os.environ.get("SQLITE_PATH", "./renamer.db") self.db_path = str(Path(default_path).resolve()) self._ensure_parent_dir() self.init_db() def _ensure_parent_dir(self) -> None: directory = Path(self.db_path).parent directory.mkdir(parents=True, exist_ok=True) def _connect(self) -> sqlite3.Connection: conn = sqlite3.connect(self.db_path) conn.row_factory = sqlite3.Row return conn def init_db(self) -> None: with self._connect() as conn: conn.execute( """ CREATE TABLE IF NOT EXISTS files ( id TEXT PRIMARY KEY, sha256 TEXT NOT NULL UNIQUE, original_key TEXT, current_key TEXT, status TEXT, template_name TEXT, period_format TEXT, document_type TEXT, emitting_company TEXT, period TEXT, emission_date TEXT, confidence REAL, optional_fields TEXT, extracted_text TEXT, created_at TEXT, updated_at TEXT ) """ ) columns = {row[1] for row in conn.execute("PRAGMA table_info(files)")} if "period_format" not in columns: conn.execute("ALTER TABLE files ADD COLUMN period_format TEXT") conn.execute( """ CREATE INDEX IF NOT EXISTS idx_files_status ON files(status) """ ) conn.execute( """ CREATE INDEX IF NOT EXISTS idx_files_template ON files(template_name) """ ) conn.execute( """ CREATE INDEX IF NOT EXISTS idx_files_period_format ON files(period_format) """ ) conn.commit() def upsert(self, record: FileRecord) -> FileRecord: record.updated_at = record.updated_at or _utc_now() record.created_at = record.created_at or record.updated_at with self._connect() as conn: conn.execute( """ INSERT INTO files ( id, sha256, original_key, current_key, status, template_name, period_format, document_type, emitting_company, period, emission_date, confidence, optional_fields, extracted_text, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET sha256 = excluded.sha256, original_key = excluded.original_key, current_key = excluded.current_key, status = excluded.status, template_name = excluded.template_name, period_format = excluded.period_format, document_type = excluded.document_type, emitting_company = excluded.emitting_company, period = excluded.period, emission_date = excluded.emission_date, confidence = excluded.confidence, optional_fields = excluded.optional_fields, extracted_text = excluded.extracted_text, updated_at = excluded.updated_at """, record.to_row(), ) conn.commit() return record def get_by_id(self, file_id: str) -> FileRecord | None: with self._connect() as conn: row = conn.execute( "SELECT * FROM files WHERE id = ? OR sha256 = ?", (file_id, file_id), ).fetchone() if row is None: return None return FileRecord.from_row(tuple(row)) def _file_conditions( self, status: str | None = None, template: str | None = None, period_format: str | None = None, q: str | None = None, current_key: str | None = None, emitting_company: str | None = None, document_type: str | None = None, period: str | None = None, confidence: str | None = None, updated_at: str | None = None, emission_date: str | None = None, additional_fields: dict[str, str] | None = None, ) -> tuple[list[str], list[Any]]: conditions: list[str] = [] params: list[Any] = [] if status: conditions.append("status = ?") params.append(status) if template: conditions.append("LOWER(template_name) LIKE ?") params.append(f"%{template.lower()}%") if current_key: value = f"%{current_key.lower()}%" conditions.append("(LOWER(current_key) LIKE ? OR LOWER(template_name) LIKE ?)") params.extend([value, value]) for column, value in ( ("emitting_company", emitting_company), ("document_type", document_type), ("period", period), ("updated_at", updated_at), ("emission_date", emission_date), ("period_format", period_format), ): if value: conditions.append(f"LOWER({column}) LIKE ?") params.append(f"%{value.lower()}%") for field_name, field_value in (additional_fields or {}).items(): if field_value: json_path = '$."' + field_name.replace('"', '\\"') + '"' conditions.append("LOWER(COALESCE(json_extract(optional_fields, ?), '')) LIKE ?") params.extend([json_path, f"%{field_value.lower()}%"]) if confidence: confidence_filter = confidence.strip().removesuffix("%") if confidence_filter: conditions.append("CAST(ROUND(confidence * 100, 1) AS TEXT) LIKE ?") params.append(f"%{confidence_filter.lower()}%") if q: conditions.append( "(LOWER(current_key) LIKE ? OR LOWER(original_key) LIKE ? OR LOWER(template_name) LIKE ? OR LOWER(document_type) LIKE ? OR LOWER(emitting_company) LIKE ? OR LOWER(period) LIKE ? OR LOWER(emission_date) LIKE ? OR LOWER(period_format) LIKE ? OR LOWER(extracted_text) LIKE ?)" ) like_value = f"%{q.lower()}%" params.extend([like_value] * 9) return conditions, params def list_files( self, status: str | None = None, template: str | None = None, period_format: str | None = None, q: str | None = None, current_key: str | None = None, emitting_company: str | None = None, document_type: str | None = None, period: str | None = None, confidence: str | None = None, updated_at: str | None = None, emission_date: str | None = None, additional_fields: dict[str, str] | None = None, sort_by: str = "updated_at", sort_order: str = "desc", page: int = 1, page_size: int = 25, ) -> list[FileRecord]: conditions, params = self._file_conditions( status=status, template=template, period_format=period_format, q=q, current_key=current_key, emitting_company=emitting_company, document_type=document_type, period=period, confidence=confidence, updated_at=updated_at, emission_date=emission_date, additional_fields=additional_fields, ) query = "SELECT * FROM files" if conditions: query += " WHERE " + " AND ".join(conditions) sortable_columns = { "status": "status", "current_key": "current_key", "template_name": "template_name", "period_format": "period_format", "emitting_company": "emitting_company", "document_type": "document_type", "period": "period", "emission_date": "emission_date", "confidence": "confidence", "updated_at": "updated_at", } if sort_by.startswith("additional_field__"): field_name = sort_by.removeprefix("additional_field__") order_column = "json_extract(optional_fields, ?)" params.append('$.' + '"' + field_name.replace('"', '\\"') + '"') else: order_column = sortable_columns.get(sort_by, "updated_at") order_direction = "ASC" if sort_order.lower() == "asc" else "DESC" query += f" ORDER BY {order_column} {order_direction} LIMIT ? OFFSET ?" params.extend([page_size, (page - 1) * page_size]) with self._connect() as conn: rows = conn.execute(query, params).fetchall() return [FileRecord.from_row(tuple(row)) for row in rows] def count_files( self, status: str | None = None, template: str | None = None, period_format: str | None = None, q: str | None = None, current_key: str | None = None, emitting_company: str | None = None, document_type: str | None = None, period: str | None = None, confidence: str | None = None, updated_at: str | None = None, emission_date: str | None = None, additional_fields: dict[str, str] | None = None, ) -> int: conditions, params = self._file_conditions( status=status, template=template, period_format=period_format, q=q, current_key=current_key, emitting_company=emitting_company, document_type=document_type, period=period, confidence=confidence, updated_at=updated_at, emission_date=emission_date, additional_fields=additional_fields, ) query = "SELECT COUNT(*) FROM files" if conditions: query += " WHERE " + " AND ".join(conditions) with self._connect() as conn: row = conn.execute(query, params).fetchone() return int(row[0]) def get_status_counts(self) -> dict[str, int]: with self._connect() as conn: row = conn.execute( """ SELECT COUNT(*) AS total, SUM(CASE WHEN status = 'processed' THEN 1 ELSE 0 END) AS processed, SUM(CASE WHEN status = 'validated' THEN 1 ELSE 0 END) AS validated, SUM(CASE WHEN status = 'not processed' THEN 1 ELSE 0 END) AS pending FROM files """ ).fetchone() return {key: int(row[key] or 0) for key in ("total", "processed", "validated", "pending")} def delete_all(self) -> None: with self._connect() as conn: conn.execute("DELETE FROM files") conn.commit() def rebuild_from_sidecars(self, storage: Any) -> list[FileRecord]: """Rebuild the SQLite index from JSON sidecar metadata in storage.""" self.delete_all() records: list[FileRecord] = [] for key in sorted(storage.list_objects(prefix="")): if not key.endswith(".meta.json"): continue raw = storage.get_object(key) if raw is None: continue if isinstance(raw, bytes): raw = raw.decode("utf-8", errors="replace") try: metadata = json.loads(raw) except json.JSONDecodeError: continue if not metadata: continue record = FileRecord.from_metadata(metadata) self.upsert(record) records.append(record) logger.info(f"Rebuilt SQLite index from {len(records)} sidecar metadata files.") return records