Refactor Sessions to clarify their usage

DBSession has been renamed to HandlerSession.

The following exist now:
- HandlerSession: a DB session scoped to the current handler;
  should therefore only be used inside handlers.
- ThreadSession: a DB session scoped to the current thread;
  should be used outside of handlers, e.g. when firing off
  tasks in threads, or inside of services.

A special HandlerSession, with knowledge of the current user, and that
verifies permissions on commit, is available as `self.Session` on the
base handler. This is the *preferred way of accessing* HandlerSession.

To make it easier to get hold of the current engine, the Sessions each
have a `.engine` attribute: `HandlerSession.engine`.

This avoids having to access the engine as
`HandlerSession.session_factory.kw["bind"]`.

Author: stefanv

app/access.py

@@ -5,7 +5,7 @@
 from sqlalchemy.orm import joinedload
 
 from baselayer.app.custom_exceptions import AccessError  # noqa: F401
-from baselayer.app.models import DBSession, Role, Token, User  # noqa: F401
+from baselayer.app.models import HandlerSession, Role, Token, User  # noqa: F401
 
 
 def auth_or_token(method):
@@ -26,7 +26,7 @@ def wrapper(self, *args, **kwargs):
         token_header = self.request.headers.get("Authorization", None)
         if token_header is not None and token_header.startswith("token "):
             token_id = token_header.replace("token", "").strip()
