Ensure worspaces/pgadmin state is saved on abrupt shutdown of pgadmin.#GH_3319

Author: yogeshmahajan-1903

web/migrations/versions/4aa49c5d8eb1_.py

@@ -0,0 +1,33 @@
+
+"""empty message
+
+Revision ID: 4aa49c5d8eb1
+Revises: 1f0eddc8fc79
+Create Date: 2025-04-17 15:20:29.605023
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = '4aa49c5d8eb1'
+down_revision = '1f0eddc8fc79'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_table(
+        'pgadmin_state_data',
+        sa.Column('uid', sa.Integer(), nullable=False),
+        sa.Column('id', sa.Integer()),
+        sa.Column('connection_info', sa.JSON()),
+        sa.Column('tool_name', sa.String()),
+        sa.Column('tool_data', sa.String()),
+        sa.ForeignKeyConstraint(['uid'], ['user.id'], ondelete='CASCADE'),
+        sa.PrimaryKeyConstraint('id', 'uid'))
+
+
+def downgrade():
+    # pgAdmin only upgrades, downgrade not implemented.
+    pass

web/pgadmin/browser/static/js/browser.js

@@ -12,9 +12,13 @@ import _ from 'lodash';
 import { checkMasterPassword, showQuickSearch } from '../../../static/js/Dialogs/index';
 import { pgHandleItemError } from '../../../static/js/utils';
 import { send_heartbeat, stop_heartbeat } from './heartbeat';
-import getApiInstance from '../../../static/js/api_instance';
+import getApiInstance, {parseApiError} from '../../../static/js/api_instance';
 import usePreferences, { setupPreferenceBroadcast } from '../../../preferences/static/js/store';
 import checkNodeVisibility from '../../../static/js/check_node_visibility';
+import * as showQueryTool from '../../../tools/sqleditor/static/js/show_query_tool';
+import {relaunchPsqlTool} from '../../../tools/psql/static/js/show_psql_tool';
+import {relaunchErdTool} from '../../../tools/erd/static/js/show_erd_tool';
+import { relaunchSchemaDiff } from '../../../tools/schema_diff/static/js/showSchemaDiffTool';
 
 define('pgadmin.browser', [
   'sources/gettext', 'sources/url_for', 'sources/pgadmin',
@@ -206,6 +210,11 @@ define('pgadmin.browser', [
     uiloaded: function() {
       this.set_master_password('');
       this.check_version_update();
+      const prefStore = usePreferences.getState();
+      let save_the_workspace = prefStore.getPreferencesForModule('misc').save_the_workspace;
+      if(save_the_workspace){
+        this.restore_pgadmin_state();
+      }
     },
     check_corrupted_db_file: function() {
       getApiInstance().get(
@@ -291,6 +300,41 @@ define('pgadmin.browser', [
       });
     },
 
+    restore_pgadmin_state: async function () {
+      getApiInstance().get(
+        url_for('settings.get_pgadmin_state')
+      ).then((res)=> {
+        if(res.data.success && res.data.data.result.length > 0){
+          //let oldIds = []
+          _.each(res.data.data.result, function(tool_data){
+            if (tool_data.tool_name == 'sqleditor'){
+              showQueryTool.relaunchSqlTool(tool_data);
+            }else if(tool_data.tool_name == 'psql'){
+              relaunchPsqlTool(tool_data);
+            }else if(tool_data.tool_name == 'ERD'){
+              relaunchErdTool(tool_data);
+            }else if(tool_data.tool_name == 'schema_diff'){
+              relaunchSchemaDiff(tool_data);
+            }
+          });
+
+          // call clear query data for which query tool has been launched.
+          try {
+            getApiInstance().delete(url_for('settings.delete_pgadmin_state'), {
+            });
+          } catch (error) {
+            console.error(error);
+            pgAdmin.Browser.notifier.error(gettext('Failed to remove query data.') + parseApiError(error));
+          }
+
+        }
+      }).catch(function(error) {
+        pgAdmin.Browser.notifier.pgRespErrorNotify(error);
+      });
+      
+    },
+
+
     bind_beforeunload: function() {
       window.addEventListener('beforeunload', function(e) {
         /* Can open you in new tab */

web/pgadmin/misc/__init__.py

@@ -124,6 +124,19 @@ def register_preferences(self):
             )
         )
 
+        self.preference.register(
+            'user_interface', 'save_the_workspace',
+            gettext("Save the workspaces and data?"),
+            'boolean', True,
+            category_label=PREF_LABEL_USER_INTERFACE,
+            help_str=gettext(
+                'Turning on this setting will save the information about the'
+                ' tools that are opened and data within it. Data will be '
+                'restored in case of abrupt shutdown/refresh of the '
+                'application'
+            )
+        )
+
         if not config.SERVER_MODE:
             self.preference.register(
                 'file_downloads', 'automatically_open_downloaded_file',

web/pgadmin/misc/workspaces/static/js/WorkspaceProvider.jsx

@@ -75,16 +75,16 @@ export function WorkspaceProvider({children}) {
     pgAdmin.Browser.docker.currentWorkspace = newVal;
     if (newVal == WORKSPACES.DEFAULT) {
       setTimeout(() => {
-        pgAdmin.Browser.tree.selectNode(lastSelectedTreeItem.current, true, 'center');
+        pgAdmin.Browser.tree?.selectNode(lastSelectedTreeItem.current, true, 'center');
         lastSelectedTreeItem.current = null;
       }, 250);
     }  else {
       // Get the selected tree node and save it into the state variable.
-      let selItem = pgAdmin.Browser.tree.selected();
+      let selItem = pgAdmin.Browser.tree?.selected();
       if (selItem)
         lastSelectedTreeItem.current = selItem;
       // Deselect the node to disable the menu options.
-      pgAdmin.Browser.tree.deselect(selItem);
+      pgAdmin.Browser.tree?.deselect(selItem);
     }
     setCurrentWorkspace(newVal);
   };

web/pgadmin/model/__init__.py

@@ -33,7 +33,7 @@
 #
 ##########################################################################
 
-SCHEMA_VERSION = 44
+SCHEMA_VERSION = 45
 
 ##########################################################################
 #
@@ -392,6 +392,17 @@ class QueryHistoryModel(db.Model):
     last_updated_flag = db.Column(db.String(), nullable=False)
 
 
+class PgadminStateData(db.Model):
+    """Define the history SQL table."""
+    __tablename__ = 'pgadmin_state_data'
+    uid = db.Column(db.Integer(), db.ForeignKey(USER_ID), nullable=False,
+                    primary_key=True)
+    id = db.Column(db.Integer(),nullable=False,primary_key=True)
+    connection_info = db.Column(MutableDict.as_mutable(types.JSON))
+    tool_name = db.Column(db.String(), nullable=False)
+    tool_data = db.Column(PgAdminDbBinaryString())
+
+
 class Database(db.Model):
     """
     Define a Database.

