models.py

from sqlalchemy import Column, Integer, String, JSON, DateTime, create_engine, func, Enum, Text, Index
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from datetime import datetime
import os
from dotenv import load_dotenv
import enum

load_dotenv()

Base = declarative_base()

class TaskStatus(enum.Enum):
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"

class ApiCall(Base):
    __tablename__ = 'api_calls'
    
    id = Column(Integer, primary_key=True)
    endpoint = Column(String(50))
    prompt = Column(JSON)
    data = Column(JSON, nullable=True)
    response = Column(JSON)
    tokens_used = Column(Integer)
    file_name = Column(String(255), nullable=True)
    file_type = Column(String(100), nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow, index=True)  # Added index
    status = Column(Enum(TaskStatus), default=TaskStatus.PENDING, index=True)  # Added index
    extracted_text = Column(Text, nullable=True)  # Changed to Text type
    task_id = Column(String(50), nullable=True)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

# Database connection
DB_ENGINE = os.getenv('DB_ENGINE')
DB_USER = os.getenv('DB_USER')
DB_PASSWORD = os.getenv('DB_PASSWORD')
DB_HOST = os.getenv('DB_HOST')
DB_NAME = os.getenv('DB_NAME')

DATABASE_URL = f"{DB_ENGINE}+pymysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}/{DB_NAME}"

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

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