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
@@ -819,13 +821,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:
@@ -839,6 +842,9 @@ def start_view_data(trans_id):
         # set fetched row count to 0 as we are executing query again.
         trans_obj.update_fetched_row_cnt(0)
 
+        # Fetch the columns for the table and store it in session
+        trans_obj.fetch_all_columns(conn)
+
         # Fetch the sql and primary_keys from the object
         sql = trans_obj.get_sql(default_conn)
         _, primary_keys = trans_obj.get_primary_keys(default_conn)
@@ -2264,7 +2270,6 @@ def get_filter_data(trans_id):
     Args:
         trans_id: unique transaction id
     """
-
     status, error_msg, conn, trans_obj, session_ob = \
         check_transaction_status(trans_id)
 

web/pgadmin/tools/sqleditor/command.py

@@ -228,6 +228,15 @@ def get_data_sorting(self):
             return self._data_sorting
         return None
 
+    def get_columns_list_with_order(self):
+        """
+        This function returns the list of columns with order.
+        """
+        if (self.all_columns_list_with_order_from_table and
+                len(self.all_columns_list_with_order_from_table) > 0):
+            return self.all_columns_list_with_order_from_table
+        return None
+
     def set_data_sorting(self, data_filter, set_from_filter_dialog=False):
         """
         This function validates the filter and set the
@@ -470,11 +479,53 @@ def __init__(self, **kwargs):
         # call base class init to fetch the table name
         super().__init__(**kwargs)
 
+        self.all_columns_list_with_order_from_table = None
+
         # Set the default sorting on table data by primary key if user
         # preference value is set
         self.data_sorting_by_pk = Preferences.module('sqleditor').preference(
             'table_view_data_by_pk').get()
 
+    def fetch_all_columns(self, conn):
+        """
+        This function fetches the list of columns for the
+        selected table and stores it locally.
+        """
+        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'])
+
+        self.all_columns_list_with_order_from_table = all_columns
+
     def get_sql(self, default_conn=None):
         """
         This method is used to create a proper SQL query
@@ -575,49 +626,17 @@ def get_all_columns_with_order(self, default_conn=None):
             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,
-        )
+        data_columns_list = self.get_columns_list_with_order()
 
-        status, result = conn.execute_dict(query)
+        # Assigns the list of columns from session for the table
+        if data_columns_list and len(data_columns_list) > 0:
+            all_columns = data_columns_list
 
-        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
 

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

@@ -38,7 +38,7 @@ def get(*args):
 
             try:
                 columns, column_list = \
-                    trans_obj.get_all_columns_with_order(conn)
+                    trans_obj.get_all_columns_with_order()
             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)