web/pgadmin/settings/__init__.py

@@ -21,9 +21,10 @@
     success_return, internal_server_error
 from pgadmin.utils.menu import MenuItem
 
-from pgadmin.model import db, Setting
+from pgadmin.model import db, Setting, PgadminStateData
 from pgadmin.utils.constants import MIMETYPE_APP_JS
 from .utils import get_dialog_type, get_file_type_setting
+from cryptography.fernet import Fernet
 
 MODULE_NAME = 'settings'
 
@@ -52,7 +53,9 @@ def get_exposed_url_endpoints(self):
             'settings.save_tree_state', 'settings.get_tree_state',
             'settings.reset_tree_state',
             'settings.save_file_format_setting',
-            'settings.get_file_format_setting'
+            'settings.get_file_format_setting',
+            'settings.save_pgadmin_state',
+            'settings.delete_pgadmin_state'
         ]
 
 
@@ -256,3 +259,111 @@ def get_file_format_setting():
 
     return make_json_response(success=True,
                               info=get_file_type_setting(list(data.values())))
+
+
+@blueprint.route(
+    '/save_pgadmin_state',
+    methods=["POST"], endpoint='save_pgadmin_state'
+)
+@pga_login_required
+def save_pgadmin_state_data():
+    """
+    Args:
+        sid: server id
+        did: database id
+    """
+    data = json.loads(request.data)
+    id = data['trans_id']
+    fernet = Fernet(current_app.config['SECRET_KEY'].encode())
+    tool_data = fernet.encrypt(json.dumps(data['tool_data']).encode())
+    connection_info = data['connection_info'] \
+        if 'connection_info' in data else None
+    try:
+        data_entry = PgadminStateData(
+            uid=current_user.id, id=id,connection_info=connection_info,
+            tool_name=data['tool_name'], tool_data=tool_data)
+
+        db.session.merge(data_entry)
+        db.session.commit()
+    except Exception as e:
+        print(e)
+        db.session.rollback()
+
+    return make_json_response(
+        data={
+            'status': True,
+            'msg': 'Success',
+        }
+    )
+
+
+@blueprint.route(
+    '/get_pgadmin_state',
+    methods=["GET"], endpoint='get_pgadmin_state'
+)
+@pga_login_required
+def get_pgadmin_state():
+    fernet = Fernet(current_app.config['SECRET_KEY'].encode())
+    result = db.session \
+        .query(PgadminStateData) \
+        .filter(PgadminStateData.uid == current_user.id) \
+        .all()
+
+    res = []
+    for row in result:
+        res.append({'tool_name': row.tool_name,
+                    'connection_info': row.connection_info,
+                    'tool_data': fernet.decrypt(row.tool_data).decode(),
+                    'id': row.id
+                    })
+    return make_json_response(
+        data={
+            'status': True,
+            'msg': '',
+            'result': res
+        }
+    )
+
+
+@blueprint.route(
+    '/delete_pgadmin_state/',
+    methods=["DELETE"], endpoint='delete_pgadmin_state')
+@pga_login_required
+def delete_pgadmin_state_data():
+    trans_id = None
+    if request.data:
+        data = json.loads(request.data)
+        trans_id = int(data['panelId'].split('_')[-1])
+    return delete_tool_data(trans_id)
+
+
+def delete_tool_data(trans_id):
+    try:
+        if trans_id:
+            results = db.session \
+                .query(PgadminStateData) \
+                .filter(PgadminStateData.uid == current_user.id,
+                        PgadminStateData.id == trans_id) \
+                .all()
+        else:
+            results = db.session \
+                .query(PgadminStateData) \
+                .filter(PgadminStateData.uid == current_user.id) \
+                .all()
+        for result in results:
+            db.session.delete(result)
+        db.session.commit()
+        return make_json_response(
+            data={
+                'status': True,
+                'msg': 'Success',
+            }
+        )
+    except Exception as e:
+        db.session.rollback()
+        return make_json_response(
+            data={
+                'status': False,
+                'msg': 'str(e)',
+            }
+        )

