common/db.py

import os
from fastapi import Depends
import logging

from typing import Annotated
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, scoped_session, Session

from common.configuration import Configuration
from components.dbo.models.base import Base
import components.dbo.models.feedback
import components.dbo.models.acronym
import components.dbo.models.dataset
import components.dbo.models.dataset_document
import components.dbo.models.document
import components.dbo.models.log
import components.dbo.models.llm_prompt
import components.dbo.models.llm_config
import components.dbo.models.entity

CONFIG_PATH = os.environ.get('CONFIG_PATH', './config_dev.yaml')
config = Configuration(CONFIG_PATH)
logger = logging.getLogger(__name__)

engine = create_engine(config.common_config.log_sql_path, connect_args={'check_same_thread': False})

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

def get_db_session_factory():
    
    db = session_factory()
    try:
        yield db
    finally:
        db.close()

logger.info("Creating tables...")
Base.metadata.create_all(bind=engine)