Reconnect a lost Query Tool connection as the user and role it was changed to. #10613

web/pgadmin/tools/sqleditor/__init__.py

@@ -30,7 +30,7 @@
 from flask import request
 from flask_babel import gettext
 from pgadmin.tools.sqleditor.utils.query_tool_connection_check \
-    import query_tool_connection_check
+    import query_tool_connection_check, connect_as_transaction_user
 from pgadmin.user_login_check import pga_login_required
 from flask_security import current_user, permissions_required
 from pgadmin.misc.file_manager import Filemanager
@@ -966,13 +966,7 @@ def check_transaction_status(trans_id, auto_comp=False):
         return False, internal_server_error(errormsg=str(e)), None, None, None
 
     if connect and conn and not conn.connected():
-        if auto_comp and getattr(trans_obj, 'conn_user', None):
-            status, msg = conn.connect(
-                user=trans_obj.conn_user,
-                role=getattr(trans_obj, 'conn_role', None),
-                password=None)
-        else:
-            status, msg = conn.connect()
+        status, msg = connect_as_transaction_user(conn, trans_obj)
 
         if not status:
             current_app.logger.error(msg)
@@ -1613,7 +1607,7 @@ def _check_and_connect(trans_obj):
 
     # Connect to the Server if not connected.
     if not conn.connected():
-        status, msg = conn.connect()
+        status, msg = connect_as_transaction_user(conn, trans_obj)
         if not status:
             return True, msg, conn
     return False, '', conn
@@ -1670,6 +1664,13 @@ def save(trans_id):
             # when password is not saved for the server
             if response is not None:
                 return response
+            # The reconnect failed (e.g. the user the connection was
+            # changed to cannot log in), so report it; there is no
+            # transaction object to carry on with.
+            if not status:
+                return make_json_response(
+                    data={'status': False, 'result': error_msg}
+                )
 
         is_error, _, conn = _check_and_connect(trans_obj)
         if is_error:

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

@@ -52,6 +52,14 @@ class TestCheckTransactionStatusReconnect(BaseTestGenerator):
         ('connect=1 reconnect succeeds', dict(
             auto_comp=False, conn_user=None, query_string='connect=1',
             connect_result=(True, None), expect_ok=True)),
+        # The main connection reconnects as the user it was changed to too,
+        # rather than the server's own user (#10613).
+        ('connect=1 reconnect as changed user fails, left to the caller',
+         dict(auto_comp=False, conn_user='jane', query_string='connect=1',
+              connect_result=(False, RECONNECT_ERROR), expect_ok=True)),
+        ('connect=1 reconnect as changed user succeeds', dict(
+            auto_comp=False, conn_user='jane', query_string='connect=1',
+            connect_result=(True, None), expect_ok=True)),
     ]
 
     def setUp(self):

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

