Ensure proper error message shown if any error occurs while restoring psql tool.

yogeshmahajan-1903 · yogeshmahajan-1903 · commit a090ac675871 · 2025-06-26T23:33:38.000+05:30
diff --git a/docs/en_US/preferences.rst b/docs/en_US/preferences.rst
@@ -329,7 +329,7 @@ Use the fields on the *User Interface* panel to set the user interface related p
   format. If the application is closed unexpectedly, the tab is accidentally closed,
   or the page is refreshed, the saved state will be automatically restored for
   each respective tool. **Note:** Saving the application state will not preserve data for tool tabs opened in
-  separate browser tabs when running in server mode.
+  separate browser tabs when running in server mode.Any tool referring ad-hoc server connection will not be restored.
 
 * Use the *Themes* drop-down listbox to select the theme for pgAdmin. You'll also get a preview just below the
   drop down. You can also submit your own themes,
diff --git a/web/pgadmin/static/js/ToolErrorView.jsx b/web/pgadmin/static/js/ToolErrorView.jsx
@@ -5,6 +5,7 @@ import { LAYOUT_EVENTS } from './helpers/Layout';
 import { styled } from '@mui/material/styles';
 import { FormHelperText, Box } from '@mui/material';
 import HTMLReactParse from 'html-react-parser';
