webapp/pgadmin4
1
0
Fork 0
mirror of synced 2026-10-23 07:08:55 +03:00
Code Activity

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

Browse Source
This commit is contained in:
Dave Page
2026-10-16 03:21:58 +01:00
parent ac76a81ad2
commit 7f21dc4610
5 changed files with 261 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
web/pgadmin/tools/sqleditor/__init__.py
View File

@@ -30,7 +30,7 @@ from flask import Response, url_for, render_template, session, current_app, \
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:

8
web/pgadmin/tools/sqleditor/tests/test_check_transaction_status_reconnect.py
View File

@@ -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):

209
web/pgadmin/tools/sqleditor/tests/test_reconnect_as_changed_user.py Normal file
View File

@@ -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()

23
web/pgadmin/tools/sqleditor/utils/query_tool_connection_check.py
View File

@@ -20,6 +20,27 @@ from flask import Response, current_app, session
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

15
web/pgadmin/tools/sqleditor/utils/start_running_query.py
View File

@@ -86,13 +86,22 @@ class StartRunningQuery:
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)