web/pgadmin/static/js/PgAdminProvider.jsx

@@ -9,6 +9,9 @@
 
 import React from 'react';
 import PropTypes from 'prop-types';
+import EventBus from './helpers/EventBus';
+import getApiInstance from './api_instance';
+import url_for from 'sources/url_for';
 
 const PgAdminContext = React.createContext();
 
@@ -19,7 +22,18 @@ export function usePgAdmin() {
 
 export function PgAdminProvider({children, value}) {
 
-  return <PgAdminContext.Provider value={value}>
+  const eventBus = React.useRef(new EventBus());
+
+  React.useEffect(()=>{
+    eventBus.current.registerListener('SAVE_TOOL_DATA', (data) => {
+      getApiInstance().post(
+        url_for('settings.save_pgadmin_state'),
+        JSON.stringify(data),
+      ).catch((error)=>{console.error(error);});
+    });
+  }, []);
+
+  return <PgAdminContext.Provider value={{pgAdminProviderEventBus: eventBus.current, ...value}}>
     {children}
   </PgAdminContext.Provider>;
 }

web/pgadmin/static/js/ToolView.jsx

@@ -13,6 +13,9 @@ import { usePgAdmin } from './PgAdminProvider';
 import { BROWSER_PANELS } from '../../browser/static/js/constants';
 import PropTypes from 'prop-types';
 import LayoutIframeTab from './helpers/Layout/LayoutIframeTab';
+import { LAYOUT_EVENTS } from './helpers/Layout';
+import getApiInstance from './api_instance';
+import url_for from 'sources/url_for';
 
 function ToolForm({actionUrl, params}) {
   const formRef = useRef(null);
@@ -56,6 +59,19 @@ export default function ToolView({dockerObj}) {
         // Handler here will return which layout instance the tool should go in
         // case of workspace layout.
         let handler = pgAdmin.Browser.getDockerHandler?.(panelId, dockerObj);
+        const dergisterRemove = handler.docker.eventBus.registerListener(LAYOUT_EVENTS.REMOVE, (closePanelId)=>{
+          if(panelId == closePanelId){
+            let api = getApiInstance();
+            api.delete(
+              url_for('settings.delete_pgadmin_state'), {data:{'panelId': panelId}}
+            ).then(()=> { /* Sona qube */}).catch(function(error) {
+              pgAdmin.Browser.notifier.pgRespErrorNotify(error);
+            });
+            dergisterRemove();
+          }
+        });
+
+
         handler.focus();
         handler.docker.openTab({
           id: panelId,