Handle query execution on View/Edit Data when server is disconnected with server reconnect option. #8608

Author: RohitBhati8269

web/pgadmin/tools/sqleditor/__init__.py

@@ -818,6 +818,9 @@ def start_view_data(trans_id):
                                   info='DATAGRID_TRANSACTION_REQUIRED',
                                   status=404)
 
+    # Check if connect is passed in the request.
+    connect = 'connect' in request.args and request.args['connect'] == '1'
+
     # get the default connection as current connection which is attached to
     # trans id holds the cursor which has query result so we cannot use that
     # connection to execute another query otherwise we'll lose query result.
@@ -834,19 +837,21 @@ def start_view_data(trans_id):
 
     # Connect to the Server if not connected.
     if not default_conn.connected():
-        # 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
+        if connect:
+            # 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
 
         status, msg = default_conn.connect()
         if not status:
-            return make_json_response(
-                data={'status': status, 'result': "{}".format(msg)}
+            return service_unavailable(
+                gettext("Connection to the server has been lost."),
+                info="CONNECTION_LOST"
             )
 
     if status and conn is not None and \
@@ -1391,6 +1396,8 @@ def save(trans_id):
         changed_data = json.loads(request.data)
     else:
         changed_data = request.args or request.form
+    # Check if connect is passed in the request.
+    connect = 'connect' in request.args and request.args['connect'] == '1'
 
     # Check the transaction and connection status
     status, error_msg, conn, trans_obj, session_obj = \
@@ -1416,10 +1423,21 @@ def save(trans_id):
                 }
             )
 
+        if connect:
+            # 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
+
         is_error, errmsg, conn = _check_and_connect(trans_obj)
         if is_error:
-            return make_json_response(
-                data={'status': status, 'result': "{}".format(errmsg)}
+            return service_unavailable(
+                gettext("Connection to the server has been lost."),
+                info="CONNECTION_LOST"
             )
 
         status, res, query_results, _rowid = trans_obj.save(

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

@@ -543,9 +543,8 @@ export default function QueryToolComponent({params, pgWindow, pgAdmin, selectedN
           <br />
           <span>{gettext('Do you want to continue and establish a new session')}</span>
         </p>,
-        function() {
-          handleParams?.connectionLostCallback?.();
-        }, null,
+        () => handleParams?.connectionLostCallback?.(),
+        () => handleParams?.cancelCallback?.(),
         gettext('Continue'),
         gettext('Cancel')
       );

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

@@ -51,6 +51,7 @@ export class ResultSetUtils {
     this.historyQuerySource = null;
     this.hasQueryCommitted = false;
     this.queryToolCtx = queryToolCtx;
+    this.setLoaderText = null;
   }
 
   static generateURLReconnectionFlag(baseUrl, transId, shouldReconnect) {
@@ -445,23 +446,72 @@ export class ResultSetUtils {
     );
   }
 
-  saveData(reqData) {
+  saveData(reqData, shouldReconnect) {
+    // Generate the URL with the optional `connect=1` parameter.
+    const url = ResultSetUtils.generateURLReconnectionFlag('sqleditor.save', this.transId, shouldReconnect);
+
     return this.api.post(
-      url_for('sqleditor.save', {
-        'trans_id': this.transId
-      }),
+      url,
       JSON.stringify(reqData)
-    ).then(response => {
+    ).then((response) => {
       if (response.data?.data?.status) {
         // Set the commit flag to true if the save was successful
         this.hasQueryCommitted = true;
       }
       return response;
-    }).catch((error) => {
-      // Set the commit flag to false if there was an error
-      this.hasQueryCommitted = false;
-      throw error;
-    });
+    })
+      .catch(async (error) => {
+        if (error.response?.status === 428) {
+          // Handle 428: Show password dialog.
+          return new Promise((resolve, reject) => {
+            this.connectServerModal(
+              error.response?.data?.result,
+              async (formData) => {
+                try {
+                  await this.connectServer(
+                    this.queryToolCtx.params.sid,
+                    this.queryToolCtx.params.user,
+                    formData,
+                    async () => {
+                      let retryRespData = await this.saveData(reqData);
+                      // Set the commit flag to true if the save was successful
+                      this.hasQueryCommitted = true;
+                      pgAdmin.Browser.notifier.success(gettext('Server Connected.'));
+                      resolve(retryRespData);
+                    }
+                  );
+
+                } catch (retryError) {
+                  reject(retryError);
+                }
+              },
+              () => this.setLoaderText(null)
+            );
+          });
+        } else if (error.response?.status === 503) {
+          // Handle 503: Fire HANDLE_API_ERROR and wait for connectionLostCallback.
+          return new Promise((resolve, reject) => {
+            this.eventBus.fireEvent(QUERY_TOOL_EVENTS.HANDLE_API_ERROR, error, {
+              connectionLostCallback: async () => {
+                try {
+                  // Retry saveData with connect=1
+                  let retryRespData = await this.saveData(reqData, true);
+                  resolve(retryRespData);
+                } catch (retryError) {
+                  reject(retryError);
+                }
+              },
+              checkTransaction: true,
+              cancelCallback: () => this.setLoaderText(null)
+            },
+            );
+          });
+        } else {
+          // Set the commit flag to false if there was an error
+          this.hasQueryCommitted = false;
+          throw error;
+        }
+      });
   }
 
   async saveResultsToFile(fileName) {
@@ -861,6 +911,8 @@ export function ResultSet() {
 
   rsu.current.setEventBus(eventBus);
   rsu.current.setQtPref(queryToolCtx.preferences?.sqleditor);
+  // To use setLoaderText to the ResultSetUtils.
+  rsu.current.setLoaderText = setLoaderText;
 
   const isDataChanged = ()=>{
     return Boolean(_.size(dataChangeStore.updated) || _.size(dataChangeStore.added) || _.size(dataChangeStore.deleted));