kiwi-vpn/api/kiwi_vpn_api/db/connection.py

from typing import Generator

from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, sessionmaker

from .models import ORMBaseModel

ENGINE: Engine | None = None
SESSION_LOCAL: sessionmaker | None = None


def reconnect(engine: Engine) -> None:
    global ENGINE, SESSION_LOCAL

    ENGINE = engine
    SESSION_LOCAL = sessionmaker(
        autocommit=False, autoflush=False, bind=engine,
    )
    ORMBaseModel.metadata.create_all(bind=engine)


async def get() -> Generator[Session | None, None, None]:
    if SESSION_LOCAL is None:
        yield None

    else:
        db = SESSION_LOCAL()
        try:
            yield db
        finally:
            db.close()