Insert tables with relations on drag and drop

Co-authored-by: Christian P. <pirnichristian@gmail.com>

Author: Lancear

web/pgadmin/browser/server_groups/servers/databases/schemas/tables/static/js/table.ui.js

@@ -416,7 +416,7 @@ export default class TableSchema extends BaseUISchema {
   static getErdSupportedData(data) {
     let newData = {...data};
     const SUPPORTED_KEYS = [
-      'name', 'schema', 'description', 'rlspolicy', 'forcerlspolicy', 'fillfactor',
+      'oid', 'name', 'schema', 'description', 'rlspolicy', 'forcerlspolicy', 'fillfactor',
       'toast_tuple_target', 'parallel_workers', 'relhasoids', 'relpersistence',
       'columns', 'primary_key', 'foreign_key', 'unique_constraint',
     ];
@@ -428,14 +428,18 @@ export default class TableSchema extends BaseUISchema {
       return c;
     });
 
-    /* Make autoindex as true if there is coveringindex since ERD works in create mode */
-    newData.foreign_key = (newData.foreign_key||[]).map((fk)=>{
+    newData.original_foreign_keys = (newData.foreign_key||[]).map((fk)=>{
+      /* Make autoindex as true if there is coveringindex since ERD works in create mode */
       fk.autoindex = false;
+
       if(fk.coveringindex) {
         fk.autoindex = true;
       }
+
       return fk;
     });
+
+    newData.foreign_key = [];
     return newData;
   }
 

web/pgadmin/tools/erd/static/js/erd_tool/ERDCore.js

@@ -400,6 +400,7 @@ export default class ERDCore {
 
     const addLink = (theFk)=>{
       if(!theFk) return;
+
       let newData = {
         local_table_uid: tableNode.getID(),
         local_column_attnum: undefined,
@@ -590,7 +591,7 @@ export default class ERDCore {
   }
 
   cloneTableData(tableData, name) {
-    const SKIP_CLONE_KEYS = ['foreign_key'];
+    const SKIP_CLONE_KEYS = ['oid', 'foreign_key', 'original_foreign_keys'];
 
     if(!tableData) {
       return tableData;
@@ -635,43 +636,95 @@ export default class ERDCore {
 
   deserializeData(data){
     let oidUidMap = {};
+    let newNodes = [];
 
     /* Add the nodes */
     data.forEach((nodeData)=>{
-      let newNode = this.addNode(TableSchema.getErdSupportedData(nodeData));
+      const newNode = this.addNode(TableSchema.getErdSupportedData(nodeData));
       oidUidMap[nodeData.oid] = newNode.getID();
+      newNodes.push(newNode);
     });
 
-    /* Lets use the oidUidMap for creating the links */
-    let tableNodesDict = this.getModel().getNodesDict();
+    // When generating for schema, there may be a reference to another schema table
+    // We'll remove the FK completely in such cases
+    newNodes.forEach((node) => {
+      const nodeData = node.getData();
+      nodeData.original_foreign_keys = nodeData.original_foreign_keys?.filter(fk => 
+        fk.columns?.[0]?.references && oidUidMap[fk.columns[0].references]
+      );
+    });
+
+    this.addLinksBetweenNodes(oidUidMap);
+  }
+
+  addNodeWithLinks(nodeData, position=[50,50], metadata={}){
+    const tableNodesDict = this.getModel().getNodesDict();
+    const oidExists = Object.values(tableNodesDict).some(node => node.getData().oid === nodeData.oid);
+
+    if (oidExists) {
+      delete nodeData.oid;
+    }
+
+    let oidUidMap = {};
+    const newNode = this.addNode(nodeData, position, metadata);
+    
+    if (!oidExists) {
+      oidUidMap[nodeData.oid] = newNode.getID();
+    }
+
     _.forIn(tableNodesDict, (node, uid)=>{
-      let nodeData = node.getData();
-      if(nodeData.foreign_key) {
-        nodeData.foreign_key = nodeData.foreign_key.filter((theFk)=>{
-          delete theFk.oid;
-          theFk = theFk.columns[0];
-          theFk.references = oidUidMap[theFk.references];
-          let newData = {
-            local_table_uid: uid,
-            local_column_attnum: undefined,
-            referenced_table_uid: theFk.references,
-            referenced_column_attnum: undefined,
-          };
-          let sourceNode = tableNodesDict[newData.referenced_table_uid];
-          let targetNode = tableNodesDict[newData.local_table_uid];
-          // When generating for schema, there may be a reference to another schema table
-          // We'll remove the FK completely in such cases.
-          if(!sourceNode || !targetNode) {
-            return false;
-          }
+      const oid = node.getData().oid;
+      if (!oid) return;
 
-          newData.local_column_attnum = _.find(targetNode.getColumns(), (col)=>col.name==theFk.local_column).attnum;
-          newData.referenced_column_attnum = _.find(sourceNode.getColumns(), (col)=>col.name==theFk.referenced).attnum;
+      oidUidMap[oid] = uid;
+    });
 
-          this.addLink(newData, 'onetomany');
-          return true;
-        });
-      }
+    this.addLinksBetweenNodes(oidUidMap, [newNode.getID()]);
+    return newNode;
+  }
+
+  addLinksBetweenNodes(oidUidMap, newNodesUids = null) {
+    const tableNodesDict = this.getModel().getNodesDict();
+
+    _.forIn(tableNodesDict, (node, uid)=>{
+      const nodeData = node.getData();
+
+      nodeData.original_foreign_keys?.forEach((theFk)=>{
+        const theFkColumn = theFk.columns[0];
+        let referencesUid = oidUidMap[theFkColumn.references];
+
+        /* Incomplete reference to missing table */
+        if (!referencesUid) {
+          return;
+        }
+
+        /* Avoid creating duplicate links */
+        if (
+          newNodesUids 
+          && !newNodesUids.includes(uid) 
+          && !newNodesUids.includes(referencesUid)
+        ) {
+          return;
+        }
+
+        const newData = {
+          local_table_uid: uid,
+          local_column_attnum: _.find(
+            tableNodesDict[uid].getColumns(), 
+            (col) => col.name == theFkColumn.local_column
+          ).attnum,
+          referenced_table_uid: referencesUid,
+          referenced_column_attnum: _.find(
+            tableNodesDict[referencesUid].getColumns(), 
+            (col) => col.name == theFkColumn.referenced
+          ).attnum,
+        };
+
+        const newForeignKey = _.cloneDeep(theFk);
+        newForeignKey.columns[0].references = referencesUid;
+        nodeData.foreign_key.push(newForeignKey);
+        this.addLink(newData, 'onetomany');
+      });
     });
   }
 

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

@@ -596,21 +596,26 @@ export default class ERDTool extends React.Component {
       if(nodeDropData.objUrl.indexOf(matchUrl) == -1) {
         pgAdmin.Browser.notifier.error(gettext('Cannot drop table from outside of the current database.'));
       } else {
-        let dataPromise = new Promise((resolve, reject)=>{
-          this.apiObj.get(nodeDropData.objUrl)
-            .then((res)=>{
-              resolve(this.diagram.cloneTableData(TableSchema.getErdSupportedData(res.data)));
-            })
-            .catch((err)=>{
-              console.error(err);
-              reject(err instanceof Error ? err : Error(gettext('Something went wrong')));
-            });
-        });
-        const {x, y} = this.diagram.getEngine().getRelativeMousePoint(e);
-        this.diagram.addNode(dataPromise, [x, y], {
-          fillColor: this.state.fill_color,
-          textColor: this.state.text_color,
-        }).setSelected(true);
+        this.apiObj.get(nodeDropData.objUrl)
+          .then((res)=>{
+            const data = TableSchema.getErdSupportedData(res.data);
+            const {x, y} = this.diagram.getEngine().getRelativeMousePoint(e);
+            const position = [x,y];
+            const metadata = {
+              fillColor: this.state.fill_color,
+              textColor: this.state.text_color,
+            };
+
+            const newNode = this.state.preferences.insert_table_with_relations
+              ? this.diagram.addNodeWithLinks(data, position, metadata)
+              : this.diagram.addNode(this.diagram.cloneTableData(data), position, metadata);
+
+            newNode.setSelected(true);
+          })
+          .catch((err)=>{
+            console.error(err);
+            throw (err instanceof Error ? err : Error(gettext('Something went wrong')));
+          });
       }
     }
   }