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


class Connection:
    engine: Engine | None = None
    session_local: sessionmaker | None = None

    @classmethod
    def connect(cls, engine: Engine) -> None:
        cls.engine = engine
        cls.session_local = sessionmaker(
            autocommit=False, autoflush=False, bind=engine,
        )
        ORMBaseModel.metadata.create_all(bind=engine)

    @classmethod
    async def get(cls) -> Generator[Session | None, None, None]:
        if cls.session_local is None:
            yield None

        else:
            db = cls.session_local()
            try:
                yield db
            finally:
                db.close()


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()