From a126f107162936de5f72f24b91b93ce15a45e62b Mon Sep 17 00:00:00 2001
From: mvdbeek <m.vandenbeek@gmail.com>
Date: Thu, 1 Aug 2024 09:31:33 +0200
Subject: [PATCH] Clean up factory handling

---
 lib/galaxy/web/framework/base.py      | 11 +++++------
 lib/galaxy/webapps/base/webapp.py     |  3 +--
 lib/galaxy/webapps/galaxy/buildapp.py |  8 ++++++++
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/lib/galaxy/web/framework/base.py b/lib/galaxy/web/framework/base.py
index f8ed3a4fb06f..8a716b2146ab 100644
--- a/lib/galaxy/web/framework/base.py
+++ b/lib/galaxy/web/framework/base.py
@@ -104,6 +104,7 @@ def __init__(self):
         self.transaction_factory = DefaultWebTransaction
         # Set if trace logging is enabled
         self.trace_logger = None
+        self.session_factories = []
 
     def add_ui_controller(self, controller_name, controller):
         """
@@ -170,14 +171,12 @@ def __call__(self, environ, start_response):
         path_info = environ.get("PATH_INFO", "")
 
         try:
-            self._model.set_request_id(request_id)  # Start SQLAlchemy session scope
-            if self._install_model:
-                self._install_model.set_request_id(request_id)
+            for session_factory in self.session_factories:
+                session_factory.set_request_id(request_id)  # Start SQLAlchemy session scope
             return self.handle_request(request_id, path_info, environ, start_response)
         finally:
-            self._model.unset_request_id(request_id)  # End SQLAlchemy session scope
-            if self._install_model:
-                self._install_model.unset_request_id(request_id)
+            for session_factory in self.session_factories:
+                session_factory.unset_request_id(request_id)  # End SQLAlchemy session scope
             self.trace(message="Handle request finished")
             if self.trace_logger:
                 self.trace_logger.context_remove("request_id")
diff --git a/lib/galaxy/webapps/base/webapp.py b/lib/galaxy/webapps/base/webapp.py
index 6fdcff60cd24..6ba4b4745a34 100644
--- a/lib/galaxy/webapps/base/webapp.py
+++ b/lib/galaxy/webapps/base/webapp.py
@@ -120,8 +120,7 @@ def __init__(
 
         # We need this to set the REQUEST_ID contextvar in model.base *BEFORE* a GalaxyWebTransaction is created.
         # This will ensure a SQLAlchemy session is request-scoped for legacy (non-fastapi) endpoints.
-        self._model = galaxy_app.model
-        self._install_model = getattr(galaxy_app, "install_model", None)
+        self.session_factories.append(galaxy_app.model)
 
     def build_apispec(self):
         """
diff --git a/lib/galaxy/webapps/galaxy/buildapp.py b/lib/galaxy/webapps/galaxy/buildapp.py
index 83f65a90e97e..ad76bc678ef3 100644
--- a/lib/galaxy/webapps/galaxy/buildapp.py
+++ b/lib/galaxy/webapps/galaxy/buildapp.py
@@ -20,6 +20,7 @@
 import galaxy.webapps.base.webapp
 from galaxy import util
 from galaxy.security.validate_user_input import VALID_PUBLICNAME_RE
+from galaxy.structured_app import MinimalApp
 from galaxy.util import asbool
 from galaxy.util.properties import load_app_properties
 from galaxy.web.framework.middleware.error import ErrorMiddleware
@@ -34,6 +35,13 @@
 class GalaxyWebApplication(galaxy.webapps.base.webapp.WebApplication):
     injection_aware = True
 
+    def __init__(
+        self, galaxy_app: MinimalApp, session_cookie: str = "galaxysession", name: Optional[str] = None
+    ) -> None:
+        super().__init__(galaxy_app, session_cookie, name)
+        self.session_factories.append(galaxy_app.install_model)
+
+
 def app_factory(*args, **kwargs):
     """
     Return a wsgi application serving the root object