From f3f41c1edc481c4316f6be76e03dd162ae9a3ced Mon Sep 17 00:00:00 2001
From: Alexander Goscinski <alexander.goscinski@posteo.de>
Date: Sun, 2 Feb 2025 09:22:24 +0100
Subject: [PATCH] less aggressively close sessions

---
 src/aiida/storage/psql_dos/backend.py  | 12 +++++++-----
 tests/storage/sqlite/test_container.py | 21 +++++++++++++++++++++
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/src/aiida/storage/psql_dos/backend.py b/src/aiida/storage/psql_dos/backend.py
index b0d2dc813a..989def2f3d 100644
--- a/src/aiida/storage/psql_dos/backend.py
+++ b/src/aiida/storage/psql_dos/backend.py
@@ -186,16 +186,18 @@ def close(self) -> None:
             return  # the instance is already closed, and so this is a no-op
         # close the connection
 
+        # PR QUESTION: can we have multiple sessons? Then we should close all of them
+        #from sqlalchemy.orm import close_all_sessions
+        #close_all_sessions()
+
         engine = self._session_factory.bind
-        if engine is not None:
-            engine.dispose()  # type: ignore[union-attr]
         self._session_factory.expunge_all()
         self._session_factory.close()
         self._session_factory = None
 
-        # Without this, sqlalchemy keeps a weakref to a session
-        # in sqlalchemy.orm.session._sessions
-        gc.collect()
+        # IMPORTANT: Dispose engine only after sessions have been closed
+        if engine is not None:
+            engine.dispose()  # type: ignore[union-attr]
 
     def _clear(self) -> None:
         from aiida.storage.psql_dos.models.settings import DbSetting
diff --git a/tests/storage/sqlite/test_container.py b/tests/storage/sqlite/test_container.py
index e0ef1e7361..0a3bbc2e19 100644
--- a/tests/storage/sqlite/test_container.py
+++ b/tests/storage/sqlite/test_container.py
@@ -2,6 +2,26 @@
 import psutil, os
 
 def test_file_descriptor_closed(aiida_profile):
+    """Checks if the number of open file descriptors change during a reset."""
+    def list_open_file_descriptors():
+        process = psutil.Process(os.getpid())
+        return process.open_files()
+    # We have some connections active due to aiida profile during the first
+    # reset these are destroyed. We check the second time changes the number of
+    # open file descriptors.
+    from aiida.manage import get_manager
+    storage_backend = get_manager().get_profile_storage()
+    migrator_context = storage_backend.migrator_context 
+    open_file_descriptors_before = list_open_file_descriptors()
+    with migrator_context(aiida_profile) as migrator:
+        migrator.initialise_repository()
+        migrator.reset_repository()
+    open_file_descriptors_after = list_open_file_descriptors()
+    assert len(open_file_descriptors_before) == len(open_file_descriptors_after), f"Before these file descriptor were open:\n{open_file_descriptors_before}\n Now these are open:\n{open_file_descriptors_after}"
+
+
+# PR COMMENT this is just a sanity check for me, I don' think that the test should be included in final PR
+def test_reset_storage_file_descriptor_closed(aiida_profile):
     """Checks if the number of open file descriptors change during a reset."""
     def list_open_file_descriptors():
         process = psutil.Process(os.getpid())
@@ -11,6 +31,7 @@ def list_open_file_descriptors():
     # open file descriptors.
     # TODO The fix should keep the existing connections alive and just reuse them
     #      then one does not need to do two resets.
+    from aiida.manage import get_manager
     aiida_profile.reset_storage() 
     open_file_descriptors_before = list_open_file_descriptors()
     aiida_profile.reset_storage()