Fixed an issue when the server/database connection was lost; the filter dialog is not getting saved. #6044

Author: anilsahoo20

web/pgadmin/tools/sqleditor/__init__.py

@@ -23,6 +23,8 @@
 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
 from pgadmin.user_login_check import pga_login_required
 from flask_security import current_user
 from pgadmin.misc.file_manager import Filemanager
@@ -821,13 +823,14 @@ def start_view_data(trans_id):
 
     # Connect to the Server if not connected.
     if not default_conn.connected():
-        view = SchemaDiffRegistry.get_node_view('server')
-        response = view.connect(trans_obj.sgid,
-                                trans_obj.sid, True)
-        if response.status_code == 428:
+        # This will check if view/edit data tool connection is lost or not,
+        # if lost then it will reconnect
+        status, error_msg, conn, trans_obj, session_obj, response = \
+            query_tool_connection_check(trans_id)
+        # This is required for asking user to enter password
+        # when password is not saved for the server
+        if response is not None:
             return response
-        else:
-            conn = manager.connection(did=trans_obj.did)
 
         status, msg = default_conn.connect()
         if not status:

web/pgadmin/tools/sqleditor/command.py

@@ -566,62 +566,23 @@ def get_primary_keys(self, default_conn=None):
     def get_all_columns_with_order(self, default_conn=None):
         """
         It is overridden method specially for Table because we all have to
-        fetch primary keys and rest of the columns both.
+        fetch primary keys.
 
         Args:
             default_conn: Connection object
 
         Returns:
             all_sorted_columns: Sorted columns for the Grid
-            all_columns: List of columns for the select2 options
         """
-        driver = get_driver(PG_DEFAULT_DRIVER)
-        if default_conn is None:
-            manager = driver.connection_manager(self.sid)
-            conn = manager.connection(did=self.did, conn_id=self.conn_id)
-        else:
-            conn = default_conn
 
         all_sorted_columns = []
         data_sorting = self.get_data_sorting()
-        all_columns = []
-        # Fetch the primary key column names
-        query = render_template(
-            "/".join([self.sql_path, 'primary_keys.sql']),
-            table_name=self.object_name,
-            table_nspname=self.nsp_name,
-            conn=conn,
-        )
-
-        status, result = conn.execute_dict(query)
 
-        if not status:
-            raise ExecuteError(result)
-
-        for row in result['rows']:
-            all_columns.append(row['attname'])
-
-        # Fetch the rest of the column names
-        query = render_template(
-            "/".join([self.sql_path, 'get_columns.sql']),
-            table_name=self.object_name,
-            table_nspname=self.nsp_name,
-            conn=conn,
-        )
-        status, result = conn.execute_dict(query)
-        if not status:
-            raise ExecuteError(result)
-
-        for row in result['rows']:
-            # Only append if not already present in the list
-            if row['attname'] not in all_columns:
-                all_columns.append(row['attname'])
-
-        # If user has custom data sorting then pass as it as it is
+        # If user has custom data sorting then pass as it is
         if data_sorting and len(data_sorting) > 0:
             all_sorted_columns = data_sorting
 
-        return all_sorted_columns, all_columns
+        return all_sorted_columns
 
     def can_edit(self):
         return True

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

@@ -12,3 +12,4 @@
 from .update_session_grid_transaction import update_session_grid_transaction
 from .start_running_query import *
 from .apply_explain_plan_wrapper import *
+from .query_tool_connection_check import *

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

@@ -37,8 +37,10 @@ def get(*args):
             msg = gettext('Success')
 
             try:
-                columns, column_list = \
-                    trans_obj.get_all_columns_with_order(conn)
+                columns = \
+                    trans_obj.get_all_columns_with_order()
+                column_list = [col_name for col_name in
+                               session_obj['columns_info'].keys()]
             except (ConnectionLost, SSHTunnelConnectionLost):
                 raise
             except Exception as e:

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