@@ -0,0 +1,209 @@
+##########################################################################
+#
+# pgAdmin 4 - PostgreSQL Tools
+#
+# Copyright (C) 2013 - 2026, The pgAdmin Development Team
+# This software is released under the PostgreSQL Licence
+#
+##########################################################################
+
+"""
+A lost Query Tool connection must be re-established as the user and role it
+was changed to, not as the server's own user (#10613).
+"""
+
+import importlib
+import pickle
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from flask import Flask
+
+from pgadmin.tools.sqleditor import _check_and_connect, save
+from pgadmin.tools.sqleditor.utils.query_tool_connection_check import \
+    connect_as_transaction_user, query_tool_connection_check
+from pgadmin.tools.sqleditor.utils.start_running_query import \
+    StartRunningQuery
+from pgadmin.utils.route import BaseTestGenerator
+
+# The utils package re-exports the function under the module's own name, so
+# fetch the module itself rather than resolving it by attribute.
+check_module = importlib.import_module(
+    'pgadmin.tools.sqleditor.utils.query_tool_connection_check')
+RECONNECT_ERROR = 'password authentication failed for user "jane"'
+
+
+def _trans_obj(conn_user, conn_role):
+    trans_obj = SimpleNamespace(sgid=1, sid=1, did=2, conn_id='conn-1',
+                                dbname='postgres')
+    if conn_user is not None or conn_role is not None:
+        trans_obj.conn_user = conn_user
+        trans_obj.conn_role = conn_role
+    return trans_obj
+
+
+def _driver_for(conn):
+    manager = MagicMock()
+    manager.connection.return_value = conn
+    driver = MagicMock()
+    driver.connection_manager.return_value = manager
+    return driver
+
+
+class TestReconnectAsChangedUser(BaseTestGenerator):
+    """Every path that reopens the Query Tool's own connection uses the
+    stored user and role, and falls back to a plain connect() only when no
+    user was recorded (e.g. View/Edit Data, or a transaction created before
+    the user was stored)."""
+
+    scenarios = [
+        ('Changed user and role', dict(
+            conn_user='jane', conn_role='reporting',
+            expected_kwargs=dict(user='jane', role='reporting',
+                                 password=None))),
+        ('Changed user, no role', dict(
+            conn_user='jane', conn_role=None,
+            expected_kwargs=dict(user='jane', role=None, password=None))),
+        ('Changed role, no user', dict(
+            conn_user=None, conn_role='reporting',
+            expected_kwargs=dict(user=None, role='reporting',
+                                 password=None))),
+        ('No user or role recorded', dict(
+            conn_user=None, conn_role=None, expected_kwargs={})),
+    ]
+
+    def setUp(self):
+        # Pure logic test with a mocked connection: skip the server setup.
+        unittest.TestCase.setUp(self)
+
+    def _conn(self, connect_result=(True, None)):
+        conn = MagicMock()
+        conn.connected.return_value = False
+        conn.connect.return_value = connect_result
+        return conn
+
+    def runTest(self):
+        self._test_helper()
+        self._test_query_tool_connection_check()
+        self._test_check_and_connect()
+
+    def _test_helper(self):
+        conn = self._conn()
+        status, _ = connect_as_transaction_user(
+            conn, _trans_obj(self.conn_user, self.conn_role))
+        self.assertTrue(status)
+        conn.connect.assert_called_once_with(**self.expected_kwargs)
+
+    def _test_query_tool_connection_check(self):
+        trans_obj = _trans_obj(self.conn_user, self.conn_role)
+        conn = self._conn()
+        view = MagicMock()
+        view.connect.return_value = MagicMock(status_code=200)
+        session_obj = {'command_obj': pickle.dumps(trans_obj)}
+
+        with patch.object(StartRunningQuery, 'retrieve_session_information',
+                          return_value=session_obj), \
+                patch.object(check_module.SchemaDiffRegistry,
+                             'get_node_view', return_value=view), \
+                patch.object(check_module, 'get_driver',
+                             return_value=_driver_for(conn)):
+            status, errmsg, res_conn, _, _, response = \
+                query_tool_connection_check(1234)
+
+        self.assertTrue(status)
+        self.assertIsNone(errmsg)
+        self.assertIs(res_conn, conn)
+        self.assertIsNone(response)
+        conn.connect.assert_called_once_with(**self.expected_kwargs)
+
+    def _test_check_and_connect(self):
+        conn = self._conn()
+        with patch('pgadmin.tools.sqleditor.get_driver',
+                   return_value=_driver_for(conn)):
+            is_error, _, res_conn = _check_and_connect(
+                _trans_obj(self.conn_user, self.conn_role))
+
+        self.assertFalse(is_error)
+        self.assertIs(res_conn, conn)
+        conn.connect.assert_called_once_with(**self.expected_kwargs)
+
+
+class TestStartRunningQueryFailedReconnect(BaseTestGenerator):
+    """When the reconnect fails, the error is returned instead of the query
+    being run on the dead connection, which would only be reported as a lost
+    connection and prompt for another reconnect."""
+
+    scenarios = [
+        ('Reconnect as changed user fails', dict()),
+    ]
+
+    def setUp(self):
+        # Pure logic test with a mocked connection: skip the server setup.
+        unittest.TestCase.setUp(self)
+
+    def runTest(self):
+        trans_obj = MagicMock(conn_id='conn-1', dbname='postgres')
+        conn = MagicMock()
+        conn.connected.return_value = False
+        session_obj = {'command_obj': pickle.dumps(None)}
+
+        app = Flask(__name__)
+        module = 'pgadmin.tools.sqleditor.utils.start_running_query'
+        with app.test_request_context(), \
+                patch(module + '.StartRunningQuery'
+                      '.retrieve_session_information',
+                      return_value=session_obj), \
+                patch(module + '.pickle.loads', return_value=trans_obj), \
+                patch(module + '.get_driver',
+                      return_value=_driver_for(conn)), \
+                patch('pgadmin.tools.sqleditor.utils'
+                      '.query_tool_connection_check',
+                      return_value=(False, RECONNECT_ERROR, None, None,
+                                    None, None)):
+            response = StartRunningQuery(MagicMock(), MagicMock()).execute(
+                dict(sql='SELECT 1', explain_plan=None), 1234, {},
+                connect=True)
+
+        data = response.get_json()['data']
+        self.assertFalse(data['status'])
+        self.assertEqual(data['result'], RECONNECT_ERROR)
+        conn.execute_async.assert_not_called()
+        conn.execute_void.assert_not_called()
+
+
+class TestSaveFailedReconnect(BaseTestGenerator):
+    """Saving with connect=1 reports a failed reconnect instead of carrying
+    on with no transaction object."""
+
+    scenarios = [
+        ('Save after reconnect as changed user fails', dict()),
+    ]
+
+    def setUp(self):
+        # Pure logic test with a mocked connection: skip the server setup.
+        unittest.TestCase.setUp(self)
+
+    def runTest(self):
+        trans_obj = _trans_obj('jane', None)
+        conn = MagicMock()
+        session_obj = {'primary_keys': {'id': 'int4'}}
+        module = 'pgadmin.tools.sqleditor'
+
+        app = Flask(__name__)
+        with app.test_request_context(
+                query_string='connect=1', method='POST',
+                data='{"updated": {"1": {"id": 1}}}'), \
+                patch(module + '.check_transaction_status',
+                      return_value=(True, None, conn, trans_obj,
+                                    session_obj)), \
+                patch(module + '.query_tool_connection_check',
+                      return_value=(False, RECONNECT_ERROR, None, None,
+                                    None, None)), \
+                patch(module + '._check_and_connect') as check_and_connect:
+            response = save.__wrapped__(1234)
+
+        data = response.get_json()['data']
+        self.assertFalse(data['status'])
+        self.assertEqual(data['result'], RECONNECT_ERROR)
+        check_and_connect.assert_not_called()

