Use SocketIO instead of REST for fetching database tables data in ERD. #5065

docs/en_US/release_notes_6_15.rst

@@ -25,6 +25,7 @@ New features
 Housekeeping
 ************
 
+  | `Issue #5065 <https://github.com/pgadmin-org/pgadmin4/issues/5065>`_ -  Use SocketIO instead of REST for fetching database tables data in ERD.
   | `Issue #5293 <https://github.com/pgadmin-org/pgadmin4/issues/5293>`_ -  Ensure that the tooltips are consistent throughout the entire application.
   | `Issue #5357 <https://github.com/pgadmin-org/pgadmin4/issues/5357>`_ -  Remove Python's 'Six' package completely.
 

web/pgadmin/authenticate/__init__.py

@@ -11,12 +11,16 @@
 
 import config
 import copy
+import functools
 
 from flask import current_app, flash, Response, request, url_for, \
     session, redirect
 from flask_babel import gettext
 from flask_security.views import _security
 from flask_security.utils import get_post_logout_redirect, logout_user
+from flask_login import current_user
+from flask_socketio import disconnect, ConnectionRefusedError
+
 
 from pgadmin.model import db, User
 from pgadmin.utils import PgAdminModule, get_safe_post_login_redirect
@@ -52,6 +56,17 @@ def get_logout_url() -> str:
         url_for('security.logout'), url_for(BROWSER_INDEX))
 
 
+def socket_login_required(f):
+    @functools.wraps(f)
+    def wrapped(*args, **kwargs):
+        if not current_user.is_authenticated:
+            disconnect()
+            raise ConnectionRefusedError("Unauthorised !")
+        else:
+            return f(*args, **kwargs)
+    return wrapped
+
+
 class AuthenticateModule(PgAdminModule):
     def get_exposed_url_endpoints(self):
         return ['authenticate.login']

web/pgadmin/static/js/socket_instance.js

@@ -0,0 +1,45 @@
+/////////////////////////////////////////////////////////////
+//
+// pgAdmin 4 - PostgreSQL Tools
+//
+// Copyright (C) 2013 - 2022, The pgAdmin Development Team
+// This software is released under the PostgreSQL Licence
+//
+//////////////////////////////////////////////////////////////
+import { io } from 'socketio';
+
+export function openSocket(namespace, options) {
+  return new Promise((resolve, reject)=>{
+    const socketObj = io(namespace, {
+      pingTimeout: 120000,
+      pingInterval: 25000,
+      ...options,
+    });
+
+    /* Once the object is created, connect event is emitted.
+    Backend must implement connect and emit connected on success,
+    connect_error on failure.
+    */
+    socketObj.on('connected', ()=>{
+      resolve(socketObj);
+    });
+    socketObj.on('connect_error', ()=>{
+      reject();
+    });
+    socketObj.on('disconnect', ()=>{
+      reject();
+    });
+  });
+}
+
+export function socketApiGet(socket, endpoint, params) {
+  return new Promise((resolve, reject)=>{
+    socket.emit(endpoint, params);
+    socket.on(`${endpoint}_success`, (data)=>{
+      resolve(data);
+    });
+    socket.on(`${endpoint}_failed`, (data)=>{
+      reject(data);
+    });
+  });
+}

web/pgadmin/tools/erd/__init__.py

@@ -31,8 +31,11 @@ from pgadmin.utils.constants import PREF_LABEL_KEYBOARD_SHORTCUTS, \
     PREF_LABEL_DISPLAY, PREF_LABEL_OPTIONS
 from .utils import ERDHelper
 from pgadmin.utils.exception import ConnectionLost
+from pgadmin.authenticate import socket_login_required
+from ... import socketio
 
 MODULE_NAME = 'erd'
+SOCKETIO_NAMESPACE = '/{0}'.format(MODULE_NAME)
 
 
 class ERDModule(PgAdminModule):
@@ -601,22 +604,36 @@ def sql(trans_id, sgid, sid, did):
     )
 
 
