From 97a87ea2e7143b892bb06acea8e8b6ca20030950 Mon Sep 17 00:00:00 2001
From: Serge Rider <serge@jkiss.org>
Date: Wed, 27 Jun 2018 19:33:55 +0300
Subject: [PATCH] ERD: add table tool

Former-commit-id: e8b52e06d7850dcd6a312070e1309993cf99808f
---
 .../ext/erd/command/EntityAddCommand.java     | 41 +++++++++++++++++++
 .../dbeaver/ext/erd/model/ERDEntity.java      | 15 +++++++
 2 files changed, 56 insertions(+)

diff --git a/plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/command/EntityAddCommand.java b/plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/command/EntityAddCommand.java
index 5301e0d81e..b8025d006b 100644
--- a/plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/command/EntityAddCommand.java
+++ b/plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/command/EntityAddCommand.java
@@ -21,8 +21,18 @@ import org.eclipse.draw2d.geometry.Point;
 import org.eclipse.draw2d.geometry.Rectangle;
 import org.eclipse.gef.commands.Command;
 import org.jkiss.dbeaver.ext.erd.model.ERDEntity;
+import org.jkiss.dbeaver.ext.erd.model.ERDUtils;
 import org.jkiss.dbeaver.ext.erd.part.DiagramPart;
 import org.jkiss.dbeaver.ext.erd.part.EntityPart;
+import org.jkiss.dbeaver.model.navigator.DBNDatabaseNode;
+import org.jkiss.dbeaver.model.navigator.DBNNode;
+import org.jkiss.dbeaver.model.runtime.VoidProgressMonitor;
+import org.jkiss.dbeaver.model.struct.DBSEntity;
+import org.jkiss.dbeaver.model.struct.DBSObjectContainer;
+import org.jkiss.dbeaver.model.struct.rdb.DBSTable;
+import org.jkiss.dbeaver.ui.UIUtils;
+import org.jkiss.dbeaver.ui.dialogs.BrowseObjectDialog;
+import org.jkiss.dbeaver.ui.navigator.NavigatorUtils;
 
 import java.util.Collection;
 
@@ -46,10 +56,41 @@ public class EntityAddCommand extends Command
     @Override
     public void execute()
 	{
+        VoidProgressMonitor monitor = new VoidProgressMonitor();
+
         Point curLocation = location == null ? null : new Point(location);
         for (ERDEntity entity : entities) {
+            boolean resolveRelations = false;
+            if (entity.getObject() == null) {
+                // Entity is not initialized
+                if (entity.getDataSource() != null) {
+                    DBNDatabaseNode dsNode = NavigatorUtils.getNodeByObject(entity.getDataSource().getContainer());
+                    if (dsNode != null) {
+                        DBNNode tableNode = BrowseObjectDialog.selectObject(
+                                UIUtils.getActiveWorkbenchShell(),
+                                "Select a table",
+                                dsNode,
+                                null,
+                                new Class[]{DBSTable.class, DBSObjectContainer.class},
+                                new Class[]{DBSTable.class});
+                        if (tableNode instanceof DBNDatabaseNode && ((DBNDatabaseNode) tableNode).getObject() instanceof DBSEntity) {
+                            entity = ERDUtils.makeEntityFromObject(
+                                    monitor,
+                                    diagramPart.getDiagram(),
+                                    (DBSEntity)((DBNDatabaseNode) tableNode).getObject());
+                            resolveRelations = true;
+                        }
+                    }
+                }
+            }
+            if (entity.getObject() == null) {
+                continue;
+            }
 		    diagramPart.getDiagram().addEntity(entity, true);
             //diagramPart.getDiagram().addModelRelations(monitor, entity, true);
+            if (resolveRelations) {
+                entity.addModelRelations(monitor, diagramPart.getDiagram().getEntityMap(), false, true);
+            }
 
             if (curLocation != null) {
                 // Put new entities in specified location
diff --git a/plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/model/ERDEntity.java b/plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/model/ERDEntity.java
index e068e91f38..8718c6bb86 100644
--- a/plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/model/ERDEntity.java
+++ b/plugins/org.jkiss.dbeaver.ext.erd/src/org/jkiss/dbeaver/ext/erd/model/ERDEntity.java
@@ -22,6 +22,7 @@ package org.jkiss.dbeaver.ext.erd.model;
 import org.jkiss.code.NotNull;
 import org.jkiss.dbeaver.DBException;
 import org.jkiss.dbeaver.Log;
+import org.jkiss.dbeaver.model.DBPDataSource;
 import org.jkiss.dbeaver.model.DBUtils;
 import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
 import org.jkiss.dbeaver.model.struct.*;
@@ -39,6 +40,7 @@ public class ERDEntity extends ERDObject<DBSEntity> {
 
     static final Log log = Log.getLog(ERDEntity.class);
 
+    private DBPDataSource dataSource;
     private String alias;
     private List<ERDEntityAttribute> attributes;
 
@@ -48,10 +50,23 @@ public class ERDEntity extends ERDObject<DBSEntity> {
 
     private boolean primary = false;
 
+    /**
+     * Special constructore for creating lazy entities.
+     * This entity will be initialized at the moment of creation within diagram.
+     */
+    public ERDEntity(DBPDataSource dataSource) {
+        super(null);
+        this.dataSource = dataSource;
+    }
+
     public ERDEntity(DBSEntity entity) {
         super(entity);
     }
 
+    public DBPDataSource getDataSource() {
+        return dataSource;
+    }
+
     public String getAlias() {
         return alias;
     }
-- 
GitLab