Merge pull request #3077 from dbeaver/#1599-mock-data_21

#1599 Mock Data. FK processing is fixed.

plugins/org.jkiss.dbeaver.ext.mockdata/src/org/jkiss/dbeaver/ext/mockdata/MockDataExecuteWizard.java

@@ -44,8 +44,8 @@ import org.jkiss.utils.CommonUtils;
 import java.io.IOException;
 import java.util.*;
 
-public class MockDataExecuteWizard  extends AbstractToolWizard<DBSDataManipulator, DBSDataManipulator> implements IImportWizard{
-
+public class MockDataExecuteWizard  extends AbstractToolWizard<DBSDataManipulator, DBSDataManipulator> implements IImportWizard
+{
     private static final Log log = Log.getLog(MockDataExecuteWizard.class);
 
     private static final int BATCH_SIZE = 1000;

plugins/org.jkiss.dbeaver.ext.mockdata/src/org/jkiss/dbeaver/ext/mockdata/generator/FKGenerator.java

@@ -18,6 +18,8 @@
 package org.jkiss.dbeaver.ext.mockdata.generator;
 
 import org.jkiss.dbeaver.DBException;
+import org.jkiss.dbeaver.Log;
+import org.jkiss.dbeaver.model.DBPNamedObject;
 import org.jkiss.dbeaver.model.DBUtils;
 import org.jkiss.dbeaver.model.data.DBDLabelValuePair;
 import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
@@ -32,6 +34,8 @@ import java.util.Map;
 
 public class FKGenerator extends AbstractMockValueGenerator
 {
+    private static final Log log = Log.getLog(FKGenerator.class);
+
     private static final int UNIQ_REF_RECORDS_LIMIT = 100000000;
     private static final int REF_RECORDS_LIMIT = 100000;
 
@@ -58,7 +62,16 @@ public class FKGenerator extends AbstractMockValueGenerator
             DBSEntityReferrer fk = attributeReferrers.get(0); // TODO only the first
             List<? extends DBSEntityAttributeRef> references = ((DBSEntityReferrer) fk).getAttributeReferences(monitor);
 
-            DBSTableForeignKeyColumn column = (DBSTableForeignKeyColumn) references.iterator().next(); // TODO only the first !!!
+            DBSTableForeignKeyColumn column = null;
+            for (DBSEntityAttributeRef ref : references) {
+                if (((DBPNamedObject) ref).getName().equals(attribute.getName())) {
+                    column = (DBSTableForeignKeyColumn) ref;
+                }
+            }
+            if (column == null) {
+                throw new DBException("Can't find reference column for '" + attribute.getName() + "'");
+            }
+
             int numberRefRecords = DBUtils.checkUnique(monitor, dbsEntity, attribute) ? UNIQ_REF_RECORDS_LIMIT : REF_RECORDS_LIMIT;
             Collection<DBDLabelValuePair> values = readColumnValues(monitor, fk.getDataSource(), (DBSAttributeEnumerable) column.getReferencedColumn(), numberRefRecords);
             for (DBDLabelValuePair value : values) {