-@blueprint.route('/tables/<int:trans_id>/<int:sgid>/<int:sid>/<int:did>',
-                 methods=["GET"],
-                 endpoint='tables')
-@login_required
-def tables(trans_id, sgid, sid, did):
-    helper = ERDHelper(trans_id, sid, did)
-    _get_connection(sid, did, trans_id)
-    status, tables = helper.get_all_tables()
+@socketio.on('connect', namespace=SOCKETIO_NAMESPACE)
+def connect():
+    """
+    Connect to the server through socket.
+    :return:
+    :rtype:
+    """
+    socketio.emit('connected', {'sid': request.sid},
+                  namespace=SOCKETIO_NAMESPACE,
+                  to=request.sid)
 
-    if not status:
-        return internal_server_error(errormsg=tables)
 
-    return make_json_response(
-        data=tables,
-        status=200
-    )
+@socketio.on('tables', namespace=SOCKETIO_NAMESPACE)
+@socket_login_required
+def tables(params):
+    try:
+        helper = ERDHelper(params['trans_id'], params['sid'], params['did'])
+        _get_connection(params['sid'], params['did'], params['trans_id'])
+        status, tables = helper.get_all_tables()
+
+        if not status:
+            socketio.emit('tables_failed', tables,
+                          namespace=SOCKETIO_NAMESPACE,
+                          to=request.sid)
+            return internal_server_error(errormsg=tables)
+        socketio.emit('tables_success', tables, namespace=SOCKETIO_NAMESPACE,
+                      to=request.sid)
+    except Exception as e:
+        socketio.emit('tables_failed', str(e), namespace=SOCKETIO_NAMESPACE,
+                      to=request.sid)
 
 
 @blueprint.route('/close/<int:trans_id>/<int:sgid>/<int:sid>/<int:did>',

web/pgadmin/tools/erd/static/js/erd_tool/components/ERDTool.jsx

@@ -32,6 +32,7 @@ import { Box, withStyles } from '@material-ui/core';
 import EventBus from '../../../../../../static/js/helpers/EventBus';
 import { ERD_EVENTS } from '../ERDConstants';
 import getApiInstance, { parseApiError } from '../../../../../../static/js/api_instance';
+import { openSocket, socketApiGet } from '../../../../../../static/js/socket_instance';
 
 /* Custom react-diagram action for keyboard events */
 export class KeyboardShortcutAction extends Action {
@@ -897,23 +898,22 @@ class ERDTool extends React.Component {
 
   async loadTablesData() {
     this.setLoading(gettext('Fetching schema data...'));
-    let url = url_for('erd.tables', {
-      trans_id: this.props.params.trans_id,
-      sgid: this.props.params.sgid,
-      sid: this.props.params.sid,
-      did: this.props.params.did,
-    });
-
+    let resData = [];
+    let socket;
     try {
-      let response = await this.apiObj.get(url);
-      this.diagram.deserializeData(response.data.data);
-      return true;
+      socket = await openSocket('/erd');
+      resData = await socketApiGet(socket, 'tables', {
+        trans_id: parseInt(this.props.params.trans_id),
+        sgid: parseInt(this.props.params.sgid),
+        sid: parseInt(this.props.params.sid),
+        did: parseInt(this.props.params.did),
+      });
     } catch (error) {
       this.handleAxiosCatch(error);
-      return false;
-    } finally {
-      this.setLoading(null);
     }
+    socket?.disconnect();
+    this.diagram.deserializeData(resData);
+    this.setLoading(null);
   }
 
   render() {

web/pgadmin/tools/erd/tests/test_tables.py

@@ -10,7 +10,7 @@
 import json
 import uuid
 import secrets
-from pgadmin.utils.route import BaseTestGenerator
+from pgadmin.utils.route import BaseTestGenerator, BaseSocketTestGenerator
 from regression.python_test_utils import test_utils as utils
 from regression import parent_node_dict
 from regression.test_setup import config_data
@@ -20,9 +20,12 @@ from pgadmin.browser.server_groups.servers.databases.schemas.tables.tests \
     import utils as tables_utils
 from pgadmin.browser.server_groups.servers.databases.schemas.tests import \
     utils as schema_utils
+from pgAdmin4 import app
+from .... import socketio
 
 
-class ERDTables(BaseTestGenerator):
+class ERDTables(BaseSocketTestGenerator):
+    SOCKET_NAMESPACE = '/erd'
 
     def dropDB(self):
         connection = utils.get_db_connection(self.server['db'],
@@ -33,6 +36,7 @@ class ERDTables(BaseTestGenerator):
         utils.drop_database(connection, self.db_name)
 
     def setUp(self):
+        super(ERDTables, self).setUp()
         self.db_name = "erdtestdb_{0}".format(str(uuid.uuid4())[1:8])
         self.sid = parent_node_dict["server"][-1]["server_id"]
         self.did = utils.create_database(self.server, self.db_name)
@@ -57,24 +61,31 @@ class ERDTables(BaseTestGenerator):
             raise
 
     def runTest(self):
+        received = self.socket_client.get_received(self.SOCKET_NAMESPACE)
+        assert received[0]['name'] == 'connected'
+
         db_con = database_utils.connect_database(self,
                                                  self.sgid,
                                                  self.sid,
                                                  self.did)
-
         if not db_con["info"] == "Database connected.":
             raise Exception("Could not connect to database to add the schema.")
 
-        trans_id = secrets.choice(range(1, 9999999))
-        url = '/erd/tables/{trans_id}/{sgid}/{sid}/{did}'.format(
-            trans_id=trans_id, sgid=self.sgid, sid=self.sid, did=self.did)
+        data = {
+            "trans_id": secrets.choice(range(1, 9999999)),
+            'sgid': self.sgid,
+            'sid': self.sid,
+            'did': self.did,
+        }
 
-        response = self.tester.get(url)
-        self.assertEqual(response.status_code, 200)
-
-        response = json.loads(response.data.decode('utf-8'))
+        self.socket_client.emit('tables', data,
+                                namespace=self.SOCKET_NAMESPACE)
+        received = self.socket_client.get_received(self.SOCKET_NAMESPACE)
+        response_data = received[0]['args'][0]
+        self.assertEqual(received[0]['name'], "tables_success", response_data)
         self.assertEqual(self.tables, [[tab['schema'], tab['name']]
-                                       for tab in response['data']])
+                                       for tab in response_data])
 
     def tearDown(self):
+        super(ERDTables, self).tearDown()
         self.dropDB()

web/pgadmin/utils/route.py

@@ -14,6 +14,7 @@ from importlib import import_module
 
 from werkzeug.utils import find_modules
 from pgadmin.utils import server_utils
+from .. import socketio
 
 import unittest
 
@@ -172,3 +173,25 @@ class BaseTestGenerator(unittest.TestCase, metaclass=TestsGeneratorRegistry):
     @classmethod
     def setForModules(cls, for_modules):
         cls.for_modules = for_modules
+
+
+class BaseSocketTestGenerator(BaseTestGenerator):
+    SOCKET_NAMESPACE = ""
+
+    def setUp(self):
+        super(BaseSocketTestGenerator, self).setUp()
+        # flask_client = self.app.test_client()
+        self.tester.get("/")
+        self.socket_client = socketio.test_client(
+            self.app, namespace=self.SOCKET_NAMESPACE,
+            flask_test_client=self.tester)
+        self.assertTrue(self.socket_client.is_connected(self.SOCKET_NAMESPACE))
+
+    def runTest(self):
+        super(BaseSocketTestGenerator, self).runTest()
+
+    def tearDown(self):
+        super(BaseSocketTestGenerator, self).tearDown()
+        self.socket_client.disconnect(namespace=self.SOCKET_NAMESPACE)
+        self.assertFalse(
+            self.socket_client.is_connected(self.SOCKET_NAMESPACE))

web/regression/javascript/erd/ui_components/ERDTool.spec.js

@@ -133,6 +133,7 @@ describe('ERDTool', ()=>{
     bodyInstance = body.instance();
     spyOn(bodyInstance, 'getDialog').and.callFake(getDialog);
     spyOn(bodyInstance, 'onChangeColors').and.callFake(()=>{/*no op*/});
+    spyOn(bodyInstance, 'loadTablesData').and.callFake(()=>{/*no op*/});
   });
 
   afterAll(() => {

web/webpack.test.config.js

@@ -129,6 +129,7 @@ module.exports = {
       'backbone': path.join(__dirname, './node_modules/backbone/backbone'),
       'react': path.join(__dirname, 'node_modules/react'),
       'react-dom': path.join(__dirname, 'node_modules/react-dom'),
+      'socketio': path.join(__dirname, './node_modules/socket.io-client/dist/socket.io.js'),
       'sources': sourcesDir + '/js',
       'translations': regressionDir + '/javascript/fake_translations',
       'pgadmin.browser.messages': regressionDir + '/javascript/fake_messages',