Skip to content

Commit

Permalink
datacube: add sample data and improve testing (#3144)
Browse files Browse the repository at this point in the history
  • Loading branch information
akphi authored Oct 1, 2024
1 parent 278d76b commit e111438
Show file tree
Hide file tree
Showing 14 changed files with 174 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@ public Client(MutableList<ReplExtension> replExtensions, MutableList<CompleterEx
this.documentation = DocumentationGeneration.buildDocumentation();
this.initialize();
replExtensions.forEach(e -> e.initialize(this));
replExtensions.forEach(e -> e.postInitialize(this));

this.printDebug("[DEV] Legend REPL v" + DeploymentStateAndVersions.sdlc.buildVersion + " (" + DeploymentStateAndVersions.sdlc.commitIdAbbreviated + ")");
if (System.getProperty("legend.repl.initializationMessage") != null)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -48,5 +48,15 @@ default MutableList<String> typeGroup()

MutableList<String> generateDynamicContent(String code);

void initialize(Client client);
default void initialize(Client client)
{
}

// This method is called after all extensions have been initialized
// This is useful for cases where we need to invoke initialization-type tasks from one extension
// that might depend on another extension. This is for now the preferred approach over specifying
// a dependency graph to determine the order of initialization.
default void postInitialize(Client client)
{
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -82,8 +82,6 @@ public static ExecutionHelper.ExecuteResultSummary executeCode(String txt, Clien
PureModel pureModel = client.getLegendInterface().compile(pmcd);

// Plan
// TODO: Since H2 does not support pivot(), when pivot() is used, the debugger will fail as it defaults to use H2
// when we switch out to use DuckDB as the core testing DB, then this issue would be resolved
Root_meta_pure_executionPlan_ExecutionPlan plan = client.getLegendInterface().generatePlan(pureModel, client.isDebug());
RichIterable<? extends Root_meta_pure_extension_Extension> extensions = PureCoreExtensionLoader.extensions().flatCollect(e -> e.extraPureCoreExtensions(pureModel.getExecutionSupport()));
String planStr = PlanGenerator.serializeToJSON(plan, "vX_X_X", pureModel, extensions, LegendPlanTransformers.transformers);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,26 +14,31 @@

package org.finos.legend.engine.repl.dataCube;

import org.eclipse.collections.api.factory.Maps;
import org.eclipse.collections.api.list.MutableList;
import org.eclipse.collections.api.map.MutableMap;
import org.eclipse.collections.impl.factory.Lists;
import org.finos.legend.engine.plan.execution.result.Result;
import org.finos.legend.engine.repl.client.Client;
import org.finos.legend.engine.repl.core.Command;
import org.finos.legend.engine.repl.core.ReplExtension;
import org.finos.legend.engine.repl.dataCube.commands.*;
import org.finos.legend.engine.repl.dataCube.server.REPLServer;
import org.finos.legend.engine.repl.dataCube.shared.DataCubeSampleData;

public class DataCubeReplExtension implements ReplExtension
{
private Client client;
public REPLServer replServer;
public MutableMap<String, DataCubeSampleData> samples = Maps.mutable.empty();

@Override
public String type()
{
return "relational";
}

@Override
public void initialize(Client client)
{
this.client = client;
Expand All @@ -49,6 +54,22 @@ public void initialize(Client client)
}
}

@Override
public void postInitialize(Client client)
{
this.client = client;

try
{
Lists.mutable.with(DataCubeSampleData.SPORT, DataCubeSampleData.TREE).forEach(sample -> this.samples.put(sample.name, sample));
this.samples.forEach(sample -> sample.load(this.client));
}
catch (Exception e)
{
throw new RuntimeException(e);
}
}

@Override
public MutableList<Command> getExtraCommands()
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
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;
Expand Down Expand Up @@ -161,10 +162,14 @@ public void beforeStep()
Statement statement = connection.createStatement())
{
MutableList<Table> tables = getTables(connection);
this.tableName = tableName == null ? "table" + (tables.size() + 1) : tableName;
if (tables.anySatisfy(table -> table.name.equals(DataCubeSampleData.SPORT.tableName)))
{
this.tableName = DataCubeSampleData.SPORT.tableName;
}
this.tableName = this.tableName == null ? "table" + (tables.size() + 1) : this.tableName;
// automatically create new table for walkthrough if it doesn't exist or somehow dropped between walkthrough steps
// NOTE: if the table has been replaced or changed somehow, the walkthrough will fail (we should consider resetting the table)
if (!tables.anySatisfy(t -> t.name.equals(this.tableName)))
if (!tables.anySatisfy(table -> table.name.equals(this.tableName)))
{
Path tempFile = Files.createTempFile("walkthrough-sample-data", ".csv");
FileOutputStream fos = new FileOutputStream(tempFile.toFile());
Expand Down Expand Up @@ -326,6 +331,11 @@ private RelationalReplExtension getRelationalExtension()
return this.client.getReplExtensions().selectInstancesOf(RelationalReplExtension.class).getFirst();
}

public String getTableName()
{
return this.tableName;
}

private String query(MutableList<String> parts, Integer indexToHighlight)
{
MutableList<String> allParts = Lists.mutable
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,9 @@ public static DataCubeExecutionResult executeQuery(Client client, LegendInterfac
{
client.printDebug("---------------------------------------- PLAN ----------------------------------------");
}
Root_meta_pure_executionPlan_ExecutionPlan _plan = legendInterface.generatePlan(pureModel, debug);
// TODO: Since H2 does not support pivot(), when pivot() is used, the debugger will fail as it defaults to use H2
// when we switch out to use DuckDB as the core testing DB, then this issue should be resolved
Root_meta_pure_executionPlan_ExecutionPlan _plan = legendInterface.generatePlan(pureModel, false);
String planStr = PlanGenerator.serializeToJSON(_plan, "vX_X_X", pureModel, extensions, LegendPlanTransformers.transformers);
if (client != null && debug)
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import org.apache.commons.io.IOUtils;
import org.finos.legend.engine.repl.dataCube.server.REPLServer;
import org.finos.legend.engine.repl.dataCube.server.model.DataCubeInfrastructureInfo;
import org.finos.legend.engine.repl.dataCube.shared.DataCubeSampleData;

import java.io.InputStream;
import java.io.OutputStream;
Expand All @@ -39,6 +40,8 @@ public HttpHandler getHandler(REPLServerState state)
{
DataCubeInfrastructureInfo info = new DataCubeInfrastructureInfo();
info.gridClientLicense = System.getProperty("legend.repl.dataCube.gridLicenseKey") == null ? "" : System.getProperty("legend.repl.dataCube.gridLicenseKey");
info.simpleSampleDataTableName = DataCubeSampleData.TREE.tableName;
info.complexSampleDataTableName = DataCubeSampleData.SPORT.tableName;
handleResponse(exchange, 200, state.objectMapper.writeValueAsString(info), state);
}
catch (Exception e)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,4 +17,6 @@
public class DataCubeInfrastructureInfo
{
public String gridClientLicense;
public String simpleSampleDataTableName;
public String complexSampleDataTableName;
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
// 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.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;
import java.nio.file.Path;
import java.sql.Connection;
import java.sql.Statement;
import java.util.Arrays;
import java.util.Objects;

import static org.finos.legend.engine.repl.relational.schema.MetadataReader.getTables;

public class DataCubeSampleData
{
public static final DataCubeSampleData SPORT = new DataCubeSampleData("sport", "sample__sport", "org/finos/legend/engine/repl/dataCube/walkthrough/sport-data.csv", Lists.mutable.with("Athlete", "Age", "Country", "Year", "Date", "Sport", "Gold", "Silver", "Bronze"));
public static final DataCubeSampleData TREE = new DataCubeSampleData("tree", "sample__tree", "org/finos/legend/engine/repl/dataCube/walkthrough/tree-data.csv", Lists.mutable.with("city", "country", "year", "tree"));

public final String name;
public final String tableName;
public final String csvFilePath;
public final MutableList<String> expectedColumns;

public DataCubeSampleData(String name, String tableName, String csvFilePath, MutableList<String> expectedColumns)
{
this.name = name;
this.tableName = tableName;
this.csvFilePath = csvFilePath;
this.expectedColumns = expectedColumns;
}

public void load(Client client)
{
DatabaseConnection databaseConnection = ConnectionHelper.getDatabaseConnection(client.getModelState().parse(), DataCube.getLocalConnectionPath());
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));
}
Path tempFile = Files.createTempFile("sample-data" + this.name, ".csv");
FileOutputStream fos = new FileOutputStream(tempFile.toFile());
IOUtils.copy(Objects.requireNonNull(inputStream, "Can't extract sample data '" + this.name + "' from " + this.csvFilePath), fos);
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);
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
city,country,year,tree
NYC,USA,2011,50
HAN,VNM,2011,12
NYC,USA,2000,-51
PEK,CHN,2012,-10
SHA,CHN,2012,-16
SAN,USA,2000,-20
HAN,VNM,2010,-5
SAN,USA,2011,12
SGP,SGP,2000,30
LDN,UK,2011,30
SAN,USA,2011,-25
NYC,USA,2000,100
NYC,USA,2012,-76
SAN,USA,2011,-17
PEK,CHN,2011,77
SHA,CHN,2000,-11
SGP,SGP,2000,20
HAN,VNM,2012,30
HAN,VNM,2010,19
LDN,UK,2012,27
NYC,USA,2012,54
SHA,CHN,2012,6
PEK,CHN,2000,60
SGP,SGP,2011,50
LDN,UK,2000,16
SGP,SGP,2012,-10
HAN,VNM,2012,-11
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,10 @@
import org.finos.legend.engine.protocol.pure.v1.model.executionPlan.SingleExecutionPlan;
import org.finos.legend.engine.repl.client.Client;
import org.finos.legend.engine.repl.dataCube.commands.DataCubeWalkthrough;
import org.finos.legend.engine.repl.dataCube.shared.DataCubeSampleData;
import org.finos.legend.engine.repl.relational.RelationalReplExtension;
import org.finos.legend.engine.repl.shared.ExecutionHelper;
import org.junit.Assert;
import org.junit.Test;

import java.nio.file.Files;
Expand Down Expand Up @@ -51,6 +53,7 @@ public TestDataCubeWalkthrough() throws Exception
@Test
public void testAllCommandsExecution()
{
Assert.assertEquals(DataCubeSampleData.SPORT.tableName, this.walkthrough.getTableName());
Lists.mutable.with(
walkthrough.query(DataCubeWalkthrough.DataCubeWalkthrough1.SELECT_ALL),
walkthrough.query(DataCubeWalkthrough.DataCubeWalkthrough1.FILTER),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ public boolean process(String line) throws Exception
String tableName = tokens[2];
try (Statement statement = connection.createStatement())
{
statement.executeUpdate(DatabaseManager.fromString(databaseConnection.type.name()).relationalDatabaseSupport().dropTable(tableName, tokens[1]));
statement.executeUpdate(DatabaseManager.fromString(databaseConnection.type.name()).relationalDatabaseSupport().dropTable(tableName));
this.client.println("Dropped table: '" + tableName + "'");
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ public String load(String tableName, String location, List<Column> columns)
}

@Override
public String dropTable(String tableName, String location)
public String dropTable(String tableName)
{
return "DROP TABLE " + tableName + ";";
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ public String load(String tableName, String location, List<Column> columns)
throw new RuntimeException("Load not implemented for " + this.getClass().getSimpleName());
}

public String dropTable(String tableName, String location)
public String dropTable(String tableName)
{
throw new RuntimeException("Drop table not implemented for " + this.getClass().getSimpleName());
}
Expand Down

0 comments on commit e111438

Please sign in to comment.