@@ -0,0 +1,68 @@
+##########################################################################
+#
+# pgAdmin 4 - PostgreSQL Tools
+#
+# Copyright (C) 2013 - 2025, The pgAdmin Development Team
+# This software is released under the PostgreSQL Licence
+#
+##########################################################################
+
+"""Check for query tool connection"""
+import pickle
+from flask_babel import gettext
+
+from config import PG_DEFAULT_DRIVER
+from pgadmin.utils.ajax import internal_server_error
+from pgadmin.tools.schema_diff.node_registry import SchemaDiffRegistry
+from pgadmin.tools.sqleditor.utils.start_running_query import StartRunningQuery
+from flask import Response, current_app, session
+
+from pgadmin.utils.driver import get_driver
+
+
+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.
+    session_obj = StartRunningQuery.retrieve_session_information(
+        session,
+        trans_id
+    )
+    if isinstance(session_obj, Response):
+        return session_obj
+
+    transaction_object = pickle.loads(session_obj['command_obj'])
+
+    # To verify if the transaction details for the specific query tool
+    # or View/Edit Data tool is available or not and if the server is
+    # disconnected from the Object Explorer then it reconnects
+    if transaction_object is not None and session_obj is not None:
+        view = SchemaDiffRegistry.get_node_view('server')
+        response = view.connect(transaction_object.sgid,
+                                transaction_object.sid, True)
+        # This is required for asking user to enter password
+        # when password is not saved for the server
+        if response.status_code == 428:
+            return False, None, None, None, None, response
+        else:
+            manager = get_driver(
+                PG_DEFAULT_DRIVER).connection_manager(
+                transaction_object.sid)
+            conn = manager.connection(
+                did=transaction_object.did,
+                conn_id=transaction_object.conn_id,
+                auto_reconnect=False,
+                use_binary_placeholder=True,
+                array_to_string=True,
+                **({"database": transaction_object.dbname} if hasattr(
+                    transaction_object, 'dbname') else {}))
+
+        status, msg = conn.connect()
+        if not status:
+            current_app.logger.error(msg)
+            return internal_server_error(errormsg=str(msg))
+        return status, None, conn, transaction_object, session_obj, None
+    else:
+        status = False
+        error_msg = gettext(
+            'Either transaction object or session object not found.')
+        return status, error_msg, None, None, None, None

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

@@ -28,7 +28,6 @@
 from pgadmin.utils.exception import ConnectionLost, SSHTunnelConnectionLost,\
     CryptKeyMissing
 from pgadmin.utils.constants import ERROR_MSG_TRANS_ID_NOT_FOUND
-from pgadmin.tools.schema_diff.node_registry import SchemaDiffRegistry
 
 
 class StartRunningQuery:
@@ -82,26 +81,15 @@ def execute(self, sql, trans_id, http_session, connect=False):
 
             # Connect to the Server if not connected.
             if connect and not conn.connected():
-                view = SchemaDiffRegistry.get_node_view('server')
-                response = view.connect(transaction_object.sgid,
-                                        transaction_object.sid, True)
-                if response.status_code == 428:
+                from pgadmin.tools.sqleditor.utils import \
+                    query_tool_connection_check
+
+                _, _, _, _, _, response = \
+                    query_tool_connection_check(trans_id)
+                # This is required for asking user to enter password
+                # when password is not saved for the server
+                if response is not None:
                     return response
-                else:
-                    conn = manager.connection(
-                        did=transaction_object.did,
-                        conn_id=self.connection_id,
-                        auto_reconnect=False,
-                        use_binary_placeholder=True,
-                        array_to_string=True,
-                        **({"database": transaction_object.dbname} if hasattr(
-                            transaction_object, 'dbname') else {}))
-
-                status, msg = conn.connect()
-                if not status:
-                    self.logger.error(msg)
-                    return internal_server_error(errormsg=str(msg))
-
             effective_sql_statement = apply_explain_plan_wrapper_if_needed(
                 manager, sql)