Ensure sessions are cleaned up when using settings

conda-store-server/conda_store_server/_internal/settings.py

@@ -2,7 +2,8 @@
 # Use of this source code is governed by a BSD-style
 # license that can be found in the LICENSE file.
 
-from typing import Any, Dict
+from typing import Any, Dict, Callable
+import functools
 
 from sqlalchemy.orm import Session
 import pydantic
@@ -16,6 +17,15 @@ def __init__(self, db: Session, deployment_default: schema.Settings):
         self.db = db
         self.deployment_default = deployment_default.model_dump()
 
+    def _ensure_closed_session(func: Callable):
+        @functools.wraps(func)
+        def wrapper(self, *args, **kwargs):
+            result = func(self, *args, **kwargs)
+            self.db.close()
+            return result
+        return wrapper
+
+    @_ensure_closed_session
     def set_settings(
         self,
         namespace: str = None,
@@ -65,6 +75,7 @@ def set_settings(
 
         api.set_kvstore_key_values(self.db, prefix, data)
 
+    @_ensure_closed_session
     def get_settings(
         self, namespace: str = None, environment_name: str = None
     ) -> schema.Settings:

conda-store-server/tests/_internal/test_settings.py

@@ -2,6 +2,8 @@
 # Use of this source code is governed by a BSD-style
 # license that can be found in the LICENSE file.
 
+from unittest import mock
+
 import pytest
 
 from conda_store_server import api
@@ -51,6 +53,27 @@ def settings(db) -> Settings:
     )
 
 
+def test_ensure_session_is_closed(settings: Settings):
+    # run a query against the db to start a transaction
+    settings.get_settings()
+    # ensure that the settings object cleans up it's transaction
+    assert settings.db.in_transaction() == False
+
+
+@mock.patch("conda_store_server.api.get_kvstore_key_values")
+def test_ensure_session_is_closed_on_error(mock_get_kvstore_key_values, settings: Settings):
+    mock_get_kvstore_key_values.side_effect = Exception
+
+    # run a query that will raise an exception
+    try:
+        settings.get_settings()
+    except:
+        pass
+
+    # ensure that the settings object cleans up it's transaction
+    assert settings.db.in_transaction() == False
+
+
 def test_get_settings_default(settings: Settings):
     test_settings = settings.get_settings()