web/pgadmin/tools/sqleditor/utils/query_tool_connection_check.py

@@ -20,6 +20,27 @@
 from pgadmin.utils.driver import get_driver
 
 
+def connect_as_transaction_user(conn, trans_obj):
+    """
+    (Re)connect a Query Tool connection as the user and role it was opened
+    with, either of which may be recorded without the other. When the
+    connection was changed to a different user or role, a bare connect()
+    would log in as the server's own user instead, so queries would silently
+    run as someone else. The password is never stored, so the
+    driver falls back to the connection's or server's saved password, and
+    the connection fails rather than switching user if that is not accepted.
+
+    :param conn: Connection to (re)connect.
+    :param trans_obj: Transaction (command) object for the Query Tool.
+    :return: (status, message) from Connection.connect().
+    """
+    user = getattr(trans_obj, 'conn_user', None)
+    role = getattr(trans_obj, 'conn_role', None)
+    if user or role:
+        return conn.connect(user=user, role=role, password=None)
+    return conn.connect()
+
+
 def query_tool_connection_check(trans_id):
     # This function will check if the query tool has the connection or not
     # if not then establishes the connection.
@@ -56,7 +77,7 @@ def query_tool_connection_check(trans_id):
                 **({"database": transaction_object.dbname} if hasattr(
                     transaction_object, 'dbname') else {}))
 
-        status, msg = conn.connect()
+        status, msg = connect_as_transaction_user(conn, transaction_object)
         if not status:
             current_app.logger.error(msg)
             return status, msg, None, None, None, None

web/pgadmin/tools/sqleditor/utils/start_running_query.py

@@ -86,13 +86,22 @@ def execute(self, sql, trans_id, http_session, connect=False):
 
                 status, errmsg, _, _, _, response = \
                     query_tool_connection_check(trans_id)
-                # If database does not exist then show error msg
-                if not status:
-                    result = errmsg
                 # This is required for asking user to enter password
                 # when password is not saved for the server
                 if response is not None:
                     return response
+                # If the reconnect failed (e.g. the database no longer
+                # exists, or the user the connection was changed to cannot
+                # log in), show the error rather than running the query on
+                # the dead connection, which would only report it as lost.
+                if not status:
+                    return make_json_response(
+                        data={
+                            'status': False, 'result': errmsg,
+                            'can_edit': False, 'can_filter': False,
+                            'notifies': None,
+                        }
+                    )
             effective_sql_statement = apply_explain_plan_wrapper_if_needed(
                 manager, sql)