From a5ff997adef2a011d65054fe44f4e88b16f40e2f Mon Sep 17 00:00:00 2001
From: Rohit Bhati <rohit.bhati@enterprisedb.com>
Date: Wed, 17 Dec 2025 16:12:26 +0530
Subject: [PATCH] Fix Query Tool state restoration for new connections and
 queries. #8987

---
 .../static/ApplicationStateProvider.jsx       |  8 +++----
 .../js/components/QueryToolComponent.jsx      | 24 +++++++++++++++++++
 .../static/js/components/sections/Query.jsx   | 22 ++++++++++++-----
 3 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/web/pgadmin/settings/static/ApplicationStateProvider.jsx b/web/pgadmin/settings/static/ApplicationStateProvider.jsx
index ca4cda6cbec..50d71b5548a 100644
--- a/web/pgadmin/settings/static/ApplicationStateProvider.jsx
+++ b/web/pgadmin/settings/static/ApplicationStateProvider.jsx
@@ -71,16 +71,16 @@ export function ApplicationStateProvider({children}){
         if(connectionInfo.is_editor_dirty){
           if(connectionInfo.external_file_changes){
             // file has external chages
-            return {loadFile: loadFile, fileName: fileName, data: toolContent, modifiedExternally: true};
+            return {loadFile: loadFile, fileName: fileName, data: toolContent, modifiedExternally: true, connectionInfo: connectionInfo};
           }
         }else if(connectionInfo.file_deleted){
-          return {loadFile: loadFile, fileName: null, data: toolContent};
+          return {loadFile: loadFile, fileName: null, data: toolContent, connectionInfo: connectionInfo};
         }else{
           loadFile = true;
-          return {loadFile: loadFile, fileName: fileName, data: null};
+          return {loadFile: loadFile, fileName: fileName, data: null, connectionInfo: connectionInfo};
         }
       }
-      return {loadFile: loadFile, fileName: fileName, data: toolContent};
+      return {loadFile: loadFile, fileName: fileName, data: toolContent, connectionInfo: connectionInfo};
 
     } catch (error) {
       let errorMsg = gettext(error?.response?.data?.errormsg || error);
diff --git a/web/pgadmin/tools/sqleditor/static/js/components/QueryToolComponent.jsx b/web/pgadmin/tools/sqleditor/static/js/components/QueryToolComponent.jsx
index cd1c3985770..69fbed08c06 100644
--- a/web/pgadmin/tools/sqleditor/static/js/components/QueryToolComponent.jsx
+++ b/web/pgadmin/tools/sqleditor/static/js/components/QueryToolComponent.jsx
@@ -305,6 +305,30 @@ export default function QueryToolComponent({params, pgWindow, pgAdmin, selectedN
   const restoreToolContent = async () =>{
     let toolContent = await getToolContent(qtState.params.trans_id);
     if(toolContent){
+      const connList = toolContent.connectionInfo?.connection_list;
+      if (Array.isArray(connList) && connList.length > 0) {
+        // ensure exactly one is_selected
+        if (!connList.some(c => c.is_selected)) {
+          connList[0].is_selected = true;
+        }
+        // pick selected connection to update params (so UI shows correct selected conn)
+        const sel = connList.find(c => c.is_selected) || connList[0];
+
+        setQtStatePartial(prev => ({
+          ...prev,
+          connection_list: connList,
+          params: {
+            ...prev.params,
+            sid: sel.sid,
+            did: sel.did,
+            user: sel.user,
+            role: sel.role,
+            server_name: sel.server_name,
+            database_name: sel.database_name,
+            title: sel.title,
+          },
+        }));
+      }
       if (toolContent?.modifiedExternally) {
         toolContent = await fmUtilsObj.warnFileReload(toolContent?.fileName, toolContent.data, '');
       }
diff --git a/web/pgadmin/tools/sqleditor/static/js/components/sections/Query.jsx b/web/pgadmin/tools/sqleditor/static/js/components/sections/Query.jsx
index 712803001f8..8897ebbe4e4 100644
--- a/web/pgadmin/tools/sqleditor/static/js/components/sections/Query.jsx
+++ b/web/pgadmin/tools/sqleditor/static/js/components/sections/Query.jsx
@@ -68,6 +68,7 @@ export default function Query({onTextSelect, setQtStatePartial}) {
   const preferencesStore = usePreferences();
   const modalId = MODAL_DIALOGS.QT_CONFIRMATIONS;
   const fmUtilsObj = useMemo(()=>new FileManagerUtils(queryToolCtx.api, {}), []);
+  const initialTransIdRef = React.useRef(queryToolCtx.params.trans_id);
 
   const highlightError = (cmObj, {errormsg: result, data}, executeCursor)=>{
     let errorLineNo = 0,
@@ -363,8 +364,9 @@ export default function Query({onTextSelect, setQtStatePartial}) {
 
   const change = useCallback(()=>{
     eventBus.fireEvent(QUERY_TOOL_EVENTS.QUERY_CHANGED, editor.current.isDirty());
-
-    if(isSaveToolDataEnabled('sqleditor') && editor.current.isDirty()){
+    const value = editor.current.getValue() || '';
+    const isDirty = editor.current.isDirty();
+    if(isSaveToolDataEnabled('sqleditor') && (isDirty || value.length === 0)){
       eventBus.fireEvent(QUERY_TOOL_EVENTS.TRIGGER_SAVE_QUERY_TOOL_DATA);
     }
 
@@ -379,10 +381,18 @@ export default function Query({onTextSelect, setQtStatePartial}) {
 
 
   const [saveQtData, setSaveQtData] = useState(false);
-  useDelayDebounce(()=>{
-    let connectionInfo = { ..._.find(queryToolCtx.connection_list, c => c.is_selected),
-      'open_file_name':queryToolCtx.current_file, 'is_editor_dirty': editor.current.isDirty() };
-    saveToolData('sqleditor', connectionInfo, queryToolCtx.params.trans_id, editor.current.getValue());
+  
+  useDelayDebounce(() => {
+    const currentTransId = initialTransIdRef.current || queryToolCtx.params.trans_id;
+    const currentConnList = queryToolCtx.connection_list;
+    const currentFile = queryToolCtx.current_file;
+
+    let connectionInfo = {
+      connection_list: currentConnList,
+      open_file_name: currentFile,
+      is_editor_dirty: editor.current.isDirty()
+    };
+    saveToolData('sqleditor', connectionInfo, currentTransId, editor.current.getValue());
     setSaveQtData(false);
   }, saveQtData, 500);