+import { deleteToolData } from '../../settings/static/ApplicationStateProvider';
 
 const StyledBox = styled(Box)(({theme}) => ({
   color: theme.palette.text.primary,
@@ -15,11 +16,14 @@ const StyledBox = styled(Box)(({theme}) => ({
   height: '100%',
 }));
 
-export default function ToolErrorView({error, panelId, panelDocker}){
-
+export default function ToolErrorView({error, panelId, panelDocker, toolDataId}){
   panelDocker.eventBus.registerListener(LAYOUT_EVENTS.CLOSING, (id)=>{
     if(panelId == id) {
       panelDocker.close(panelId, true);
+      if(toolDataId){
+        let transId = toolDataId.replace('-','_');
+        deleteToolData(transId, transId);
+      }
     }
   });
 
@@ -33,4 +37,5 @@ ToolErrorView.propTypes = {
   error: PropTypes.string,
   panelId: PropTypes.string,
   panelDocker: PropTypes.object,
+  toolDataId: PropTypes.string
 };
diff --git a/web/pgadmin/tools/erd/__init__.py b/web/pgadmin/tools/erd/__init__.py
@@ -10,7 +10,7 @@
 """A blueprint module implementing the erd tool."""
 import json
 
-from flask import request, Response
+from flask import request, Response, session
 from flask import render_template, current_app as app
 from flask_security import permissions_required
 from pgadmin.user_login_check import pga_login_required
@@ -20,7 +20,7 @@
     SHORTCUT_FIELDS as shortcut_fields
 from pgadmin.utils.ajax import make_json_response, internal_server_error
 from pgadmin.model import Server
-from config import PG_DEFAULT_DRIVER
+from config import PG_DEFAULT_DRIVER, ALLOW_SAVE_PASSWORD
 from pgadmin.utils.driver import get_driver
 from pgadmin.browser.utils import underscore_unescape
 from pgadmin.browser.server_groups.servers.databases.schemas.utils \
@@ -541,10 +541,27 @@ def initialize_erd(trans_id, sgid, sid, did):
     # Read the data if present. Skipping read may cause connection
     # reset error if data is sent from the client
     if request.data:
-        _ = request.data
-
-    conn = _get_connection(sid, did, trans_id)
-
+        data = json.loads(request.data)
+        user = data.get('user', None)
+        server_label = data.get('server_name', None),
+        server_type = data.get('server_type', None),
+        db_name = data.get('db_name', None)
+    try:
+        conn = _get_connection(sid, did, trans_id,db_name)
+    except ConnectionLost as e:
+        return make_json_response(
+            success=0,
+            status=428,
+            result={"server_label": server_label,
+                    "username": user,
+                    "server_type": server_type,
+                    "errmsg": str(e),
+                    "prompt_password": True,
+                    "allow_save_password": True
+                    if ALLOW_SAVE_PASSWORD and
+                    session.get('allow_save_password', None) else False,
+                    }
+        )
     return make_json_response(
         data={
             'connId': str(trans_id),
@@ -554,7 +571,7 @@ def initialize_erd(trans_id, sgid, sid, did):
     )
 
 
-def _get_connection(sid, did, trans_id):
+def _get_connection(sid, did, trans_id, db_name=None):
     """
     Get the connection object of ERD.
     :param sid:
@@ -564,9 +581,13 @@ def _get_connection(sid, did, trans_id):
     """
     manager = get_driver(PG_DEFAULT_DRIVER).connection_manager(sid)
     try:
-        conn = manager.connection(did=did, conn_id=trans_id,
+        conn = manager.connection(conn_id=trans_id,
                                   auto_reconnect=True,
-                                  use_binary_placeholder=True)
+                                  use_binary_placeholder=True,
+                                  **({"database": db_name}
+                                     if db_name is not None
+                                     else {"did": did})
+                                  )
         status, msg = conn.connect()
         if not status:
             app.logger.error(msg)
diff --git a/web/pgadmin/tools/erd/static/js/ERDModule.js b/web/pgadmin/tools/erd/static/js/ERDModule.js
@@ -96,14 +96,19 @@ export default class ERDModule {
       
       parentData = {
         server_group: {
-          _id: connectionInfo.sgid || 0,
-          server_type: connectionInfo.server_type
+          _id: connectionInfo.sgid || 0
         },
         server: {
           _id: connectionInfo.sid,
+          server_type: connectionInfo.server_type,
+          label: connectionInfo.server_name,
+          user: {
+            name: connectionInfo.user
+          }
         },
         database: {
           _id: connectionInfo.did,
+          label: connectionInfo.db_name
         },
         schema: {
           _id: connectionInfo.scid || null,
@@ -142,7 +147,7 @@ export default class ERDModule {
       'pgadmin:tool:show',
       `${BROWSER_PANELS.ERD_TOOL}_${transId}`,
       panelUrl,
-      {sql_id: toolDataId, title: _.escape(panelTitle)},
+      {sql_id: toolDataId, title: _.escape(panelTitle), db_name:parentData.database.label, server_name: parentData.server.label, user: parentData.server.user.name, server_type: parentData.server.server_type},
       {title: 'Untitled', icon: 'fa fa-sitemap'},
       Boolean(open_new_tab?.includes('erd_tool'))
     );
diff --git a/web/pgadmin/tools/erd/static/js/erd_tool/components/ConnectionBar.jsx b/web/pgadmin/tools/erd/static/js/erd_tool/components/ConnectionBar.jsx
@@ -42,7 +42,7 @@ export const STATUS = {
 function ConnectionStatusIcon({status}) {
   if(status == STATUS.CONNECTING) {
     return <CircularProgress style={{height: '18px', width: '18px'}} />;
-  } else if(status == STATUS.CONNECTED || status == STATUS.FAILED) {
+  } else if(status == STATUS.CONNECTED) {
     return <ConnectedIcon />;
   } else {
     return <DisconnectedIcon />;
diff --git a/web/pgadmin/tools/erd/static/js/erd_tool/components/ERDTool.jsx b/web/pgadmin/tools/erd/static/js/erd_tool/components/ERDTool.jsx
@@ -39,6 +39,7 @@ import BeforeUnload from './BeforeUnload';
 import { isMac } from '../../../../../../static/js/keyboard_shortcuts';
 import DownloadUtils from '../../../../../../static/js/DownloadUtils';
 import { getToolData } from '../../../../../../settings/static/ApplicationStateProvider';
+import { connectServerModal, connectServer } from '../../../../../sqleditor/static/js/components/connectServer';
 
 /* Custom react-diagram action for keyboard events */
 export class KeyboardShortcutAction extends Action {
@@ -348,10 +349,12 @@ export default class ERDTool extends React.Component {
     });
 
     let done = await this.initConnection();
-    if(!done) return;
+    if(!done && !this.props.params.sql_id) return;
 
-    done = await this.loadPrequisiteData();
-    if(!done) return;
+    if(done){
+      done = await this.loadPrequisiteData();
+      if(!done && !this.props.params.sql_id) return;
+    }
 
 
     if(this.props.params.sql_id){
@@ -361,7 +364,7 @@ export default class ERDTool extends React.Component {
         this.diagram.clearSelection();
         this.registerModelEvents();
         this.setState({dirty: true});
-        this.eventBus.fireEvent(ERD_EVENTS.DIRTY, true, this.serializeFile());
+        this.eventBus.fireEvent(ERD_EVENTS.DIRTY, true, sqlValue);
       }
     }
     else if(this.props.params.gen) {
@@ -597,6 +600,7 @@ export default class ERDTool extends React.Component {
         current_file: fileName,
         dirty: false,
       });
+      this.eventBus.fireEvent(ERD_EVENTS.DIRTY, true, res.data);
       this.eventBus.fireEvent(ERD_EVENTS.DIRTY, false);
       this.setTitle(fileName);
       this.diagram.deserialize(res.data);
@@ -857,7 +861,12 @@ export default class ERDTool extends React.Component {
     });
 
     try {
-      let response = await this.apiObj.post(initUrl);
+      let response = await this.apiObj.post(
+        initUrl, 
+        {server_name: this.props.params.server_name,
+          server_type : this.props.params.server_type,
+          user: this.props.params.user,
+          db_name: this.props.params.db_name});
       this.setState({
         conn_status: CONNECT_STATUS.CONNECTED,
         server_version: response.data.data.serverVersion,
@@ -866,7 +875,16 @@ export default class ERDTool extends React.Component {
       return true;
     } catch (error) {
       this.setState({conn_status: CONNECT_STATUS.FAILED});
-      this.handleAxiosCatch(error);
+
+      connectServerModal(this.context, error.response?.data?.result, async (passwordData)=>{
+        await connectServer(this.apiObj, this.context, this.props.params.sid, this.props.params.sid, passwordData, async ()=>{
+          await this.initConnection();
+          await this.loadPrequisiteData();
+        });
+      }, ()=>{
+        this.setState({conn_status: CONNECT_STATUS.FAILED});
+      });
+
       return false;
     } finally {
       this.setLoading(null);
@@ -975,6 +993,9 @@ ERDTool.propTypes = {
     fgcolor: PropTypes.string,
     gen: PropTypes.bool.isRequired,
     sql_id: PropTypes.string,
+    server_name: PropTypes.string,
+    user: PropTypes.string,
+    db_name: PropTypes.string,
   }),
   pgWindow: PropTypes.object.isRequired,
   pgAdmin: PropTypes.object.isRequired,
diff --git a/web/pgadmin/tools/psql/__init__.py b/web/pgadmin/tools/psql/__init__.py
@@ -101,9 +101,10 @@ def panel(trans_id):
     s = Server.query.filter_by(id=int(params['sid'])).first()
     if s:
         data = _get_database_role(params['sid'], params['did'])
-        params['db'] = underscore_escape(data['db_name']) \
-            if 'db_name' in data else 'postgres'
-        params['role'] = underscore_escape(data['role'])
+        if data:
+            params['db'] = underscore_escape(data['db_name']) \
+                if 'db_name' in data else 'postgres'
+            params['role'] = underscore_escape(data['role'])
         set_env_variables(is_win=_platform == 'win32')
         return render_template("psql/index.html",
                                params=json.dumps(params))
@@ -303,6 +304,7 @@ def read_and_forward_pty_output(sid, data):
     # Check user is authenticated and PSQL is enabled in config.
     if current_user.is_authenticated and config.ENABLE_PSQL:
         connection_data = []
+        connection_successful = False
         try:
             db = ''
             if data['db']:
@@ -326,33 +328,39 @@ def read_and_forward_pty_output(sid, data):
                 return
             connection_data = get_connection_str(psql_utility, db,
                                                  manager)
+            connection_successful = True
         except Exception as e:
             # If any error raised during the start the PSQL emit error to UI.
             # request.sid: This sid is socket id.
-            sio.emit(
-                'conn_error',
-                {
-                    'error': 'Error while running psql command: {0}'.format(e),
-                }, namespace='/pty', room=request.sid)
+            error_msg = 'Error while running psql command: {0}'.format(e)
+            if str(e) == 'Server is not connected.':
+                error_msg = 'Error while opening psql tool: {0}'.format(e)
 
-        try:
-            if str(data['sid']) not in app.config['sid_soid_mapping']:
-                # request.sid: refer request.sid as socket id.
-                app.config['sid_soid_mapping'][str(data['sid'])] = list()
-                app.config['sid_soid_mapping'][str(data['sid'])].append(
-                    request.sid)
-            else:
-                app.config['sid_soid_mapping'][str(data['sid'])].append(
-                    request.sid)
+            sio.emit('conn_error',
+                     {'error': error_msg},
+                     namespace='/pty',
+                     room=request.sid)
 
-            sio.start_background_task(read_and_forward_pty_output,
-                                      request.sid, data)
-        except Exception as e:
-            sio.emit(
-                'conn_error',
-                {
-                    'error': 'Error while running psql command: {0}'.format(e),
-                }, namespace='/pty', room=request.sid)
+        if connection_successful:
+            try:
+                if str(data['sid']) not in app.config['sid_soid_mapping']:
+                    # request.sid: refer request.sid as socket id.
+                    app.config['sid_soid_mapping'][str(data['sid'])] = list()
+                    app.config['sid_soid_mapping'][str(data['sid'])].append(
+                        request.sid)
+                else:
+                    app.config['sid_soid_mapping'][str(data['sid'])].append(
+                        request.sid)
+
+                sio.start_background_task(read_and_forward_pty_output,
+                                          request.sid, data)
+            except Exception as e:
+                sio.emit(
+                    'conn_error',
+                    {'error':'Error while running psql command: {0}'.
+                        format(e)},
+                    namespace='/pty',
+                    room=request.sid)
     else:
         # Show error if user is not authenticated.
         sio.emit('conn_not_allow', {'sid': request.sid}, namespace='/pty',
@@ -378,12 +386,6 @@ def _get_connection(sid, data):
             status, msg = conn.connect()
         if not status:
             app.logger.error(msg)
-            sio.emit(sio.emit(
-                'conn_error',
-                {
-                    'error': 'Error while running psql command: {0}'
-                             ''.format('Server connection not present.'),
-                }, namespace='/pty', room=request.sid))
             raise RuntimeError('Server is not connected.')
 
         return conn, manager
diff --git a/web/pgadmin/tools/psql/static/js/PsqlModule.js b/web/pgadmin/tools/psql/static/js/PsqlModule.js
@@ -122,7 +122,7 @@ export default class Psql {
         },
         database: {
           _id: connectionInfo.did,
-          label: connectionInfo.db
+          _label: connectionInfo.db
         },
         schema: {
           _id: connectionInfo.scid || null,
@@ -161,7 +161,7 @@ export default class Psql {
     const transId = getRandomInt(1, 9999999);
     // Set psql tab title as per prefrences setting.
     let title_data = {
-      'database': parentData.database ? _.unescape(parentData.database.label) : 'postgres' ,
+      'database': parentData.database ? _.unescape(parentData.database._label) : 'postgres' ,
       'username': parentData.server.user.name,
       'server': parentData.server.label,
       'type': 'psql_tool',
diff --git a/web/pgadmin/tools/sqleditor/static/js/SQLEditorModule.js b/web/pgadmin/tools/sqleditor/static/js/SQLEditorModule.js
@@ -254,6 +254,7 @@ export default class SQLEditor {
                   error={params.error}
                   panelId={`${BROWSER_PANELS.QUERY_TOOL}_${params.trans_id}`}
                   panelDocker={panelDocker}
+                  toolDataId={params.toolDataId}
                 /> :
                 <QueryToolComponent params={params} 
                   pgWindow={pgWindow} 
diff --git a/web/pgadmin/tools/sqleditor/static/js/components/QueryToolComponent.jsx b/web/pgadmin/tools/sqleditor/static/js/components/QueryToolComponent.jsx
diff --git a/web/pgadmin/tools/sqleditor/static/js/components/connectServer.js b/web/pgadmin/tools/sqleditor/static/js/components/connectServer.js
diff --git a/web/pgadmin/tools/sqleditor/static/js/components/sections/ResultSet.jsx b/web/pgadmin/tools/sqleditor/static/js/components/sections/ResultSet.jsx
