cloud-api/src/database.py

"""
Database connections and init

Exports:
 - db_dependency
 - Base (sqlalchemy base model)
"""
from typing import Annotated
from sqlalchemy import create_engine, StaticPool
from sqlalchemy.orm import DeclarativeBase, sessionmaker, Session

from fastapi import Depends

from src.constants import Environment
from src.config import SQLALCHEMY_DATABASE_URI, settings as global_settings

if global_settings.ENVIRONMENT == Environment.TESTING:
	connect_args = {"check_same_thread": False}
	engine = create_engine(SQLALCHEMY_DATABASE_URI.get_secret_value(), connect_args=connect_args, poolclass=StaticPool)
else:
	engine = create_engine(SQLALCHEMY_DATABASE_URI.get_secret_value())


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


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


db_dependency = Annotated[Session, Depends(get_db)]
class Base(DeclarativeBase):
	pass