fix: Refactor DB session logic
@@ -9,12 +9,15 @@ import sys
import logging
from sqlalchemy import create_engine, exc, func
-from sqlalchemy.orm import Session as _Session
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import scoped_session, sessionmaker
from owtf.db.model_base import Model
from owtf.settings import DATABASE_IP, DATABASE_NAME, DATABASE_PASS, DATABASE_USER, DATABASE_PORT
+DB_URI = "postgresql+psycopg2://{}:{}@{}:{}/{}".format(
+ DATABASE_USER, DATABASE_PASS, DATABASE_IP, DATABASE_PORT, DATABASE_NAME
+)
+
def get_count(q):
count_q = q.statement.with_only_columns([func.count()]).order_by(None)
@@ -44,34 +47,16 @@ def flush_transaction(method):
def get_db_engine():
try:
- engine = create_engine(
- "postgresql+psycopg2://{}:{}@{}:{}/{}".format(
- DATABASE_USER, DATABASE_PASS, DATABASE_IP, DATABASE_PORT, DATABASE_NAME
- ),
- pool_recycle=120,
- )
+ engine = create_engine(DB_URI, pool_recycle=120)
Model.metadata.create_all(engine)
return engine
except exc.OperationalError as e:
- logging.error("Could not create database engine - Exception occured\n%s", str(e))
+ logging.error("Could not create engine - Exception occured\n%s", str(e))
sys.exit(1)
def get_scoped_session():
- Session.configure(bind=get_db_engine())
- return Session()
-
-
-class Session(_Session):
- """ Custom session meant to utilize add on the model.
- This Session overrides the add/add_all methods to prevent them
- from being used. This is to for using the add methods on the
- models themselves where overriding is available.
- """
-
- _add = _Session.add
- _add_all = _Session.add_all
- _delete = _Session.delete
+ return scoped_session(Session)
-Session = sessionmaker(class_=Session)
+Session = sessionmaker(bind=get_db_engine())
