File size: 2,067 Bytes
af24ae8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """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()
|