ds_quickbooks/database.py

from typing import Annotated

from fastapi import Depends
from sqlalchemy import Column, DateTime, Float, Integer, String, create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session, sessionmaker

SQLALCHEMY_DATABASE_URL = "sqlite:///./sql_app.db"

engine = create_engine(
    SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
)

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


db_dependency = Annotated[Session, Depends(get_db)]
Base = declarative_base()


def create_db_tables():
    Base.metadata.create_all(bind=engine)


def clear_all_data():
    """Clear all data from the database (useful for testing)"""
    db = SessionLocal()
    try:
        db.query(Transaction).delete()
        db.query(Receipt).delete()
        db.commit()
    finally:
        db.close()


# Transactions table
class Transaction(Base):
    __tablename__ = "transactions"

    id = Column(Integer, primary_key=True, index=True)
    transaction_id = Column(String, unique=True, index=True)
    amount = Column(Float, nullable=False)
    date = Column(DateTime, nullable=False)
    vendor = Column(String, nullable=False)
    description = Column(String, nullable=True)
    category = Column(String, nullable=True)
    tax_amount = Column(Float, nullable=True)
    categorisation_id = Column(String, nullable=True)
    user_id = Column(String, nullable=True)


# Receipts table
class Receipt(Base):
    __tablename__ = "receipts"

    id = Column(Integer, primary_key=True, index=True)
    receipt_id = Column(String, unique=True, index=True)
    file_id = Column(String, unique=True, index=True)
    amount = Column(Float, nullable=False)
    date = Column(DateTime, nullable=False)
    vendor = Column(String, nullable=False)
    description = Column(String, nullable=True)
    category = Column(String, nullable=True)
    tax_amount = Column(Float, nullable=True)
    confidence = Column(Float, nullable=True)
    extraction_success = Column(String, nullable=True)
    error_message = Column(String, nullable=True)