From 3f598990409d905755322301693851ba4861b551 Mon Sep 17 00:00:00 2001 From: Sai Sriharsha Annepu <72639930+gs-ssh16@users.noreply.github.com> Date: Sun, 22 Sep 2024 23:30:22 +0530 Subject: [PATCH] Use result set metadata types for REPL caching (#3116) * Use result set metadata types for REPL caching * Update TestDuckDBCommands.java --- .../repl/dataCube/commands/DataCubeCache.java | 5 +- .../repl/relational/commands/Cache.java | 5 +- .../pom.xml | 28 +++ .../driver/vendors/duckdb/DuckDBCommands.java | 8 + .../vendors/duckdb/TestDuckDBCommands.java | 215 ++++++++++++++++++ .../src/test/resources/personTable.csv | 8 + .../commands/RelationalDatabaseCommands.java | 5 + .../driver/vendors/h2/H2Commands.java | 6 + .../relational/result/RelationalResult.java | 25 ++ 9 files changed, 303 insertions(+), 2 deletions(-) create mode 100644 legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/src/test/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/vendors/duckdb/TestDuckDBCommands.java create mode 100644 legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/src/test/resources/personTable.csv diff --git a/legend-engine-config/legend-engine-repl/legend-engine-repl-data-cube/src/main/java/org/finos/legend/engine/repl/dataCube/commands/DataCubeCache.java b/legend-engine-config/legend-engine-repl/legend-engine-repl-data-cube/src/main/java/org/finos/legend/engine/repl/dataCube/commands/DataCubeCache.java index a45a5175eaf..75aad4faa1f 100644 --- a/legend-engine-config/legend-engine-repl/legend-engine-repl-data-cube/src/main/java/org/finos/legend/engine/repl/dataCube/commands/DataCubeCache.java +++ b/legend-engine-config/legend-engine-repl/legend-engine-repl-data-cube/src/main/java/org/finos/legend/engine/repl/dataCube/commands/DataCubeCache.java @@ -20,6 +20,7 @@ import org.finos.legend.engine.plan.execution.result.serialization.TemporaryFile; import org.finos.legend.engine.plan.execution.stores.StoreType; import org.finos.legend.engine.plan.execution.stores.relational.connection.driver.DatabaseManager; +import org.finos.legend.engine.plan.execution.stores.relational.connection.driver.commands.Column; import org.finos.legend.engine.plan.execution.stores.relational.plugin.RelationalStoreState; import org.finos.legend.engine.plan.execution.stores.relational.result.RelationalResult; import org.finos.legend.engine.plan.execution.stores.relational.serialization.RelationalResultToCSVSerializerWithTransformersApplied; @@ -38,6 +39,7 @@ 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; @@ -101,6 +103,7 @@ public boolean process(String line) throws Exception if (res instanceof RelationalResult) { RelationalResult relationalResult = (RelationalResult) res; + List relationalResultColumns = relationalResult.getResultSetColumns(); String tempDir = ((RelationalStoreState) this.client.getPlanExecutor().getExecutorsOfType(StoreType.Relational).getOnly().getStoreState()).getRelationalExecutor().getRelationalExecutionConfiguration().tempPath; try (TemporaryFile tempFile = new TemporaryFile(tempDir)) { @@ -119,7 +122,7 @@ public boolean process(String line) throws Exception String tableName = specifiedTableName != null ? specifiedTableName : "test" + (getTables(connection).size() + 1); try (Statement statement = connection.createStatement()) { - statement.executeUpdate(DatabaseManager.fromString(databaseConnection.type.name()).relationalDatabaseSupport().load(tableName, tempFile.getTemporaryPathForFile())); + statement.executeUpdate(DatabaseManager.fromString(databaseConnection.type.name()).relationalDatabaseSupport().load(tableName, tempFile.getTemporaryPathForFile(), relationalResultColumns)); this.client.println("Cached into table: '" + tableName + "'. Launching DataCube..."); String functionBodyCode = "#>{" + DataCube.getLocalDatabasePath() + "." + tableName + "}#->select()->from(" + DataCube.getLocalRuntimePath() + ")"; diff --git a/legend-engine-config/legend-engine-repl/legend-engine-repl-relational/src/main/java/org/finos/legend/engine/repl/relational/commands/Cache.java b/legend-engine-config/legend-engine-repl/legend-engine-repl-relational/src/main/java/org/finos/legend/engine/repl/relational/commands/Cache.java index 50444ff3191..f8d9a71d106 100644 --- a/legend-engine-config/legend-engine-repl/legend-engine-repl-relational/src/main/java/org/finos/legend/engine/repl/relational/commands/Cache.java +++ b/legend-engine-config/legend-engine-repl/legend-engine-repl-relational/src/main/java/org/finos/legend/engine/repl/relational/commands/Cache.java @@ -23,6 +23,7 @@ import org.finos.legend.engine.plan.execution.result.serialization.TemporaryFile; import org.finos.legend.engine.plan.execution.stores.StoreType; import org.finos.legend.engine.plan.execution.stores.relational.connection.driver.DatabaseManager; +import org.finos.legend.engine.plan.execution.stores.relational.connection.driver.commands.Column; import org.finos.legend.engine.plan.execution.stores.relational.plugin.RelationalStoreState; import org.finos.legend.engine.plan.execution.stores.relational.result.RelationalResult; import org.finos.legend.engine.plan.execution.stores.relational.serialization.RelationalResultToCSVSerializerWithTransformersApplied; @@ -41,6 +42,7 @@ 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; @@ -94,6 +96,7 @@ public boolean process(String line) throws Exception if (res instanceof RelationalResult) { RelationalResult relationalResult = (RelationalResult) res; + List relationalResultColumns = relationalResult.getResultSetColumns(); String tempDir = ((RelationalStoreState) this.client.getPlanExecutor().getExecutorsOfType(StoreType.Relational).getOnly().getStoreState()).getRelationalExecutor().getRelationalExecutionConfiguration().tempPath; try (TemporaryFile tempFile = new TemporaryFile(tempDir)) { @@ -112,7 +115,7 @@ public boolean process(String line) throws Exception String tableName = specifiedTableName != null ? specifiedTableName : "test" + (getTables(connection).size() + 1); try (Statement statement = connection.createStatement()) { - statement.executeUpdate(DatabaseManager.fromString(databaseConnection.type.name()).relationalDatabaseSupport().load(tableName, tempFile.getTemporaryPathForFile())); + statement.executeUpdate(DatabaseManager.fromString(databaseConnection.type.name()).relationalDatabaseSupport().load(tableName, tempFile.getTemporaryPathForFile(), relationalResultColumns)); this.client.println("Cached into table: '" + tableName + "'"); } } diff --git a/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/pom.xml b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/pom.xml index 5fe4e598b93..4f5ebdecd21 100644 --- a/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/pom.xml +++ b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/pom.xml @@ -61,5 +61,33 @@ eclipse-collections + + + + junit + junit + test + + + org.apache.commons + commons-csv + test + + + org.finos.legend.engine + legend-engine-executionPlan-execution + test + + + org.finos.legend.engine + legend-engine-protocol-pure + test + + + org.finos.legend.engine + legend-engine-identity-core + test + + diff --git a/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/vendors/duckdb/DuckDBCommands.java b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/vendors/duckdb/DuckDBCommands.java index 064b3407771..331dde931aa 100644 --- a/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/vendors/duckdb/DuckDBCommands.java +++ b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/vendors/duckdb/DuckDBCommands.java @@ -20,6 +20,7 @@ import org.finos.legend.engine.plan.execution.stores.relational.connection.driver.commands.RelationalDatabaseCommandsVisitor; import java.util.List; +import java.util.stream.Collectors; public class DuckDBCommands extends RelationalDatabaseCommands { @@ -47,6 +48,13 @@ public String load(String tableName, String location) return "CREATE TABLE " + tableName + " AS SELECT * FROM read_csv('" + location + "', header=true);"; } + @Override + public String load(String tableName, String location, List columns) + { + String columnTypesString = columns.stream().map(c -> String.format("'%s': '%s'", c.name, c.type)).collect(Collectors.joining(", ", "{", "}")); + return "CREATE TABLE " + tableName + " AS SELECT * FROM read_csv('" + location + "', header = true, columns = " + columnTypesString + ");"; + } + @Override public String dropTable(String tableName, String location) { diff --git a/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/src/test/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/vendors/duckdb/TestDuckDBCommands.java b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/src/test/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/vendors/duckdb/TestDuckDBCommands.java new file mode 100644 index 00000000000..0264696ffcf --- /dev/null +++ b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/src/test/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/vendors/duckdb/TestDuckDBCommands.java @@ -0,0 +1,215 @@ +// Copyright 2024 Goldman Sachs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package org.finos.legend.engine.plan.execution.stores.relational.connection.driver.vendors.duckdb; + +import org.apache.commons.csv.CSVFormat; +import org.apache.commons.csv.CSVPrinter; +import org.eclipse.collections.impl.factory.Lists; +import org.finos.legend.engine.plan.execution.result.serialization.TemporaryFile; +import org.finos.legend.engine.plan.execution.stores.relational.activity.RelationalExecutionActivity; +import org.finos.legend.engine.plan.execution.stores.relational.config.TemporaryTestDbConfiguration; +import org.finos.legend.engine.plan.execution.stores.relational.connection.driver.commands.Column; +import org.finos.legend.engine.plan.execution.stores.relational.connection.ds.state.ConnectionStateManager; +import org.finos.legend.engine.plan.execution.stores.relational.connection.manager.ConnectionManagerSelector; +import org.finos.legend.engine.plan.execution.stores.relational.result.RelationalResult; +import org.finos.legend.engine.plan.execution.stores.relational.result.SQLExecutionResult; +import org.finos.legend.engine.plan.execution.stores.relational.serialization.RelationalResultToCSVSerializer; +import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.RelationalTdsInstantiationExecutionNode; +import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.nodes.SQLExecutionNode; +import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.result.DataTypeResultType; +import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.DatabaseType; +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.connection.authentication.TestDatabaseAuthenticationStrategy; +import org.finos.legend.engine.protocol.pure.v1.model.packageableElement.store.relational.connection.specification.DuckDBDatasourceSpecification; +import org.finos.legend.engine.shared.core.identity.Identity; +import org.junit.*; +import org.junit.rules.TemporaryFolder; + +import javax.security.auth.Subject; +import java.io.IOException; +import java.sql.*; +import java.util.*; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +public class TestDuckDBCommands +{ + @Rule + public final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder(); + + private static final DuckDBCommands DUCK_DB_COMMANDS = new DuckDBCommands(); + private static final ConnectionManagerSelector CONNECTION_MANAGER_SELECTOR = new ConnectionManagerSelector(new TemporaryTestDbConfiguration(-1), Collections.emptyList(), Optional.empty()); + + @AfterClass + public static void tearDown() throws IOException + { + ConnectionStateManager.getInstance().close(); + } + + @Test + public void testLoadCommand() throws Exception + { + try ( + Connection connection = CONNECTION_MANAGER_SELECTOR.getDatabaseConnection((Subject) null, this.testDuckDBConnection()); + Statement statement = connection.createStatement() + ) + { + String loadSql = DUCK_DB_COMMANDS.load("load_command_test_table", Objects.requireNonNull(TestDuckDBCommands.class.getClassLoader().getResource("personTable.csv")).getFile()); + statement.execute(loadSql); + try (ResultSet rs = statement.executeQuery("select * from load_command_test_table")) + { + assertOnColumnCountAndColumnTypes(rs.getMetaData(), 4, "(id:BIGINT)|(firstName:VARCHAR)|(lastName:VARCHAR)|(age:BIGINT)"); + assertOnResultSetCSV( + rs, + "1,Peter,Smith,23\r\n" + + "2,John,Johnson,22\r\n" + + "3,John,Hill,12\r\n" + + "4,Anthony,Allen,22\r\n" + + "5,Fabrice,Roberts,34\r\n" + + "6,Oliver,Hill,32\r\n" + + "7,David,Harris,35\r\n" + ); + } + } + } + + @Test + public void testLoadCommandWithProvidedTypes() throws Exception + { + try ( + Connection connection = CONNECTION_MANAGER_SELECTOR.getDatabaseConnection((Subject) null, this.testDuckDBConnection()); + Statement statement = connection.createStatement() + ) + { + List columns = Arrays.asList( + new Column("id", "INT"), + new Column("firstName", "VARCHAR"), + new Column("lastName", "VARCHAR(32)"), + new Column("age", "DOUBLE") + ); + String loadSql = DUCK_DB_COMMANDS.load("load_command_with_types_test_table", Objects.requireNonNull(TestDuckDBCommands.class.getClassLoader().getResource("personTable.csv")).getFile(), columns); + statement.execute(loadSql); + try (ResultSet rs = statement.executeQuery("select * from load_command_with_types_test_table")) + { + assertOnColumnCountAndColumnTypes(rs.getMetaData(), 4, "(id:INTEGER)|(firstName:VARCHAR)|(lastName:VARCHAR)|(age:DOUBLE)"); + assertOnResultSetCSV( + rs, + "1,Peter,Smith,23.0\r\n" + + "2,John,Johnson,22.0\r\n" + + "3,John,Hill,12.0\r\n" + + "4,Anthony,Allen,22.0\r\n" + + "5,Fabrice,Roberts,34.0\r\n" + + "6,Oliver,Hill,32.0\r\n" + + "7,David,Harris,35.0\r\n" + ); + } + } + } + + @Test + public void testLoadCommandWithProvidedTypesFromRelationalResult() throws Exception + { + RelationalResult relationalResult; + try (TemporaryFile tempFile = new TemporaryFile(TEMPORARY_FOLDER.getRoot().getAbsolutePath())) + { + RelationalDatabaseConnection duckDbConnection = this.testDuckDBConnection(); + try ( + Connection connection = CONNECTION_MANAGER_SELECTOR.getDatabaseConnection((Subject) null, duckDbConnection); + Statement statement = connection.createStatement() + ) + { + statement.execute(DUCK_DB_COMMANDS.load("load_command_with_types_test_table_1", Objects.requireNonNull(TestDuckDBCommands.class.getClassLoader().getResource("personTable.csv")).getFile())); + SQLExecutionNode sqlExecutionNode = new SQLExecutionNode(); + sqlExecutionNode.isResultColumnsDynamic = true; + sqlExecutionNode.connection = duckDbConnection; + RelationalTdsInstantiationExecutionNode tdsNode = new RelationalTdsInstantiationExecutionNode(); + tdsNode.resultType = new DataTypeResultType(); + SQLExecutionResult sqlExecutionResult = new SQLExecutionResult( + Lists.mutable.of(new RelationalExecutionActivity("select firstName as \"First Name\", avg(age) as \"Average Age\" from load_command_with_types_test_table_1 group by firstName", "")), + sqlExecutionNode, + "DuckDB", + "UTC", + connection, + Identity.getAnonymousIdentity(), + Lists.mutable.empty(), + null + ); + relationalResult = new RelationalResult(sqlExecutionResult, tdsNode); + RelationalResultToCSVSerializer csvSerializer = new RelationalResultToCSVSerializer(relationalResult, true); + tempFile.writeFile(csvSerializer); + relationalResult.close(); + } + + try ( + Connection connection = CONNECTION_MANAGER_SELECTOR.getDatabaseConnection((Subject) null, duckDbConnection); + Statement statement = connection.createStatement() + ) + { + String loadSql = DUCK_DB_COMMANDS.load("load_command_with_types_test_table_2", tempFile.path.toString(), relationalResult.getResultSetColumns()); + statement.execute(loadSql); + try (ResultSet rs = statement.executeQuery("select * from load_command_with_types_test_table_2 order by \"Average Age\"")) + { + assertOnColumnCountAndColumnTypes(rs.getMetaData(), 2, "(First Name:VARCHAR)|(Average Age:DOUBLE)"); + assertOnResultSetCSV( + rs, + "John,17.0\r\n" + + "Anthony,22.0\r\n" + + "Peter,23.0\r\n" + + "Oliver,32.0\r\n" + + "Fabrice,34.0\r\n" + + "David,35.0\r\n" + ); + } + } + } + } + + private static void assertOnColumnCountAndColumnTypes(ResultSetMetaData resultSetMetaData, int expectedColumnCount, String expectedColumnTypes) throws SQLException + { + int columnCount = resultSetMetaData.getColumnCount(); + Assert.assertEquals(expectedColumnCount, columnCount); + + String columnsAndTypes = IntStream.range(1, expectedColumnCount + 1).mapToObj(i -> + { + try + { + return "(" + resultSetMetaData.getColumnLabel(i) + ":" + resultSetMetaData.getColumnTypeName(i) + ")"; + } + catch (SQLException e) + { + throw new RuntimeException(e); + } + }).collect(Collectors.joining("|")); + Assert.assertEquals(expectedColumnTypes, columnsAndTypes); + } + + private static void assertOnResultSetCSV(ResultSet resultSet, String expectedResult) throws SQLException, IOException + { + StringBuilder sb = new StringBuilder(); + try (CSVPrinter csvPrinter = new CSVPrinter(sb, CSVFormat.DEFAULT)) + { + csvPrinter.printRecords(resultSet); + } + Assert.assertEquals(expectedResult, sb.toString()); + } + + private RelationalDatabaseConnection testDuckDBConnection() throws Exception + { + DuckDBDatasourceSpecification duckDBDatasourceSpecification = new DuckDBDatasourceSpecification(); + duckDBDatasourceSpecification.path = TEMPORARY_FOLDER.newFolder("duckDB").getAbsolutePath() + "/duck_db_file"; + TestDatabaseAuthenticationStrategy testDatabaseAuthSpec = new TestDatabaseAuthenticationStrategy(); + return new RelationalDatabaseConnection(duckDBDatasourceSpecification, testDatabaseAuthSpec, DatabaseType.DuckDB); + } +} diff --git a/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/src/test/resources/personTable.csv b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/src/test/resources/personTable.csv new file mode 100644 index 00000000000..9cd9912e93d --- /dev/null +++ b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-dbExtension/legend-engine-xt-relationalStore-duckdb/legend-engine-xt-relationalStore-duckdb-execution/src/test/resources/personTable.csv @@ -0,0 +1,8 @@ +id,firstName,lastName,age +1,Peter,Smith,23 +2,John,Johnson,22 +3,John,Hill,12 +4,Anthony,Allen,22 +5,Fabrice,Roberts,34 +6,Oliver,Hill,32 +7,David,Harris,35 \ No newline at end of file diff --git a/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-execution/legend-engine-xt-relationalStore-executionPlan-connection/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/commands/RelationalDatabaseCommands.java b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-execution/legend-engine-xt-relationalStore-executionPlan-connection/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/commands/RelationalDatabaseCommands.java index 4bdc40f8aae..4488a4fd564 100644 --- a/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-execution/legend-engine-xt-relationalStore-executionPlan-connection/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/commands/RelationalDatabaseCommands.java +++ b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-execution/legend-engine-xt-relationalStore-executionPlan-connection/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/commands/RelationalDatabaseCommands.java @@ -34,6 +34,11 @@ public String load(String tableName, String location) throw new RuntimeException("Load not implemented for " + this.getClass().getSimpleName()); } + public String load(String tableName, String location, List columns) + { + throw new RuntimeException("Load not implemented for " + this.getClass().getSimpleName()); + } + public String dropTable(String tableName, String location) { throw new RuntimeException("Drop table not implemented for " + this.getClass().getSimpleName()); diff --git a/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-execution/legend-engine-xt-relationalStore-executionPlan-connection/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/vendors/h2/H2Commands.java b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-execution/legend-engine-xt-relationalStore-executionPlan-connection/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/vendors/h2/H2Commands.java index a9720f895fe..a5067e06e76 100644 --- a/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-execution/legend-engine-xt-relationalStore-executionPlan-connection/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/vendors/h2/H2Commands.java +++ b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-execution/legend-engine-xt-relationalStore-executionPlan-connection/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/connection/driver/vendors/h2/H2Commands.java @@ -51,6 +51,12 @@ public String load(String tableName, String location) return "CREATE TABLE " + tableName + " AS SELECT * FROM CSVREAD('" + location + "');"; } + @Override + public String load(String tableName, String location, List columns) + { + return "CREATE TABLE " + tableName + "(" + columns.stream().map(c -> c.name + " " + c.type).collect(Collectors.joining(", ")) + ") AS SELECT * FROM CSVREAD('" + location + "');"; + } + @Override public T accept(RelationalDatabaseCommandsVisitor visitor) { diff --git a/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-execution/legend-engine-xt-relationalStore-executionPlan/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/result/RelationalResult.java b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-execution/legend-engine-xt-relationalStore-executionPlan/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/result/RelationalResult.java index 4ae60f8d2ca..ff1394aa737 100644 --- a/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-execution/legend-engine-xt-relationalStore-executionPlan/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/result/RelationalResult.java +++ b/legend-engine-xts-relationalStore/legend-engine-xt-relationalStore-execution/legend-engine-xt-relationalStore-executionPlan/src/main/java/org/finos/legend/engine/plan/execution/stores/relational/result/RelationalResult.java @@ -19,12 +19,14 @@ import io.opentracing.Span; import java.sql.Connection; +import java.sql.JDBCType; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Statement; import java.sql.Timestamp; import java.sql.Types; +import java.util.ArrayList; import java.util.Calendar; import java.util.GregorianCalendar; import java.util.List; @@ -64,6 +66,7 @@ import org.finos.legend.engine.plan.execution.stores.StoreExecutableManager; import org.finos.legend.engine.plan.execution.stores.relational.activity.RelationalExecutionActivity; import org.finos.legend.engine.plan.execution.stores.relational.connection.driver.DatabaseManager; +import org.finos.legend.engine.plan.execution.stores.relational.connection.driver.commands.Column; import org.finos.legend.engine.plan.execution.stores.relational.result.builder.relation.RelationBuilder; import org.finos.legend.engine.plan.execution.stores.relational.serialization.RelationalResultToCSVSerializer; import org.finos.legend.engine.plan.execution.stores.relational.serialization.RelationalResultToCSVSerializerWithTransformersApplied; @@ -694,4 +697,26 @@ public void cancel() LOGGER.error(new LogInfo(Identity.getAnonymousIdentity().getName(), LoggingEventType.EXECUTABLE_CANCELLATION_ERROR, "Unable to cancel RelationalResult for session " + RequestContext.getSessionID(this.requestContext) + " " + e.getMessage()).toString()); } } + + public List getResultSetColumns() + { + try + { + List columns = new ArrayList<>(this.resultSetMetaData.getColumnCount()); + Function unquote = s -> s.startsWith("\"") && s.endsWith("\"") ? s.substring(1, s.length() - 1) : s; + for (int i = 1; i <= this.resultSetMetaData.getColumnCount(); i++) + { + String columnType = JDBCType.valueOf(this.resultSetMetaData.getColumnType(i)).getName(); + String updatedColumnType = columnType.equals("TIMESTAMP_WITH_TIMEZONE") ? "TIMESTAMP WITH TIME ZONE" : + columnType.equals("TIME_WITH_TIMEZONE") ? "TIME WITH TIME ZONE" : columnType; + columns.add(new Column(unquote.valueOf(this.resultSetMetaData.getColumnLabel(i)), updatedColumnType)); + } + return columns; + } + catch (SQLException e) + { + this.close(); + throw new RuntimeException(e); + } + } } \ No newline at end of file