diff --git a/web/pgadmin/tools/sqleditor/__init__.py b/web/pgadmin/tools/sqleditor/__init__.py index 5768393f9..1293eab94 100644 --- a/web/pgadmin/tools/sqleditor/__init__.py +++ b/web/pgadmin/tools/sqleditor/__init__.py @@ -84,9 +84,22 @@ MODULE_NAME = 'sqleditor' TRANSACTION_STATUS_CHECK_FAILED = gettext("Transaction status check failed.") _NODES_SQL = 'nodes.sql' sqleditor_close_session_lock = Lock() +# SQLAutoComplete objects, keyed by _auto_complete_key(). Each one holds a +# live connection belonging to the session that created it, so the key must +# include the session as well as the client-chosen transaction id; otherwise +# another session using the same id would be given someone else's connection. auto_complete_objects = dict() +def _auto_complete_key(trans_id): + """ + Return the key for the current session's SQLAutoComplete object for the + given transaction id. Connection managers are held per session (see + connection_manager() in the driver), so the cache is scoped the same way. + """ + return session.sid, int(trans_id) + + class SqlEditorModule(PgAdminModule): """ class SqlEditorModule(PgAdminModule) @@ -760,9 +773,8 @@ def close(trans_id): trans_id: unique transaction id """ with sqleditor_close_session_lock: - # delete the SQLAutoComplete object - if trans_id in auto_complete_objects: - del auto_complete_objects[trans_id] + # delete this session's SQLAutoComplete object + auto_complete_objects.pop(_auto_complete_key(trans_id), None) if 'gridData' not in session: return make_json_response(data={'status': True}) @@ -829,6 +841,10 @@ def close_sqleditor_session(trans_id): :param trans_id: Transaction id :return: """ + # Drop this session's autocomplete object, which holds the autocomplete + # connection about to be released. + auto_complete_objects.pop(_auto_complete_key(trans_id), None) + if 'gridData' in session and str(trans_id) in session['gridData']: cmd_obj_str = session['gridData'][str(trans_id)]['command_obj'] # Use pickle.loads function to get the command object @@ -2137,14 +2153,18 @@ def auto_complete(trans_id): trans_obj is not None and session_obj is not None: with sqleditor_close_session_lock: - if trans_id not in auto_complete_objects: + ac_key = _auto_complete_key(trans_id) + auto_complete_obj = auto_complete_objects.get(ac_key) + # Only reuse the cached object if it was built on the very + # connection that this session's transaction uses now. + if auto_complete_obj is None or auto_complete_obj.conn is not conn: # Create object of SQLAutoComplete class and pass # connection object - auto_complete_objects[trans_id] = \ + auto_complete_obj = \ SQLAutoComplete(sid=trans_obj.sid, did=trans_obj.did, conn=conn) + auto_complete_objects[ac_key] = auto_complete_obj - auto_complete_obj = auto_complete_objects[trans_id] # # Get the auto completion suggestions. res = auto_complete_obj.get_completions(full_sql, text_before_cursor) diff --git a/web/pgadmin/tools/sqleditor/tests/test_autocomplete_per_session.py b/web/pgadmin/tools/sqleditor/tests/test_autocomplete_per_session.py new file mode 100644 index 000000000..d92f83409 --- /dev/null +++ b/web/pgadmin/tools/sqleditor/tests/test_autocomplete_per_session.py @@ -0,0 +1,98 @@ +########################################################################## +# +# pgAdmin 4 - PostgreSQL Tools +# +# Copyright (C) 2013 - 2026, The pgAdmin Development Team +# This software is released under the PostgreSQL Licence +# +########################################################################## + +import json +from unittest.mock import MagicMock, patch + +from pgadmin.utils.route import BaseTestGenerator +import pgadmin.tools.sqleditor as sqleditor + + +class _FakeSession(dict): + """A Flask session stand-in with its own session id.""" + def __init__(self, sid): + super().__init__() + self.sid = sid + + +class TestAutocompletePerSession(BaseTestGenerator): + """ + The cached SQLAutoComplete object holds a live connection, so it must + only ever be used by the session that created it, even when another + session picks the same transaction id. + """ + scenarios = [ + ('Autocomplete objects are not shared between sessions', dict()), + ] + + TRANS_ID = 4242 + + def setUp(self): + self.saved_objects = dict(sqleditor.auto_complete_objects) + sqleditor.auto_complete_objects.clear() + # One connection per session; the transaction id is the same. + self.conns = {'session-a': MagicMock(name='conn-a'), + 'session-b': MagicMock(name='conn-b')} + + def tearDown(self): + sqleditor.auto_complete_objects.clear() + sqleditor.auto_complete_objects.update(self.saved_objects) + + def _fake_auto_complete(self, sid, did, conn): + obj = MagicMock() + obj.conn = conn + obj.get_completions.return_value = conn._mock_name + return obj + + def _call(self, view, sess, *args): + conn = self.conns[sess.sid] + trans_obj = MagicMock(sid=1, did=2) + with self.app.test_request_context( + data=json.dumps(['SELECT * FROM ', 'SELECT * FROM '])), \ + patch.object(sqleditor, 'session', sess), \ + patch.object(sqleditor, 'check_and_upgrade_to_qt', + return_value=(False, None)), \ + patch.object(sqleditor, 'check_transaction_status', + return_value=(True, None, conn, trans_obj, {})), \ + patch.object(sqleditor, 'SQLAutoComplete', + side_effect=self._fake_auto_complete): + response = view.__wrapped__(*args) + if hasattr(response, 'data'): + return json.loads(response.data.decode('utf-8'))['data'] + return response + + def runTest(self): + session_a = _FakeSession('session-a') + session_b = _FakeSession('session-b') + + data = self._call(sqleditor.auto_complete, session_a, self.TRANS_ID) + self.assertEqual(data['result'], 'conn-a') + + # Same transaction id, different session: must get its own object + # on its own connection, not session A's. + data = self._call(sqleditor.auto_complete, session_b, self.TRANS_ID) + self.assertEqual(data['result'], 'conn-b') + self.assertEqual(len(sqleditor.auto_complete_objects), 2) + + # Session A's object is still its own. + data = self._call(sqleditor.auto_complete, session_a, self.TRANS_ID) + self.assertEqual(data['result'], 'conn-a') + + # Closing the transaction in session B leaves session A's alone. + self._call(sqleditor.close, session_b, self.TRANS_ID) + self.assertEqual( + list(sqleditor.auto_complete_objects), + [('session-a', self.TRANS_ID)]) + + # Releasing the transaction in session A (as on logout or a + # connection change) drops its object too. + with self.app.test_request_context(), \ + patch.object(sqleditor, 'session', session_a): + sqleditor.close_sqleditor_session(str(self.TRANS_ID)) + self.assertEqual(sqleditor.auto_complete_objects, {})