-            with DBSession() as session:
+            with HandlerSession() as session:
                 token = session.scalars(
                     sa.select(Token)
                     .options(

app/handlers/base.py

@@ -22,7 +22,7 @@
 from ..env import load_env
 from ..flow import Flow
 from ..json_util import to_json
-from ..models import DBSession, User, VerifiedSession, bulk_verify, session_context_id
+from ..models import HandlerSession, User, VerifiedSession, bulk_verify, session_context_id
 
 env, cfg = load_env()
 log = make_log("basehandler")
@@ -49,7 +49,7 @@ def get_current_user(self):
         user_id = int(self.user_id())
         oauth_uid = self.get_secure_cookie("user_oauth_uid")
         if user_id and oauth_uid:
-            with DBSession() as session:
+            with HandlerSession() as session:
                 try:
                     user = session.scalars(
                         sqlalchemy.select(User).where(User.id == user_id)
@@ -74,7 +74,7 @@ def get_current_user(self):
             return None
 
     def login_user(self, user):
-        with DBSession() as session:
+        with HandlerSession() as session:
             try:
                 self.set_secure_cookie("user_id", str(user.id))
                 user = session.scalars(
@@ -120,7 +120,7 @@ def log_exception(self, typ=None, value=None, tb=None):
             )
 
     def on_finish(self):
-        DBSession.remove()
+        HandlerSession.remove()
 
 
 class BaseHandler(PSABaseHandler):
@@ -153,7 +153,7 @@ def Session(self):
             # must merge the user object with the current session
             # ref: https://docs.sqlalchemy.org/en/14/orm/session_basics.html#adding-new-or-existing-items
             session.add(self.current_user)
-            session.bind = DBSession.session_factory.kw["bind"]
+            session.bind = HandlerSession.engine
             yield session
 
     def verify_permissions(self):
@@ -164,20 +164,20 @@ def verify_permissions(self):
         """
 
         # get items to be inserted
-        new_rows = [row for row in DBSession().new]
+        new_rows = [row for row in HandlerSession().new]
 
         # get items to be updated
         updated_rows = [
-            row for row in DBSession().dirty if DBSession().is_modified(row)
+            row for row in HandlerSession().dirty if HandlerSession().is_modified(row)
         ]
 
         # get items to be deleted
-        deleted_rows = [row for row in DBSession().deleted]
+        deleted_rows = [row for row in HandlerSession().deleted]
 
         # get items that were read
         read_rows = [
             row
-            for row in set(DBSession().identity_map.values())
+            for row in set(HandlerSession().identity_map.values())
             - (set(updated_rows) | set(new_rows) | set(deleted_rows))
         ]
 
@@ -194,15 +194,15 @@ def verify_permissions(self):
         # update transaction state in DB, but don't commit yet. this updates
         # or adds rows in the database and uses their new state in joins,
         # for permissions checking purposes.
-        DBSession().flush()
+        HandlerSession().flush()
         bulk_verify("create", new_rows, self.current_user)
 
     def verify_and_commit(self):
         """Verify permissions on the current database session and commit if
         successful, otherwise raise an AccessError.
         """
         self.verify_permissions()
-        DBSession().commit()
+        HandlerSession().commit()
 
     def prepare(self):
         self.cfg = self.application.cfg
@@ -225,7 +225,7 @@ def prepare(self):
         N = 5
         for i in range(1, N + 1):
             try:
-                assert DBSession.session_factory.kw["bind"] is not None
+                assert HandlerSession.engine is not None
             except Exception as e:
                 if i == N:
                     raise e

app/model_util.py

@@ -21,15 +21,15 @@ def status(message):
     else:
         print(f"\r[✓] {message}")
     finally:
-        models.DBSession().commit()
+        models.HandlerSession().commit()
 
 
 def drop_tables():
-    conn = models.DBSession.session_factory.kw["bind"]
-    print(f"Dropping tables on database {conn.url.database}")
+    engine = models.HandlerSession.engine
+    print(f"Dropping tables on database {engine.url.database}")
     meta = sa.MetaData()
-    meta.reflect(bind=conn)
-    meta.drop_all(bind=conn)
+    meta.reflect(bind=engine)
+    meta.drop_all(bind=engine)
 
 
 def create_tables(retry=5, add=True):
@@ -45,17 +45,16 @@ def create_tables(retry=5, add=True):
         tables.
 
     """
-    conn = models.DBSession.session_factory.kw["bind"]
     tables = models.Base.metadata.sorted_tables
     if tables and not add:
         print("Existing tables found; not creating additional tables")
         return
 
     for i in range(1, retry + 1):
         try:
-            conn = models.DBSession.session_factory.kw["bind"]
-            print(f"Creating tables on database {conn.url.database}")
-            models.Base.metadata.create_all(conn)
+            engine = models.HandlerSession.engine
+            print(f"Creating tables on database {engine.url.database}")
+            models.Base.metadata.create_all(engine)
 
             table_list = ", ".join(list(models.Base.metadata.tables.keys()))
             print(f"Refreshed tables: {table_list}")

app/models.py

@@ -28,9 +28,24 @@
 log_database = cfg.get("log.database", False)
 log_database_pool = cfg.get("log.database_pool", False)
 
+
+# This provides one session per *thread*
+ThreadSession = scoped_session(sessionmaker())
+
+
+# This provides one session per *handler*
+# It is recommended to use the handler session via
+# self.Session, which has some knowledge
+# of the current user. See `handlers/base.py`
+#
+# DBSession has been renamed to HandlerSession
+# to make it clearer what it is doing.
+
+# We've renamed DBSession
+# It is not recommended to use DBSession directly;
 session_context_id = contextvars.ContextVar("request_id", default=None)
-# left here for backward compatibility:
 DBSession = scoped_session(sessionmaker(), scopefunc=session_context_id.get)
+HandlerSession = DBSession
 
 
 class _VerifiedSession(sa.orm.session.Session):
@@ -61,7 +76,7 @@ def __init__(self, user_or_token, **kwargs):
         or be generating an unverified session to only query
         the user with a certain id. Example:
 
-        with DBSession() as session:
+        with HandlerSession() as session:
             user = session.scalars(
                 sa.select(User).where(User.id == user_id)
             ).first()
@@ -159,7 +174,7 @@ def bulk_verify(mode, collection, accessor):
         ).subquery()
 
         inaccessible_row_ids = (
-            DBSession()
+            HandlerSession()
             .scalars(
                 sa.select(record_cls.id)
                 .outerjoin(
@@ -258,7 +273,7 @@ def init_db(
         "max_overflow": 10,
         "pool_recycle": 3600,
     }
-    conn = sa.create_engine(
+    engine = sa.create_engine(
         url,
         client_encoding="utf8",
         executemany_mode="values_plus_batch",
@@ -268,8 +283,15 @@ def init_db(
         **{**default_engine_args, **engine_args},
     )
 
-    DBSession.configure(bind=conn, autoflush=autoflush, future=True)
-    Base.metadata.bind = conn
+    HandlerSession.configure(bind=engine, autoflush=autoflush, future=True)
+    # Convenience attribute to easily access the engine, otherwise would need
+    # HandlerSession.session_factory.kw["bind"]
+    HandlerSession.engine = engine
+
+    ThreadSession.configure(bind=engine, autoflush=autoflush, future=True)
+    ThreadSession.engine = engine
+
+    Base.metadata.bind = engine
 
     return conn
 
@@ -478,8 +500,8 @@ def query_accessible_rows(self, cls, user_or_token, columns=None):
         """
         # return only selected columns if requested
         if columns is not None:
-            return DBSession().query(*columns).select_from(cls)
-        return DBSession().query(cls)
+            return HandlerSession().query(*columns).select_from(cls)
+        return HandlerSession().query(cls)
 
     def select_accessible_rows(self, cls, user_or_token, columns=None):
         """Construct a Select object that, when executed, returns the rows of a
@@ -571,9 +593,9 @@ def query_accessible_rows(self, cls, user_or_token, columns=None):
 
         # return only selected columns if requested
         if columns is not None:
-            query = DBSession().query(*columns).select_from(cls)
+            query = HandlerSession().query(*columns).select_from(cls)
         else:
-            query = DBSession().query(cls)
+            query = HandlerSession().query(cls)
 
         # traverse the relationship chain via sequential JOINs
         for relationship_name in self.relationship_names:
@@ -735,9 +757,9 @@ def query_accessible_rows(self, cls, user_or_token, columns=None):
 
         # return only selected columns if requested
         if columns is None:
-            base = DBSession().query(cls)
+            base = HandlerSession().query(cls)
         else:
-            base = DBSession().query(*columns).select_from(cls)
+            base = HandlerSession().query(*columns).select_from(cls)
 
         # ensure the target class has all the relationships referred to
         # in this instance
@@ -922,9 +944,9 @@ def query_accessible_rows(self, cls, user_or_token, columns=None):
 
         # retrieve specified columns if requested
         if columns is not None:
-            query = DBSession().query(*columns).select_from(cls)
+            query = HandlerSession().query(*columns).select_from(cls)
         else:
-            query = DBSession().query(cls)
+            query = HandlerSession().query(cls)
 
         # keep track of columns that will be null in the case of an unsuccessful
         # match for OR logic.
@@ -1076,9 +1098,9 @@ def query_accessible_rows(self, cls, user_or_token, columns=None):
         # otherwise, all records are inaccessible
         if columns is not None:
             return (
-                DBSession().query(*columns).select_from(cls).filter(sa.literal(False))
+                HandlerSession().query(*columns).select_from(cls).filter(sa.literal(False))
             )
-        return DBSession().query(cls).filter(sa.literal(False))
+        return HandlerSession().query(cls).filter(sa.literal(False))
 
     def select_accessible_rows(self, cls, user_or_token, columns=None):
         """Construct a Select object that, when executed, returns the rows of a
@@ -1148,7 +1170,7 @@ def __init__(self, query_or_query_generator):
 
             Query (SQLA 1.4):
             >>>> CustomUserAccessControl(
-                DBSession().query(Department).join(Employee).group_by(
+                HandlerSession().query(Department).join(Employee).group_by(
                     Department.id
                 ).having(sa.func.bool_and(Employee.is_manager.is_(True)))
             )
@@ -1166,8 +1188,8 @@ def __init__(self, query_or_query_generator):
             Query (SQLA 1.4):
             >>>> def access_logic(cls, user_or_token):
              ...      if user_or_token.is_system_admin:
-             ...         return DBSession().query(cls)
-             ...      return DBSession().query(cls).join(Employee).group_by(
+             ...         return HandlerSession().query(cls)
+             ...      return HandlerSession().query(cls).join(Employee).group_by(
              ...             cls.id
              ...      ).having(sa.func.bool_and(Employee.is_manager.is_(True)))
             >>>> CustomUserAccessControl(access_logic)
@@ -1303,7 +1325,7 @@ def is_accessible_by(self, user_or_token, mode="read"):
 
         # Query for the value of the access_func for this particular record and
         # return the result.
-        result = DBSession().execute(stmt).scalar_one() > 0
+        result = HandlerSession().execute(stmt).scalar_one() > 0
         if result is None:
             result = False
 
@@ -1353,7 +1375,7 @@ def get_if_accessible_by(
 
         # TODO: vectorize this
         for pk in standardized:
-            instance = DBSession().query(cls).options(options).get(pk.item())
+            instance = HandlerSession().query(cls).options(options).get(pk.item())
             if instance is None or not instance.is_accessible_by(
                 user_or_token, mode=mode
             ):
@@ -1467,7 +1489,7 @@ def get(
         standardized = np.atleast_1d(id_or_list)
         result = []
 
-        with DBSession() as session:
+        with HandlerSession() as session:
             # TODO: vectorize this
             for pk in standardized:
                 if options:
@@ -1519,7 +1541,7 @@ def get_all(
             If columns is specified, will return a list of tuples
             containing the data from each column requested.
         """
-        with DBSession() as session:
+        with HandlerSession() as session:
             stmt = cls.select(user_or_token, mode, options, columns)
             values = session.scalars(stmt).all()
 
@@ -1566,7 +1588,7 @@ def select(
             stmt = stmt.options(option)
         return stmt
 
-    query = DBSession.query_property()
+    query = HandlerSession.query_property()
 
     id = sa.Column(
         sa.Integer,
@@ -1608,8 +1630,8 @@ def __repr__(self):
     def to_dict(self):
         """Serialize this object to a Python dictionary."""
         if sa.inspection.inspect(self).expired:
-            self = DBSession().merge(self)
-            DBSession().refresh(self)
+            self = HandlerSession().merge(self)
+            HandlerSession().refresh(self)
         return {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
 
     @classmethod
@@ -1632,7 +1654,7 @@ def get_if_readable_by(cls, ident, user_or_token, options=[]):
         obj : baselayer.app.models.Base
            The requested entity.
         """
-        obj = DBSession().query(cls).options(options).get(ident)
+        obj = HandlerSession().query(cls).options(options).get(ident)
 
         if obj is not None and not obj.is_readable_by(user_or_token):
             raise AccessError("Insufficient permissions.")
@@ -1659,7 +1681,7 @@ def is_readable_by(self, user_or_token):
     def create_or_get(cls, id):
         """Return a new `cls` if an instance with the specified primary key
         does not exist, else return the existing instance."""
-        obj = DBSession().query(cls).get(id)
+        obj = HandlerSession().query(cls).get(id)
         if obj is not None:
             return obj
         else:
@@ -1876,7 +1898,7 @@ class User(Base):
     role_ids = association_proxy(
         "roles",
         "id",
-        creator=lambda r: DBSession().query(Role).get(r),
+        creator=lambda r: HandlerSession().query(Role).get(r),
     )
     tokens = relationship(
         "Token",
@@ -1982,7 +2004,7 @@ class Token(Base):
         lazy="selectin",
     )
     acl_ids = association_proxy(
-        "acls", "id", creator=lambda acl: DBSession().query(ACL).get(acl)
+        "acls", "id", creator=lambda acl: HandlerSession().query(ACL).get(acl)
     )
     permissions = acl_ids