Fix session handling for db backend using Flask-SQLAlchemy

datagateway_api/common/database/helpers.py

@@ -17,7 +17,7 @@
     INVESTIGATIONINSTRUMENT,
     SESSION,
 )
-from datagateway_api.common.database.session_manager import session_manager
+from datagateway_api.common.database.session_manager import db
 from datagateway_api.common.exceptions import (
     AuthenticationError,
     BadRequestError,
@@ -41,12 +41,11 @@ def requires_session_id(method):
     @wraps(method)
     def wrapper_requires_session(*args, **kwargs):
         log.info(" Authenticating consumer")
-        session = session_manager.get_icat_db_session()
+        session = db.session()
         query = session.query(SESSION).filter(SESSION.ID == args[1]).first()
         if query is not None:
             log.info(" Closing DB session")
             session.close()
-            session.close()
             log.info(" Consumer authenticated")
             return method(*args, **kwargs)
         else:
@@ -66,7 +65,7 @@ class Query(ABC):
 
     @abstractmethod
     def __init__(self, table):
-        self.session = session_manager.get_icat_db_session()
+        self.session = db.session
         self.table = table
         self.base_query = self.session.query(table)
 
@@ -86,7 +85,12 @@ def commit_changes(self):
         Commits all changes to the database and closes the session
         """
         log.info(" Committing changes to %s", self.table)
-        self.session.commit()
+        try:
+            self.session.commit()
+        except Exception as e:
+            log.error(f"Error whilst committing changes to {self.table}, rolling back")
+            self.session.rollback()
+            raise e
 
 
 class CountQuery(Query):

datagateway_api/common/database/session_manager.py

@@ -2,6 +2,7 @@
 from sqlalchemy.orm import scoped_session, sessionmaker
 from sqlalchemy.pool import QueuePool
 
+from flask_sqlalchemy import SQLAlchemy
 from datagateway_api.common.constants import Constants
 
 engine = create_engine(
@@ -24,3 +25,4 @@ def get_icat_db_session(self):
 
 
 session_manager = SessionManager()
+db = SQLAlchemy()

datagateway_api/src/main.py

@@ -7,10 +7,13 @@
 from flask_cors import CORS
 from flask_restful import Api
 from flask_swagger_ui import get_swaggerui_blueprint
+from flask_sqlalchemy import SQLAlchemy
 
 from datagateway_api.common.backends import create_backend
 from datagateway_api.common.config import config
 from datagateway_api.common.logger_setup import setup_logger
+from datagateway_api.common.database.session_manager import db
+from datagateway_api.common.constants import Constants
 from datagateway_api.src.resources.entities.entity_endpoint import (
     get_count_endpoint,
     get_endpoint,
@@ -63,6 +66,8 @@ def create_app_infrastructure(flask_app):
     CORS(flask_app)
     flask_app.url_map.strict_slashes = False
     api = CustomErrorHandledApi(flask_app)
+    app.config["SQLALCHEMY_DATABASE_URI"] = Constants.DATABASE_URL
+    db.init_app(app)
 
     initialise_spec(spec)
 
@@ -168,10 +173,10 @@ def specs():
     resp.mimetype = "application/json"
     return resp
 
+api, spec = create_app_infrastructure(app)
+create_api_endpoints(app, api, spec)
 
 if __name__ == "__main__":
-    api, spec = create_app_infrastructure(app)
-    create_api_endpoints(app, api, spec)
     openapi_config(spec)
     app.run(
         host=config.get_host(), port=config.get_port(), debug=config.is_debug_mode(),

pyproject.toml

@@ -19,6 +19,7 @@ flask-swagger-ui = "3.25.0"
 PyYAML = "5.3.1"
 python-icat = "0.17.0"
 suds-community = "^0.8.4"
+Flask-SQLAlchemy = "^2.4.4"
 
 [tool.poetry.dev-dependencies]
 pip-tools = "5.3.1"