Fix Query Tool autocomplete objects being shared between sessions that use the same transaction id. #10607 (#10608)

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)

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, {})