Reuse schema exploration on REPL to generate database to support same…

… Db to model translation (#3299)
finos · Dec 12, 2024 · feae92b · feae92b
1 parent 323cca3
commit feae92b
Show file tree

Hide file tree

Showing 22 changed files with 236 additions and 357 deletions.
diff --git a/...data-cube/src/main/java/org/finos/legend/engine/repl/dataCube/commands/DataCubeCache.java b/...data-cube/src/main/java/org/finos/legend/engine/repl/dataCube/commands/DataCubeCache.java
@@ -14,6 +14,11 @@
 
 package org.finos.legend.engine.repl.dataCube.commands;
 
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.List;
+import java.util.stream.Collectors;
 import org.eclipse.collections.api.list.MutableList;
 import org.finos.legend.engine.language.pure.compiler.toPureGraph.PureModel;
 import org.finos.legend.engine.plan.execution.result.Result;
@@ -26,25 +31,19 @@
 import org.finos.legend.engine.plan.execution.stores.relational.serialization.RelationalResultToCSVSerializerWithTransformersApplied;
 import org.finos.legend.engine.protocol.pure.v1.model.context.PureModelContextData;
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.SingleExecutionPlan;
-import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.DatabaseConnection;
+import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.RelationalDatabaseConnection;
+import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.model.Table;
 import org.finos.legend.engine.repl.client.Client;
 import org.finos.legend.engine.repl.core.Command;
 import org.finos.legend.engine.repl.dataCube.server.REPLServer;
 import org.finos.legend.engine.repl.relational.shared.ConnectionHelper;
+import static org.finos.legend.engine.repl.shared.ExecutionHelper.REPL_RUN_FUNCTION_SIGNATURE;
+import static org.finos.legend.engine.repl.shared.ExecutionHelper.executeCode;
 import org.finos.legend.engine.shared.core.operational.errorManagement.EngineException;
 import org.jline.reader.Candidate;
 import org.jline.reader.LineReader;
 import org.jline.reader.ParsedLine;
 
-import java.sql.Connection;
-import java.sql.SQLException;
-import java.sql.Statement;
-import java.util.List;
-
-import static org.finos.legend.engine.repl.relational.schema.MetadataReader.getTables;
-import static org.finos.legend.engine.repl.shared.ExecutionHelper.REPL_RUN_FUNCTION_SIGNATURE;
-import static org.finos.legend.engine.repl.shared.ExecutionHelper.executeCode;
-
 public class DataCubeCache implements Command
 {
     private final DataCube parentCommand;
@@ -94,7 +93,7 @@ public boolean process(String line) throws Exception
                 this.client.printError("Failed to retrieve the last command");
                 return true;
             }
-            DatabaseConnection databaseConnection = ConnectionHelper.getDatabaseConnection(this.client.getModelState().parse(), DataCube.getLocalConnectionPath());
+            RelationalDatabaseConnection databaseConnection = ConnectionHelper.getDatabaseConnection(this.client.getModelState().parse(), DataCube.getLocalConnectionPath());
 
             try
             {
@@ -117,9 +116,11 @@ public boolean process(String line) throws Exception
                                 throw new RuntimeException(e);
                             }
 
+                            List<Table> tables = ConnectionHelper.getTables(databaseConnection, client.getPlanExecutor()).collect(Collectors.toList());
+
                             try (Connection connection = ConnectionHelper.getConnection(databaseConnection, client.getPlanExecutor()))
                             {
-                                String tableName = specifiedTableName != null ? specifiedTableName : "test" + (getTables(connection).size() + 1);
+                                String tableName = specifiedTableName != null ? specifiedTableName : "test" + (tables.size() + 1);
                                 try (Statement statement = connection.createStatement())
                                 {
                                     statement.executeUpdate(DatabaseManager.fromString(databaseConnection.type.name()).relationalDatabaseSupport().load(tableName, tempFile.getTemporaryPathForFile(), relationalResultColumns));

diff --git a/...data-cube/src/main/java/org/finos/legend/engine/repl/dataCube/commands/DataCubeTable.java b/...data-cube/src/main/java/org/finos/legend/engine/repl/dataCube/commands/DataCubeTable.java
@@ -14,26 +14,25 @@
 
 package org.finos.legend.engine.repl.dataCube.commands;
 
+import java.util.stream.Collectors;
 import org.eclipse.collections.api.list.MutableList;
 import org.eclipse.collections.impl.factory.Lists;
 import org.eclipse.collections.impl.utility.ListIterate;
 import org.finos.legend.engine.language.pure.grammar.to.PureGrammarComposerUtility;
 import org.finos.legend.engine.protocol.pure.v1.model.context.PureModelContextData;
 import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.connection.PackageableConnection;
 import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.DatabaseConnection;
+import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.RelationalDatabaseConnection;
+import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.model.Database;
 import org.finos.legend.engine.repl.client.Client;
 import org.finos.legend.engine.repl.core.Command;
 import org.finos.legend.engine.repl.dataCube.server.REPLServer;
 import org.finos.legend.engine.repl.relational.shared.ConnectionHelper;
+import static org.finos.legend.engine.repl.shared.ExecutionHelper.REPL_RUN_FUNCTION_SIGNATURE;
 import org.jline.reader.Candidate;
 import org.jline.reader.LineReader;
 import org.jline.reader.ParsedLine;
 
-import java.sql.Connection;
-
-import static org.finos.legend.engine.repl.relational.schema.MetadataReader.getTables;
-import static org.finos.legend.engine.repl.shared.ExecutionHelper.REPL_RUN_FUNCTION_SIGNATURE;
-
 public class DataCubeTable implements Command
 {
     private final DataCube parentCommand;
@@ -101,14 +100,7 @@ public MutableList<Candidate> complete(String inScope, LineReader lineReader, Pa
                         .select(c -> PureGrammarComposerUtility.convertPath(c.getPath()).equals(DataCube.getLocalConnectionPath()));
                 if (!foundConnections.isEmpty() && foundConnections.getFirst().connectionValue instanceof DatabaseConnection)
                 {
-                    try (Connection connection = ConnectionHelper.getConnection((DatabaseConnection) foundConnections.getFirst().connectionValue, client.getPlanExecutor()))
-                    {
-                        return getTables(connection).select(c -> c.name.startsWith(start)).collect(c -> c.name).collect(Candidate::new);
-                    }
-                    catch (Exception e)
-                    {
-                        // do nothing
-                    }
+                    return ConnectionHelper.getTables((RelationalDatabaseConnection) foundConnections.getFirst().connectionValue, client.getPlanExecutor()).filter(c -> c.name.startsWith(start)).map(c -> c.name).map(Candidate::new).collect(Collectors.toCollection(Lists.mutable::empty));
                 }
                 return null;
             }

diff --git a/...ube/src/main/java/org/finos/legend/engine/repl/dataCube/commands/DataCubeWalkthrough.java b/...ube/src/main/java/org/finos/legend/engine/repl/dataCube/commands/DataCubeWalkthrough.java
@@ -14,37 +14,39 @@
 
 package org.finos.legend.engine.repl.dataCube.commands;
 
+import java.io.FileOutputStream;
+import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.sql.Connection;
+import java.sql.Statement;
+import java.util.Objects;
+import java.util.stream.Collectors;
 import org.apache.commons.io.IOUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.eclipse.collections.api.block.function.Function0;
 import org.eclipse.collections.api.factory.Lists;
 import org.eclipse.collections.api.list.MutableList;
 import org.finos.legend.engine.plan.execution.stores.relational.connection.driver.DatabaseManager;
-import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.DatabaseConnection;
+import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.RelationalDatabaseConnection;
+import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.model.Table;
 import org.finos.legend.engine.repl.client.Client;
 import org.finos.legend.engine.repl.core.Command;
 import org.finos.legend.engine.repl.core.commands.Execute;
 import org.finos.legend.engine.repl.dataCube.server.REPLServer;
 import org.finos.legend.engine.repl.dataCube.shared.DataCubeSampleData;
 import org.finos.legend.engine.repl.relational.RelationalReplExtension;
-import org.finos.legend.engine.repl.relational.schema.Table;
 import org.finos.legend.engine.repl.relational.shared.ConnectionHelper;
 import org.finos.legend.engine.repl.shared.DocumentationHelper;
+import static org.finos.legend.engine.repl.shared.REPLHelper.ansiDim;
+import static org.finos.legend.engine.repl.shared.REPLHelper.ansiGreen;
+import static org.finos.legend.engine.repl.shared.REPLHelper.ansiYellow;
+import static org.finos.legend.engine.repl.shared.REPLHelper.getLineWidth;
+import static org.finos.legend.engine.repl.shared.REPLHelper.wrap;
 import org.jline.reader.Candidate;
 import org.jline.reader.LineReader;
 import org.jline.reader.ParsedLine;
 
-import java.io.FileOutputStream;
-import java.io.InputStream;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.sql.Connection;
-import java.sql.Statement;
-import java.util.Objects;
-
-import static org.finos.legend.engine.repl.relational.schema.MetadataReader.getTables;
-import static org.finos.legend.engine.repl.shared.REPLHelper.*;
-
 public class DataCubeWalkthrough implements Command
 {
     private final DataCube parentCommand;
@@ -155,13 +157,15 @@ public static class DataCubeWalkthrough1 extends DocumentationHelper.Walkthrough
         @Override
         public void beforeStep()
         {
-            DatabaseConnection databaseConnection = ConnectionHelper.getDatabaseConnection(this.client.getModelState().parse(), DataCube.getLocalConnectionPath());
+            RelationalDatabaseConnection databaseConnection = ConnectionHelper.getDatabaseConnection(this.client.getModelState().parse(), DataCube.getLocalConnectionPath());
+
+            MutableList<Table> tables = ConnectionHelper.getTables(databaseConnection, this.client.getPlanExecutor()).collect(Collectors.toCollection(Lists.mutable::empty));
+
             try (
                     InputStream inputStream = Thread.currentThread().getContextClassLoader().getResourceAsStream("org/finos/legend/engine/repl/dataCube/walkthrough/sport-data.csv");
                     Connection connection = ConnectionHelper.getConnection(databaseConnection, client.getPlanExecutor());
                     Statement statement = connection.createStatement())
             {
-                MutableList<Table> tables = getTables(connection);
                 if (tables.anySatisfy(table -> table.name.equals(DataCubeSampleData.SPORT.tableName)))
                 {
                     this.tableName = DataCubeSampleData.SPORT.tableName;

diff --git a/...a-cube/src/main/java/org/finos/legend/engine/repl/dataCube/shared/DataCubeSampleData.java b/...a-cube/src/main/java/org/finos/legend/engine/repl/dataCube/shared/DataCubeSampleData.java
@@ -14,16 +14,6 @@
 
 package org.finos.legend.engine.repl.dataCube.shared;
 
-import org.apache.commons.io.IOUtils;
-import org.eclipse.collections.api.list.MutableList;
-import org.eclipse.collections.impl.factory.Lists;
-import org.finos.legend.engine.plan.execution.stores.relational.connection.driver.DatabaseManager;
-import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.DatabaseConnection;
-import org.finos.legend.engine.repl.client.Client;
-import org.finos.legend.engine.repl.dataCube.commands.DataCube;
-import org.finos.legend.engine.repl.relational.schema.Table;
-import org.finos.legend.engine.repl.relational.shared.ConnectionHelper;
-
 import java.io.FileOutputStream;
 import java.io.InputStream;
 import java.nio.file.Files;
@@ -32,8 +22,16 @@
 import java.sql.Statement;
 import java.util.Arrays;
 import java.util.Objects;
-
-import static org.finos.legend.engine.repl.relational.schema.MetadataReader.getTables;
+import java.util.stream.Collectors;
+import org.apache.commons.io.IOUtils;
+import org.eclipse.collections.api.list.MutableList;
+import org.eclipse.collections.impl.factory.Lists;
+import org.finos.legend.engine.plan.execution.stores.relational.connection.driver.DatabaseManager;
+import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.RelationalDatabaseConnection;
+import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.model.Table;
+import org.finos.legend.engine.repl.client.Client;
+import org.finos.legend.engine.repl.dataCube.commands.DataCube;
+import org.finos.legend.engine.repl.relational.shared.ConnectionHelper;
 
 public class DataCubeSampleData
 {
@@ -55,13 +53,13 @@ public DataCubeSampleData(String name, String tableName, String csvFilePath, Mut
 
     public void load(Client client)
     {
-        DatabaseConnection databaseConnection = ConnectionHelper.getDatabaseConnection(client.getModelState().parse(), DataCube.getLocalConnectionPath());
+        RelationalDatabaseConnection databaseConnection = ConnectionHelper.getDatabaseConnection(client.getModelState().parse(), DataCube.getLocalConnectionPath());
+        MutableList<Table> tables = ConnectionHelper.getTables(databaseConnection, client.getPlanExecutor()).collect(Collectors.toCollection(Lists.mutable::empty));
         try (
                 InputStream inputStream = Thread.currentThread().getContextClassLoader().getResourceAsStream(this.csvFilePath);
                 Connection connection = ConnectionHelper.getConnection(databaseConnection, client.getPlanExecutor());
                 Statement statement = connection.createStatement())
         {
-            MutableList<Table> tables = getTables(connection);
             if (tables.anySatisfy(t -> t.name.equals(this.tableName)))
             {
                 statement.executeUpdate(DatabaseManager.fromString(databaseConnection.type.name()).relationalDatabaseSupport().dropTable(tableName));
@@ -71,18 +69,18 @@ public void load(Client client)
             IOUtils.copy(Objects.requireNonNull(inputStream, "Can't extract sample data '" + this.name + "' from " + this.csvFilePath), outputStream);
             outputStream.close(); // explicitly close output stream to let database access this file, else this would throw `java.sql.SQLException: IO Error: File is already open in` error on Windows
             statement.executeUpdate(DatabaseManager.fromString(databaseConnection.type.name()).relationalDatabaseSupport().load(tableName, tempFile.toString()));
-
-            // post check
-            tables = getTables(connection);
-            Table table = tables.detect(t -> t.name.equals(this.tableName));
-            if (!Arrays.equals(table.columns.collect(column -> column.name).toArray(), this.expectedColumns.toArray()))
-            {
-                throw new RuntimeException("Sample data '" + this.name + "' does not have the expected columns " + this.expectedColumns.makeString("(", ",", ")") + " (got: " + table.columns.collect(column -> column.name).makeString(",") + ")");
-            }
         }
         catch (Exception e)
         {
             throw new RuntimeException(e);
         }
+
+        // post check
+        tables = ConnectionHelper.getTables(databaseConnection, client.getPlanExecutor()).collect(Collectors.toCollection(Lists.mutable::empty));;
+        Table table = tables.detect(t -> t.name.equals(this.tableName));
+        if (!Arrays.equals(table.columns.stream().map(column -> column.name).toArray(), this.expectedColumns.toArray()))
+        {
+            throw new RuntimeException("Sample data '" + this.name + "' does not have the expected columns " + this.expectedColumns.makeString("(", ",", ")") + " (got: " + table.columns.stream().map(column -> column.name).collect(Collectors.joining(",")) + ")");
+        }
     }
 }
diff --git a/...repl-relational/src/main/java/org/finos/legend/engine/repl/relational/commands/Cache.java b/...repl-relational/src/main/java/org/finos/legend/engine/repl/relational/commands/Cache.java
@@ -14,6 +14,11 @@
 
 package org.finos.legend.engine.repl.relational.commands;
 
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.List;
+import java.util.stream.Collectors;
 import org.eclipse.collections.api.list.MutableList;
 import org.eclipse.collections.impl.factory.Lists;
 import org.eclipse.collections.impl.utility.ListIterate;
@@ -30,24 +35,18 @@
 import org.finos.legend.engine.protocol.pure.v1.model.context.PureModelContextData;
 import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.SingleExecutionPlan;
 import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.connection.PackageableConnection;
-import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.DatabaseConnection;
+import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.RelationalDatabaseConnection;
+import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.model.Table;
 import org.finos.legend.engine.repl.client.Client;
 import org.finos.legend.engine.repl.core.Command;
 import org.finos.legend.engine.repl.relational.shared.ConnectionHelper;
+import static org.finos.legend.engine.repl.shared.ExecutionHelper.executeCode;
+import static org.finos.legend.engine.repl.shared.ExecutionHelper.printExecutionTime;
 import org.finos.legend.engine.shared.core.operational.errorManagement.EngineException;
 import org.jline.reader.Candidate;
 import org.jline.reader.LineReader;
 import org.jline.reader.ParsedLine;
 
-import java.sql.Connection;
-import java.sql.SQLException;
-import java.sql.Statement;
-import java.util.List;
-
-import static org.finos.legend.engine.repl.relational.schema.MetadataReader.getTables;
-import static org.finos.legend.engine.repl.shared.ExecutionHelper.executeCode;
-import static org.finos.legend.engine.repl.shared.ExecutionHelper.printExecutionTime;
-
 public class Cache implements Command
 {
     private final Client client;
@@ -88,7 +87,7 @@ public boolean process(String line) throws Exception
                 this.client.printError("Failed to retrieve the last command");
                 return true;
             }
-            DatabaseConnection databaseConnection = ConnectionHelper.getDatabaseConnection(this.client.getModelState().parse(), connectionPath);
+            RelationalDatabaseConnection databaseConnection = ConnectionHelper.getDatabaseConnection(this.client.getModelState().parse(), connectionPath);
 
             try
             {
@@ -112,9 +111,11 @@ public boolean process(String line) throws Exception
                                 throw new RuntimeException(e);
                             }
 
+                            List<Table> tables = ConnectionHelper.getTables(databaseConnection, client.getPlanExecutor()).collect(Collectors.toList());
+
                             try (Connection connection = ConnectionHelper.getConnection(databaseConnection, client.getPlanExecutor()))
                             {
-                                String tableName = specifiedTableName != null ? specifiedTableName : "test" + (getTables(connection).size() + 1);
+                                String tableName = specifiedTableName != null ? specifiedTableName : "test" + (tables.size() + 1);
                                 try (Statement statement = connection.createStatement())
                                 {
                                     statement.executeUpdate(DatabaseManager.fromString(databaseConnection.type.name()).relationalDatabaseSupport().load(tableName, tempFile.getTemporaryPathForFile(), relationalResultColumns));