"""Database models and session management.""" from __future__ import annotations import os from datetime import datetime from enum import Enum from typing import Generator from sqlalchemy import JSON, DateTime, String, Text, create_engine from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column, sessionmaker class Base(DeclarativeBase): pass class ReviewStatus(str, Enum): PENDING = "pending" APPROVED = "approved" REJECTED = "rejected" class InvoiceReview(Base): __tablename__ = "invoice_reviews" id: Mapped[str] = mapped_column(String(36), primary_key=True) vendor_name: Mapped[str | None] = mapped_column(String(255), nullable=True) invoice_number: Mapped[str | None] = mapped_column(String(100), nullable=True) total_amount: Mapped[float | None] = mapped_column(nullable=True) currency: Mapped[str] = mapped_column(String(10), default="IRR") confidence: Mapped[float] = mapped_column(default=0.0) status: Mapped[str] = mapped_column(String(20), default=ReviewStatus.PENDING.value) invoice_data: Mapped[dict] = mapped_column(JSON) validation_data: Mapped[dict | None] = mapped_column(JSON, nullable=True) reviewer: Mapped[str | None] = mapped_column(String(100), nullable=True) review_notes: Mapped[str | None] = mapped_column(Text, nullable=True) created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) reviewed_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) def get_database_url() -> str: return os.getenv( "DATABASE_URL", "postgresql://docflow:docflow_secret@postgres:5432/docflow", ) engine = create_engine(get_database_url(), pool_pre_ping=True) SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False) def init_db() -> None: Base.metadata.create_all(bind=engine) def get_db() -> Generator[Session, None, None]: db = SessionLocal() try: yield db finally: db.close()