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