36 lines
913 B
Python
36 lines
913 B
Python
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
from os import getenv, path
|
|
|
|
from sqlmodel import Session, SQLModel, create_engine
|
|
|
|
SQLALCHEMY_DATABASE_URL = getenv("SQLALCHEMY_DATABASE_URL", "sqlite:///./gatekeeper.db")
|
|
|
|
engine = create_engine(SQLALCHEMY_DATABASE_URL)
|
|
|
|
def create_db_and_tables():
|
|
if not path.exists(SQLALCHEMY_DATABASE_URL):
|
|
SQLModel.metadata.create_all(engine)
|
|
from alembic.config import Config
|
|
|
|
from alembic import command
|
|
alembic_cfg = Config("./alembic.ini")
|
|
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 |