43 lines
1.0 KiB
Python
43 lines
1.0 KiB
Python
import logging
|
|
from os import getenv, path
|
|
|
|
from sqlmodel import Session, SQLModel, create_engine
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
SQLALCHEMY_DATABASE_PATH = getenv("SQLALCHEMY_DATABASE_PATH", "./gatekeeper.db")
|
|
SQLALCHEMY_DATABASE_URL = "sqlite:///" + SQLALCHEMY_DATABASE_PATH
|
|
|
|
engine = create_engine(SQLALCHEMY_DATABASE_URL)
|
|
|
|
|
|
def create_db_and_tables():
|
|
if not path.exists(SQLALCHEMY_DATABASE_PATH):
|
|
SQLModel.metadata.create_all(engine)
|
|
from alembic.config import Config
|
|
|
|
from alembic import command
|
|
|
|
alembic_cfg = Config("./alembic.ini")
|
|
alembic_cfg.attributes["sqlalchemy.url"] = SQLALCHEMY_DATABASE_URL
|
|
command.stamp(alembic_cfg, "head")
|
|
logger.info("Database created and tables initialized.")
|
|
else:
|
|
logger.info("Database already exists")
|
|
|
|
|
|
def get_session():
|
|
with Session(engine) as db:
|
|
yield db
|
|
|
|
|
|
def get_db_session():
|
|
return Session(engine)
|
|
|
|
|
|
def add_and_refresh(db: Session, obj):
|
|
db.add(obj)
|
|
db.commit()
|
|
db.refresh(obj)
|
|
return obj
|