PillarTool/pillar_tool/db/database.py

34 lines
729 B
Python

from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker, declarative_base
from sqlalchemy.sql.annotation import Annotated
from pillar_tool.util import config
cfg = config()
SessionLocal = sessionmaker(
autocommit=False,
autoflush=False,
bind=create_engine(
url=f"postgresql+psycopg2://{cfg.db.user}:{cfg.db.password}@{cfg.host}:{cfg.port}/{cfg.database}"
)
)
Base = declarative_base()
def get_connection():
session = SessionLocal()
# noinspection PyBroadException
try:
yield session
except:
session.rollback()
session.close()
else:
session.commit()
session.close()
DB = Annotated[Session, get_connection()]