diff --git a/web/pgadmin/tools/sqleditor/__init__.py b/web/pgadmin/tools/sqleditor/__init__.py index 54d7658e6..37395e294 100644 --- a/web/pgadmin/tools/sqleditor/__init__.py +++ b/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) diff --git a/web/pgadmin/tools/sqleditor/tests/test_autocomplete_session_gc.py b/web/pgadmin/tools/sqleditor/tests/test_autocomplete_session_gc.py new file mode 100644 index 000000000..eeaffa3ed --- /dev/null +++ b/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) diff --git a/web/pgadmin/utils/driver/abstract.py b/web/pgadmin/utils/driver/abstract.py index 413723228..44223ac0a 100644 --- a/web/pgadmin/utils/driver/abstract.py +++ b/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): diff --git a/web/pgadmin/utils/driver/psycopg3/__init__.py b/web/pgadmin/utils/driver/psycopg3/__init__.py index aabc6da2d..2cca18705 100644 --- a/web/pgadmin/utils/driver/psycopg3/__init__.py +++ b/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): """