multi_layer_operation_predictor/db/db_setup.py

from sqlalchemy import create_engine, Column, Integer, String, JSON, DateTime
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
from typing import Generator
from contextlib import contextmanager
from datetime import datetime
import pytz



# Create database URL (update this according to your database type)
SQLALCHEMY_DATABASE_URL ="postgresql+psycopg2://codedb:codedb@localhost:5432/codedb"

# Create engine with connection pooling
engine = create_engine(
    SQLALCHEMY_DATABASE_URL,
    pool_size=5,
    max_overflow=10,
    pool_timeout=30,
    pool_recycle=1800
)

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

class ExtractedFile(Base):
    __tablename__ = "extracted_files"
    
    id = Column(Integer, primary_key=True)
    file_name = Column(String(255), nullable=False)
    file_data = Column(JSON, nullable=False)
    repository_url = Column(String, nullable=False)
    created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(pytz.UTC))

    def __repr__(self):
        return f"<ExtractedFile(id={self.id}, file_name={self.file_name}, repository_url={self.repository_url})>"

    @classmethod
    def get_by_extension(cls, db: Session, extension: str):
        """Get all files with specific extension"""
        return db.query(cls).filter(
            cls.file_name.endswith(extension)
        ).all()

@contextmanager
def get_db() -> Generator[Session, None, None]:
    """Context manager for database sessions"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()