Fix the query tool restore connection issue on the server disconnection from the left side object explorer. #6502

web/pgadmin/browser/server_groups/servers/__init__.py

@@ -1377,7 +1377,7 @@ class ServerNode(PGChildNodeView):
             }
         )
 
-    def connect(self, gid, sid):
+    def connect(self, gid, sid, is_qt=False):
         """
         Connect the Server and return the connection object.
         Verification Process before Connection:
@@ -1443,7 +1443,9 @@ class ServerNode(PGChildNodeView):
 
         # Connect the Server
         manager = get_driver(PG_DEFAULT_DRIVER).connection_manager(sid)
-        if not manager.connection().connected():
+        # Update the manager with the server details if not connected and
+        # the API call is not made from SQL Editor or View/Edit Data tool
+        if not manager.connection().connected() and not is_qt:
             manager.update(server)
         conn = manager.connection()
 

web/pgadmin/tools/sqleditor/__init__.py

@@ -802,6 +802,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:
+            return response
+        else:
+            conn = manager.connection(did=trans_obj.did)
+
         status, msg = default_conn.connect()
         if not status:
             return make_json_response(

web/pgadmin/tools/sqleditor/static/js/components/QueryToolComponent.jsx

@@ -333,7 +333,7 @@ export default function QueryToolComponent({params, pgWindow, pgAdmin, selectedN
     }
   };
 
-  const initializeQueryTool = (password)=>{
+  const initializeQueryTool = (password, explainObject=null, macroSQL='', executeCursor=false)=>{
     let selectedConn = _.find(qtState.connection_list, (c)=>c.is_selected);
     let baseUrl = '';
     if(qtState.params.is_query_tool) {
@@ -364,9 +364,9 @@ export default function QueryToolComponent({params, pgWindow, pgAdmin, selectedN
           obtaining_conn: false,
         });
 
-        if(!qtState.params.is_query_tool) {
+        eventBus.current.fireEvent(QUERY_TOOL_EVENTS.TRIGGER_EXECUTION, explainObject, macroSQL, executeCursor);
-          eventBus.current.fireEvent(QUERY_TOOL_EVENTS.TRIGGER_EXECUTION);
+        let msg = `${selectedConn['server_name']}/${selectedConn['database_name']} - Database connected`;
-        }
+        pgAdmin.Browser.notifier.success(_.escape(msg));
       }).catch((error)=>{
         if(error.response?.request?.responseText?.search('Ticket expired') !== -1) {
           Kerberos.fetch_ticket()
@@ -415,6 +415,8 @@ export default function QueryToolComponent({params, pgWindow, pgAdmin, selectedN
     getSQLScript();
     initializeQueryTool();
 
+    eventBus.current.registerListener(QUERY_TOOL_EVENTS.REINIT_QT_CONNECTION, initializeQueryTool);
+
     eventBus.current.registerListener(QUERY_TOOL_EVENTS.FOCUS_PANEL, (qtPanelId)=>{
       docker.current.focus(qtPanelId);
     });
@@ -684,6 +686,10 @@ export default function QueryToolComponent({params, pgWindow, pgAdmin, selectedN
               params: {
                 ...prev.params,
                 trans_id: respData.data.trans_id,
+                server_name: connectionData.server_name,
+                database_name: connectionData.database_name,
+                dbname: connectionData.database_name,
+                user: connectionData.user,
                 sid: connectionData.sid,
                 did: connectionData.did,
                 title: connectionData.title,

web/pgadmin/tools/sqleditor/static/js/components/QueryToolConstants.js

@@ -57,6 +57,7 @@ export const QUERY_TOOL_EVENTS = {
   HANDLE_API_ERROR: 'HANDLE_API_ERROR',
   SET_FILTER_INFO: 'SET_FILTER_INFO',
   FETCH_MORE_ROWS: 'FETCH_MORE_ROWS',
+  REINIT_QT_CONNECTION:'REINIT_QT_CONNECTION',
 
   EDITOR_LAST_FOCUS: 'EDITOR_LAST_FOCUS',
   EDITOR_FIND_REPLACE: 'EDITOR_FIND_REPLACE',

web/pgadmin/tools/sqleditor/static/js/components/sections/MainToolBar.jsx

@@ -276,10 +276,10 @@ export function MainToolBar({containerRef, onFilterClick, onManageMacros, onAddT
   }, [queryToolConnCtx.connectionStatus]);
 
   const onCommitClick=()=>{
-    eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, 'COMMIT;', null, true);
+    eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, 'COMMIT;', null, '', true);
   };
   const onRollbackClick=()=>{
-    eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, 'ROLLBACK;', null, true);
+    eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, 'ROLLBACK;', null, '', true);
   };
   const executeMacro = (m)=>{
     eventBus.fireEvent(QUERY_TOOL_EVENTS.TRIGGER_EXECUTION, null, m.sql);

web/pgadmin/tools/sqleditor/static/js/components/sections/Query.jsx

@@ -152,10 +152,10 @@ export default function Query({onTextSelect}) {
         query = query || editor.current?.getValue() || '';
       }
       if(query) {
-        eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, query, explainObject, external, null, executeCursor);
+        eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, query, explainObject, macroSQL, external, null, executeCursor);
       }
     } else {
-      eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, null, null);
+      eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, null, null, '');
     }
   };
 

web/pgadmin/tools/sqleditor/static/js/components/sections/ResultSet.jsx

@@ -30,6 +30,7 @@ import EmptyPanelMessage from '../../../../../../static/js/components/EmptyPanel
 import { GraphVisualiser } from './GraphVisualiser';
 import { usePgAdmin } from '../../../../../../static/js/BrowserComponent';
 import pgAdmin from 'sources/pgadmin';
+import ConnectServerContent from '../../../../../../static/js/Dialogs/ConnectServerContent';
 
 const StyledBox = styled(Box)(({theme}) => ({
   display: 'flex',
@@ -39,7 +40,7 @@ const StyledBox = styled(Box)(({theme}) => ({
 }));
 
 export class ResultSetUtils {
-  constructor(api, transId, isQueryTool=true) {
+  constructor(api, queryToolCtx, transId, isQueryTool=true) {
     this.api = api;
     this.transId = transId;
     this.startTime = new Date();
@@ -48,6 +49,7 @@ export class ResultSetUtils {
     this.clientPKLastIndex = 0;
     this.historyQuerySource = null;
     this.hasQueryCommitted = false;
+    this.queryToolCtx = queryToolCtx;
   }
 
   static generateURLReconnectionFlag(baseUrl, transId, shouldReconnect) {
@@ -187,8 +189,49 @@ export class ResultSetUtils {
       );
     }
   }
+  connectServerModal (modalData, connectCallback, cancelCallback) {
+    this.queryToolCtx.modal.showModal(gettext('Connect to server'), (closeModal)=>{
+      return (
+        <ConnectServerContent
+          closeModal={()=>{
+            cancelCallback?.();
+            closeModal();
+          }}
+          data={modalData}
+          onOK={(formData)=>{
+            connectCallback(Object.fromEntries(formData));
+            closeModal();
+          }}
+        />
+      );
+    }, {
+      onClose: cancelCallback,
+    });
+  };
+  async connectServer (sid, user, formData, connectCallback) {
+    try {
+      let {data: respData} = await this.api({
+        method: 'POST',
+        url: url_for('sqleditor.connect_server', {
+          'sid': sid,
+          ...(user ? {
+            'usr': user,
+          }:{}),
+        }),
+        headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
+        data: formData
+      });
+      connectCallback?.(respData.data);
+    } catch (error) {
+      this.connectServerModal(error.response?.data?.result, async (data)=>{
+        this.connectServer(sid, user, data, connectCallback);
+      }, ()=>{
+        /*This is intentional (SonarQube)*/
+      });
+    }
+  };
 
-  async startExecution(query, explainObject, onIncorrectSQL, flags={
+  async startExecution(query, explainObject, macroSQL, onIncorrectSQL, flags={
     isQueryTool: true, external: false, reconnect: false, executeCursor: false
   }) {
     let startTime = new Date();
@@ -244,17 +287,27 @@ export class ResultSetUtils {
         }
       }
     } catch(e) {
+      if(e?.response?.status == 428){
+        this.connectServerModal(e.response?.data?.result, async (passwordData)=>{
+          await this.connectServer(this.queryToolCtx.params.sid, this.queryToolCtx.params.user, passwordData, async ()=>{
+            await this.eventBus.fireEvent(QUERY_TOOL_EVENTS.REINIT_QT_CONNECTION, '', explainObject, macroSQL, flags.executeCursor);
+          });
+        }, ()=>{
+          /*This is intentional (SonarQube)*/
+        });
+      } else {
         this.eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_END);
         this.eventBus.fireEvent(QUERY_TOOL_EVENTS.HANDLE_API_ERROR,
           e,
           {
             connectionLostCallback: ()=>{
-            this.eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, query, explainObject, flags.external, true, flags.executeCursor);
+              this.eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, query, explainObject, '', flags.external, true, flags.executeCursor);
             },
             checkTransaction: true,
           }
         );
       }
+    }
     return false;
   }
 
@@ -304,7 +357,7 @@ export class ResultSetUtils {
     });
     this.eventBus.fireEvent(QUERY_TOOL_EVENTS.HANDLE_API_ERROR, error, {
       connectionLostCallback: ()=>{
-        this.eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, this.query, explainObject, flags.external, true, flags.executeCursor);
+        this.eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, this.query, explainObject, '', flags.external, true, flags.executeCursor);
       },
       checkTransaction: true,
     });
@@ -767,7 +820,7 @@ export function ResultSet() {
   const [columns, setColumns] = useState([]);
   const [isLoadingMore, setIsLoadingMore] = useState(false);
   const api = getApiInstance();
-  const rsu = React.useRef(new ResultSetUtils(api, queryToolCtx.params.trans_id, queryToolCtx.params.is_query_tool));
+  const rsu = React.useRef(new ResultSetUtils(api, queryToolCtx, queryToolCtx.params.trans_id, queryToolCtx.params.is_query_tool));
   const [dataChangeStore, dispatchDataChange] = React.useReducer(dataChangeReducer, {});
   const [selectedRows, setSelectedRows] = useState(new Set());
   const [selectedColumns, setSelectedColumns] = useState(new Set());
@@ -802,7 +855,7 @@ export function ResultSet() {
     eventBus.fireEvent(QUERY_TOOL_EVENTS.SELECTED_ROWS_COLS_CELL_CHANGED, selectedRows.size, selectedColumns.size, selectedRange.current, selectedCell.current?.length);
   };
 
-  const executionStartCallback = async (query, explainObject, external=false, reconnect=false, executeCursor=false)=>{
+  const executionStartCallback = async (query, explainObject, macroSQL, external=false, reconnect=false, executeCursor=false)=>{
     const yesCallback = async ()=>{
       /* Reset */
       eventBus.fireEvent(QUERY_TOOL_EVENTS.HIGHLIGHT_ERROR, null);
@@ -813,7 +866,7 @@ export function ResultSet() {
       setLoaderText(gettext('Waiting for the query to complete...'));
       setDataOutputQuery(query);
       return await rsu.current.startExecution(
-        query, explainObject,
+        query, explainObject, macroSQL,
         ()=>{
           setColumns([]);
           setRows([]);
@@ -852,8 +905,15 @@ export function ResultSet() {
     };
 
     const executeAndPoll = async ()=>{
-      await yesCallback();
+      await yesCallback().then((res)=>{
+        if(res){
           pollCallback();
+        } else {
+          eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_END);
+        }
+      }).catch((err)=>{
+        console.error(err);
+      });
     };
 
     if(isDataChanged()) {
@@ -989,7 +1049,7 @@ export function ResultSet() {
           e,
           {
             connectionLostCallback: ()=>{
-              eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, rsu.current.query, null, false, true);
+              eventBus.fireEvent(QUERY_TOOL_EVENTS.EXECUTION_START, rsu.current.query, null, '', false, true);
             },
             checkTransaction: true,
           }

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

@@ -28,6 +28,7 @@ from pgadmin.utils.driver import get_driver
 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:
@@ -81,6 +82,21 @@ class StartRunningQuery:
 
             # 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:
+                    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)