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

31 lines
685 B
Python
Raw Normal View History

2022-03-18 18:22:17 +00:00
from typing import Generator
from sqlalchemy.engine import Engine
from sqlalchemy.orm import sessionmaker, Session
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()