Free Query Tool autocomplete objects when an idle session's connections are released. #10612 (#10625)

web/pgadmin/tools/sqleditor/__init__.py

@@ -47,6 +47,7 @@ from pgadmin.utils import get_storage_directory
 from pgadmin.utils.ajax import make_json_response, bad_request, \
     success_return, internal_server_error, service_unavailable, gone
 from pgadmin.utils.driver import get_driver
+from pgadmin.utils.driver.abstract import BaseDriver
 from pgadmin.utils.crypto import encrypt
 from pgadmin.utils.master_password import get_crypt_key
 from pgadmin.utils.exception import ConnectionLost, SSHTunnelConnectionLost, \
@@ -100,6 +101,20 @@ def _auto_complete_key(trans_id):
     return session.sid, int(trans_id)
 
 
+def _release_session_auto_complete(session_id):
+    """
+    Drop every SQLAutoComplete object belonging to the given session. Called
+    by the driver when it releases the connections of a session that has
+    expired without logging out, as their connections are no longer usable.
+    """
+    for key in [k for k in list(auto_complete_objects)
+                if k[0] == session_id]:
+        auto_complete_objects.pop(key, None)
+
+
+BaseDriver.register_session_gc_callback(_release_session_auto_complete)
+
+
 class SqlEditorModule(PgAdminModule):
     """
     class SqlEditorModule(PgAdminModule)

web/pgadmin/tools/sqleditor/tests/test_autocomplete_session_gc.py

@@ -0,0 +1,87 @@
+##########################################################################
+#
+# pgAdmin 4 - PostgreSQL Tools
+#
+# Copyright (C) 2013 - 2026, The pgAdmin Development Team
+# This software is released under the PostgreSQL Licence
+#
+##########################################################################
+
+import datetime
+from unittest.mock import MagicMock, patch
+
+from config import PG_DEFAULT_DRIVER
+from pgadmin.utils.route import BaseTestGenerator
+from pgadmin.utils.driver import get_driver
+from pgadmin.utils.driver.abstract import BaseDriver
+import pgadmin.utils.driver.psycopg3 as psycopg3_driver
+from pgadmin.utils.driver.psycopg3.server_manager import ServerManager
+import pgadmin.tools.sqleditor as sqleditor
+
+
+class TestAutocompleteSessionGC(BaseTestGenerator):
+    """
+    When the driver releases the connections of a session that has gone
+    idle without logging out, that session's cached SQLAutoComplete objects
+    must be dropped too, whilst other sessions' objects are kept.
+    """
+    scenarios = [
+        ('Autocomplete objects are dropped for expired sessions', dict()),
+    ]
+
+    def setUp(self):
+        self.saved_objects = dict(sqleditor.auto_complete_objects)
+        sqleditor.auto_complete_objects.clear()
+
+    def tearDown(self):
+        sqleditor.auto_complete_objects.clear()
+        sqleditor.auto_complete_objects.update(self.saved_objects)
+
+    def runTest(self):
+        now = datetime.datetime.now()
+        expired_mgr = MagicMock(spec=ServerManager)
+        active_mgr = MagicMock(spec=ServerManager)
+        managers = {
+            'expired': {'pinged': now - datetime.timedelta(days=1),
+                        '1': expired_mgr},
+            'active': {'pinged': now, '1': active_mgr},
+        }
+        sqleditor.auto_complete_objects.update({
+            ('expired', 1): MagicMock(),
+            ('expired', 2): MagicMock(),
+            ('active', 1): MagicMock(),
+        })
+
+        callback = MagicMock()
+
+        with self.app.test_request_context():
+            driver = get_driver(PG_DEFAULT_DRIVER)
+            with patch.object(driver, 'managers', managers), \
+                    patch.object(psycopg3_driver, 'session',
+                                 MagicMock(sid='current')), \
+                    patch.object(BaseDriver, '_session_gc_callbacks',
+                                 [sqleditor._release_session_auto_complete,
+                                  callback]):
+                driver.gc_timeout()
+
+                expired_mgr.release.assert_called_once_with()
+                active_mgr.release.assert_not_called()
+                self.assertEqual(list(sqleditor.auto_complete_objects),
+                                 [('active', 1)])
+                callback.assert_called_once_with('expired')
+
+                # Later sweeps still release the expired session's
+                # connections, but do not run the callbacks again.
+                driver.gc_timeout()
+                self.assertEqual(expired_mgr.release.call_count, 2)
+                callback.assert_called_once_with('expired')
+
+                # Once the session is active again, the callbacks run the
+                # next time it goes idle.
+                with patch.object(psycopg3_driver, 'session',
+                                  MagicMock(sid='expired')):
+                    driver.gc_timeout()
+                managers['expired']['pinged'] = \
+                    now - datetime.timedelta(days=1)
+                driver.gc_timeout()
+                self.assertEqual(callback.call_count, 2)

web/pgadmin/utils/driver/abstract.py

@@ -10,6 +10,7 @@
 """Implement the Base class for Driver and Connection"""
 
 from abc import ABCMeta, abstractmethod, abstractproperty
+from flask import current_app
 from .registry import DriverRegistry
 
 
@@ -44,8 +45,32 @@ class BaseDriver(metaclass=DriverRegistry):
     - Implement this function to release the connections assigned in the
       session, which has not been pinged from more than the idle timeout
       configuration.
+
+    Callbacks registered with register_session_gc_callback() are called
+    with the session id whenever gc_timeout() releases that session's
+    connections, so modules caching per-session state that holds those
+    connections can drop it.
     """
 
+    _session_gc_callbacks = []
+
+    @classmethod
+    def register_session_gc_callback(cls, callback):
+        """
+        Register callback(session_id) to be run when the connections of an
+        idle session are released by gc_timeout().
+        """
+        if callback not in BaseDriver._session_gc_callbacks:
+            BaseDriver._session_gc_callbacks.append(callback)
+
+    @staticmethod
+    def run_session_gc_callbacks(session_id):
+        for callback in BaseDriver._session_gc_callbacks:
+            try:
+                callback(session_id)
+            except Exception as e:
+                current_app.logger.exception(e)
+
     @property
     @abstractmethod
     def version(cls):

web/pgadmin/utils/driver/psycopg3/__init__.py

@@ -320,6 +320,9 @@ class Driver(BaseDriver):
                                 manager_source)
 
         managers['pinged'] = datetime.datetime.now()
+        # The session is active again, so gc_timeout() should run the
+        # session GC callbacks the next time it goes idle.
+        managers.pop('gc_callbacks_run', None)
         if str(sid) not in managers:
             # server_data was already access-checked above;
             # it cannot be None at this point.
@@ -450,6 +453,7 @@ class Driver(BaseDriver):
 
             if sess == session.sid:
                 sess_mgr['pinged'] = curr_time
+                sess_mgr.pop('gc_callbacks_run', None)
                 continue
             if curr_time - sess_mgr['pinged'] >= session_idle_timeout:
                 for mgr in [
@@ -457,6 +461,11 @@ class Driver(BaseDriver):
                                                                ServerManager)
                 ]:
                     mgr.release()
+                # Expired sessions stay in self.managers, so only run the
+                # callbacks once per idle period rather than on every sweep.
+                if not sess_mgr.get('gc_callbacks_run'):
+                    self.run_session_gc_callbacks(sess)
+                    sess_mgr['gc_callbacks_run'] = True
 
     def gc_own(self):
         """