| 1 | # ! NEW WITH MYSQL ! # |
| 2 | |
| 3 | from contextlib import asynccontextmanager |
| 4 | from contextlib import contextmanager |
| 5 | |
| 6 | # from settings import SQLALCHEMY_DATABASE_URI |
| 7 | from pathlib import Path |
| 8 | |
| 9 | from dotenv import load_dotenv |
| 10 | from environs import Env |
| 11 | from loguru import logger |
| 12 | from sqlalchemy.ext.asyncio import AsyncSession |
| 13 | from sqlalchemy.ext.asyncio import create_async_engine |
| 14 | from sqlalchemy.orm import sessionmaker |
| 15 | from sqlmodel import Session |
| 16 | from sqlmodel import create_engine |
| 17 | |
| 18 | # Repo root is four parents up from this file (backend/app/db/db_session.py). |
| 19 | # environs >= 14 stopped mutating os.environ from read_env(); use python-dotenv so |
| 20 | # os.environ.get() callsites (e.g. _load_jwt_secret) see values from .env. |
| 21 | _DOTENV_PATH = Path(__file__).parent.parent.parent.parent / ".env" |
| 22 | load_dotenv(_DOTENV_PATH) |
| 23 | env = Env() |
| 24 | logger.info(f"Loading environment from {_DOTENV_PATH}") |
| 25 | |
| 26 | db_user = env.str("MYSQL_USER", default="copilot") |
| 27 | db_password = env.str("MYSQL_PASSWORD") |
| 28 | db_root_password = env.str("MYSQL_ROOT_PASSWORD") |
| 29 | db_url = env.str("MYSQL_URL", default="copilot-mysql") |
| 30 | |
| 31 | logger.info(f"DB User: {db_user} and password: {db_password}") |
| 32 | |
| 33 | # Update the SQLALCHEMY_DATABASE_URI to a MySQL compatible one in settings.py |
| 34 | # For this example, let's assume it has been updated. copilot-mysql |
| 35 | SQLALCHEMY_DATABASE_URI_NO_DB = f"mysql+pymysql://root:{db_root_password}@{db_url}" |
| 36 | SQLALCHEMY_DATABASE_URI = f"mysql+aiomysql://{db_user}:{db_password}@{db_url}/copilot" |
| 37 | |
| 38 | |
| 39 | session = "placeholder" |
| 40 | |
| 41 | # Create async engine for MySQL using aiomysql |
| 42 | async_engine = create_async_engine( |
| 43 | SQLALCHEMY_DATABASE_URI, |
| 44 | echo=False, |
| 45 | # Additional MySQL-specific options can be set here if needed |
| 46 | ) |
| 47 | # If you still need sync sessions for some operations, set it appropriately |
| 48 | # This would typically require a different sync driver since SQLAlchemy doesn't use aiomysql for sync operations |
| 49 | # ! THIS IS USED BY THE SCHEDULER ! # |
| 50 | sync_engine = create_engine( |
| 51 | SQLALCHEMY_DATABASE_URI.replace("+aiomysql", "+pymysql"), |
| 52 | echo=False, |
| 53 | # Additional MySQL-specific options can be set here if needed |
| 54 | ) |
| 55 | |
| 56 | # Create a configured "AsyncSession" class |
| 57 | AsyncSessionLocal = sessionmaker(bind=async_engine, class_=AsyncSession, expire_on_commit=False) |
| 58 | SyncSessionLocal = sessionmaker(bind=sync_engine, class_=Session, expire_on_commit=False) |
| 59 | |
| 60 | |
| 61 | @asynccontextmanager |
| 62 | async def get_db_session(): |
| 63 | async with AsyncSessionLocal() as session: |
| 64 | logger.info("DB session created") |
| 65 | try: |
| 66 | yield session |
| 67 | except Exception as e: |
| 68 | logger.error(f"Error during DB session: {e}") |
| 69 | await session.rollback() |
| 70 | raise e |
| 71 | finally: |
| 72 | logger.info("Closing DB session") |
| 73 | await session.close() |
| 74 | |
| 75 | |
| 76 | @contextmanager |
| 77 | def get_sync_db_session(): |
| 78 | session = SyncSessionLocal() |
| 79 | logger.info("Sync DB session created") |
| 80 | try: |
| 81 | yield session |
| 82 | except Exception as e: |
| 83 | logger.error(f"Error during sync DB session: {e}") |
| 84 | session.rollback() |
| 85 | raise e |
| 86 | finally: |
| 87 | logger.info("Closing sync DB session") |
| 88 | session.close() |
| 89 | |
| 90 | |
| 91 | @asynccontextmanager |
| 92 | async def get_session(): |
| 93 | async with get_db_session() as session: |
| 94 | yield session |
| 95 | |
| 96 | |
| 97 | async def get_db(): |
| 98 | async with get_session() as session: |
| 99 | yield session |