datacube: update API (#3073)

finos · Sep 10, 2024 · 027c1ae · 027c1ae
1 parent 8768edc
commit 027c1ae
Show file tree

Hide file tree

Showing 9 changed files with 135 additions and 19 deletions.
diff --git a/legend-engine-config/legend-engine-repl/legend-engine-repl-app-assembly/pom.xml b/legend-engine-config/legend-engine-repl/legend-engine-repl-app-assembly/pom.xml
@@ -25,6 +25,7 @@
     </parent>
 
     <artifactId>legend-engine-repl-app-assembly</artifactId>
+    <name>Legend Engine - REPL - App Assembly</name>
     <packaging>pom</packaging>
 
     <build>

diff --git a/...engine-repl-client/src/main/java/org/finos/legend/engine/repl/autocomplete/Completer.java b/...engine-repl-client/src/main/java/org/finos/legend/engine/repl/autocomplete/Completer.java
@@ -411,7 +411,7 @@ private static boolean checkIfCurrent(SourceInformation sourceInformation, int _
 //                _line <= sourceInformation.endLine &&
 //                (sourceInformation.startColumn - 1) <= _column &&
 //                _column <= sourceInformation.endColumn);
-        return sourceInformation.startLine <= _line &&
+        return sourceInformation != null && sourceInformation.startLine <= _line &&
                 _line <= sourceInformation.endLine &&
                 (sourceInformation.startColumn - 1) <= _column &&
                 _column <= sourceInformation.endColumn;

diff --git a/...data-cube/src/main/java/org/finos/legend/engine/repl/dataCube/server/DataCubeHelpers.java b/...data-cube/src/main/java/org/finos/legend/engine/repl/dataCube/server/DataCubeHelpers.java
@@ -139,7 +139,7 @@ public static String getQueryCode(ValueSpecification valueSpecification, Boolean
         return valueSpecification.accept(DEPRECATED_PureGrammarComposerCore.Builder.newInstance().withRenderStyle(pretty != null && pretty ? RenderStyle.PRETTY : RenderStyle.STANDARD).build());
     }
 
-    public static CompletionResult getCodeTypeahead(String code, Boolean isPartial, PureModelContextData data, MutableList<CompleterExtension> extensions, LegendInterface legendInterface)
+    public static CompletionResult getCodeTypeahead(String code, String baseQueryCode, PureModelContextData data, MutableList<CompleterExtension> extensions, LegendInterface legendInterface)
     {
         try
         {
@@ -149,14 +149,9 @@ public static CompletionResult getCodeTypeahead(String code, Boolean isPartial,
                     .withElements(ListIterate.select(data.getElements(), el -> !el.getPath().equals(REPL_RUN_FUNCTION_QUALIFIED_PATH)))
                     .build();
             String graphCode = PureGrammarComposer.newInstance(PureGrammarComposerContext.Builder.newInstance().build()).renderPureModelContextData(newData);
-            String queryCode = code;
-            if (isPartial != null && isPartial)
-            {
-                Function func = (Function) ListIterate.select(data.getElements(), el -> el.getPath().equals(REPL_RUN_FUNCTION_QUALIFIED_PATH)).getFirst();
-                String existingCode = func.body.get(0).accept(DEPRECATED_PureGrammarComposerCore.Builder.newInstance().build());
-                queryCode = existingCode + code;
-            }
-            CompletionResult result = new Completer(graphCode, extensions, legendInterface).complete(queryCode);
+            String queryCode = (baseQueryCode != null ? baseQueryCode : "") + code;
+            Completer completer = new Completer(graphCode, extensions, legendInterface);
+            CompletionResult result = completer.complete(queryCode);
             if (result.getEngineException() != null)
             {
                 return new CompletionResult(Lists.mutable.empty());

diff --git a/...repl-data-cube/src/main/java/org/finos/legend/engine/repl/dataCube/server/REPLServer.java b/...repl-data-cube/src/main/java/org/finos/legend/engine/repl/dataCube/server/REPLServer.java
@@ -86,6 +86,7 @@ public void initialize() throws Exception
                 .withKeyValue("/api/dataCube/getQueryCode/batch", new DataCubeQueryBuilder.GetQueryCodeBatch())
                 .withKeyValue("/api/dataCube/getBaseQuery", new DataCubeQueryBuilder.GetBaseQuery())
                 .withKeyValue("/api/dataCube/getRelationReturnType", new DataCubeQueryBuilder.GetRelationReturnType())
+                .withKeyValue("/api/dataCube/getRelationReturnType/code", new DataCubeQueryBuilder.GetQueryCodeRelationReturnType())
                 .withKeyValue("/api/dataCube/executeQuery", new DataCubeQueryExecutor.ExecuteQuery())
                 .keyValuesView().collect(config -> server.createContext(config.getOne(), config.getTwo().getHandler(this.state))).toList();
 

diff --git a/.../main/java/org/finos/legend/engine/repl/dataCube/server/handler/DataCubeQueryBuilder.java b/.../main/java/org/finos/legend/engine/repl/dataCube/server/handler/DataCubeQueryBuilder.java
@@ -15,8 +15,15 @@
 package org.finos.legend.engine.repl.dataCube.server.handler;
 
 import com.sun.net.httpserver.HttpHandler;
+import org.apache.commons.lang3.StringUtils;
 import org.eclipse.collections.impl.map.mutable.MapAdapter;
+import org.eclipse.collections.impl.utility.ListIterate;
+import org.finos.legend.engine.language.pure.grammar.from.PureGrammarParser;
+import org.finos.legend.engine.language.pure.grammar.to.PureGrammarComposer;
+import org.finos.legend.engine.language.pure.grammar.to.PureGrammarComposerContext;
+import org.finos.legend.engine.protocol.pure.v1.model.SourceInformation;
 import org.finos.legend.engine.protocol.pure.v1.model.context.PureModelContextData;
+import org.finos.legend.engine.protocol.pure.v1.model.relationType.RelationType;
 import org.finos.legend.engine.protocol.pure.v1.model.valueSpecification.ValueSpecification;
 import org.finos.legend.engine.protocol.pure.v1.model.valueSpecification.raw.Lambda;
 import org.finos.legend.engine.repl.autocomplete.CompletionResult;
@@ -30,6 +37,8 @@
 import java.util.stream.Collectors;
 
 import static org.finos.legend.engine.repl.dataCube.server.REPLServerHelpers.*;
+import static org.finos.legend.engine.repl.shared.ExecutionHelper.REPL_RUN_FUNCTION_QUALIFIED_PATH;
+import static org.finos.legend.engine.repl.shared.ExecutionHelper.REPL_RUN_FUNCTION_SIGNATURE;
 
 public class DataCubeQueryBuilder
 {
@@ -53,7 +62,7 @@ public HttpHandler getHandler(REPLServerState state)
                     }
                     catch (Exception e)
                     {
-                        handleResponse(exchange, 400, e instanceof EngineException ? state.objectMapper.writeValueAsString(e) : e.getMessage(), state);
+                        handleResponse(exchange, 400, e instanceof EngineException ? state.objectMapper.writeValueAsString(new DataCubeQueryBuilderError((EngineException) e)) : e.getMessage(), state);
                     }
                 }
             };
@@ -140,7 +149,7 @@ public HttpHandler getHandler(REPLServerState state)
                         String requestBody = bufferReader.lines().collect(Collectors.joining());
                         DataCubeQueryTypeaheadInput input = state.objectMapper.readValue(requestBody, DataCubeQueryTypeaheadInput.class);
                         PureModelContextData data = state.getCurrentPureModelContextData();
-                        CompletionResult result = DataCubeHelpers.getCodeTypeahead(input.code, input.isPartial, data, state.client.getCompleterExtensions(), state.legendInterface);
+                        CompletionResult result = DataCubeHelpers.getCodeTypeahead(input.code, DataCubeHelpers.getQueryCode(input.baseQuery, false), data, state.client.getCompleterExtensions(), state.legendInterface);
                         handleResponse(exchange, 200, state.objectMapper.writeValueAsString(result.getCompletion()), state);
                     }
                     catch (Exception e)
@@ -173,7 +182,60 @@ public HttpHandler getHandler(REPLServerState state)
                     }
                     catch (Exception e)
                     {
-                        handleResponse(exchange, 500, e instanceof EngineException ? state.objectMapper.writeValueAsString(e) : e.getMessage(), state);
+                        handleResponse(exchange, 500, e instanceof EngineException ? state.objectMapper.writeValueAsString(new DataCubeQueryBuilderError((EngineException) e)) : e.getMessage(), state);
+                    }
+                }
+            };
+        }
+    }
+
+    public static class GetQueryCodeRelationReturnType implements DataCubeServerHandler
+    {
+        @Override
+        public HttpHandler getHandler(REPLServerState state)
+        {
+            return exchange ->
+            {
+                if ("POST".equals(exchange.getRequestMethod()))
+                {
+                    try
+                    {
+                        InputStreamReader inputStreamReader = new InputStreamReader(exchange.getRequestBody(), StandardCharsets.UTF_8);
+                        BufferedReader bufferReader = new BufferedReader(inputStreamReader);
+                        String requestBody = bufferReader.lines().collect(Collectors.joining());
+                        DataCubeGetQueryCodeRelationReturnTypeInput input = state.objectMapper.readValue(requestBody, DataCubeGetQueryCodeRelationReturnTypeInput.class);
+                        PureModelContextData currentData = state.getCurrentPureModelContextData();
+                        PureModelContextData newData = PureModelContextData.newBuilder()
+                                .withOrigin(currentData.getOrigin())
+                                .withSerializer(currentData.getSerializer())
+                                .withElements(ListIterate.select(currentData.getElements(), el -> !el.getPath().equals(REPL_RUN_FUNCTION_QUALIFIED_PATH)))
+                                .build();
+                        String graphCode = PureGrammarComposer.newInstance(PureGrammarComposerContext.Builder.newInstance().build()).renderPureModelContextData(newData);
+                        graphCode += "\n###Pure\n" +
+                                "import meta::pure::functions::relation::*;\n" +
+                                "function " + REPL_RUN_FUNCTION_SIGNATURE + "{\n";
+                        graphCode += DataCubeHelpers.getQueryCode(input.baseQuery, false) + "\n";
+                        int lineOffset = StringUtils.countMatches(graphCode, "\n");
+                        graphCode += input.code;
+                        graphCode += "\n}";
+
+                        try
+                        {
+                            PureModelContextData data = PureGrammarParser.newInstance().parseModel(graphCode);
+                            RelationType relationType = DataCubeHelpers.getRelationReturnType(state.legendInterface, data);
+                            handleResponse(exchange, 200, state.objectMapper.writeValueAsString(relationType), state);
+                        }
+                        catch (EngineException e)
+                        {
+                            SourceInformation sourceInformation = e.getSourceInformation();
+                            sourceInformation.startLine -= lineOffset;
+                            sourceInformation.endLine -= lineOffset;
+                            handleResponse(exchange, 400, state.objectMapper.writeValueAsString(new DataCubeQueryBuilderError(new EngineException(e.getMessage(), sourceInformation, e.getErrorType()))), state);
+                        }
+                    }
+                    catch (Exception e)
+                    {
+                        handleResponse(exchange, 500, e instanceof EngineException ? state.objectMapper.writeValueAsString(new DataCubeQueryBuilderError((EngineException) e)) : e.getMessage(), state);
                     }
                 }
             };

diff --git a/...legend/engine/repl/dataCube/server/model/DataCubeGetQueryCodeRelationReturnTypeInput.java b/...legend/engine/repl/dataCube/server/model/DataCubeGetQueryCodeRelationReturnTypeInput.java
@@ -0,0 +1,23 @@
+// 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.server.model;
+
+import org.finos.legend.engine.protocol.pure.v1.model.valueSpecification.ValueSpecification;
+
+public class DataCubeGetQueryCodeRelationReturnTypeInput
+{
+    public String code;
+    public ValueSpecification baseQuery;
+}
diff --git a/...in/java/org/finos/legend/engine/repl/dataCube/server/model/DataCubeQueryBuilderError.java b/...in/java/org/finos/legend/engine/repl/dataCube/server/model/DataCubeQueryBuilderError.java
@@ -0,0 +1,32 @@
+// 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.server.model;
+
+import org.finos.legend.engine.protocol.pure.v1.model.SourceInformation;
+import org.finos.legend.engine.shared.core.operational.errorManagement.EngineException;
+
+public class DataCubeQueryBuilderError
+{
+    public String type;
+    public String message;
+    public SourceInformation sourceInformation;
+
+    public DataCubeQueryBuilderError(EngineException engineException)
+    {
+        this.type = engineException.getErrorType() != null ? engineException.getErrorType().name() : null;
+        this.message = engineException.getMessage();
+        this.sourceInformation = engineException.getSourceInformation();
+    }
+}
diff --git a/.../java/org/finos/legend/engine/repl/dataCube/server/model/DataCubeQueryTypeaheadInput.java b/.../java/org/finos/legend/engine/repl/dataCube/server/model/DataCubeQueryTypeaheadInput.java
@@ -14,8 +14,10 @@
 
 package org.finos.legend.engine.repl.dataCube.server.model;
 
+import org.finos.legend.engine.protocol.pure.v1.model.valueSpecification.ValueSpecification;
+
 public class DataCubeQueryTypeaheadInput
 {
     public String code;
-    public Boolean isPartial;
+    public ValueSpecification baseQuery;
 }
diff --git a/...pl-data-cube/src/test/java/org/finos/legend/engine/repl/dataCube/TestDataCubeHelpers.java b/...pl-data-cube/src/test/java/org/finos/legend/engine/repl/dataCube/TestDataCubeHelpers.java
@@ -197,30 +197,30 @@ public void testTypeaheadPartial()
     {
         String code = "->extend(~[newCol:c|'ok', colX: c|$c.";
         String expectedResult = "{\"completion\":[{\"completion\":\"FIRSTNAME\",\"display\":\"FIRSTNAME\"}]}";
-        testTypeahead(expectedResult, code, true);
+        testTypeahead(expectedResult, code, "#>{test::TestDatabase.TEST0}#->filter(c | $c.FIRSTNAME != 'Doe')->select(~FIRSTNAME)->from(test::test)");
     }
 
     @Test
     public void testTypeaheadFull()
     {
         String code = "#>{test::TestDatabase.TEST0}#->extend(~[newCol:c|'ok', colX: c|$c.";
         String expectedResult = "{\"completion\":[{\"completion\":\"FIRSTNAME\",\"display\":\"FIRSTNAME\"},{\"completion\":\"LASTNAME\",\"display\":\"LASTNAME\"}]}";
-        testTypeahead(expectedResult, code, false);
+        testTypeahead(expectedResult, code, null);
     }
 
     @Test
     public void testTypeaheadFullWithError()
     {
         String code = "#>{test::TestDatabase.TEST0}#-->extend(~[newCol:c|'ok', colX: c|$c.";
         String expectedResult = "{\"completion\":[]}";
-        testTypeahead(expectedResult, code, false);
+        testTypeahead(expectedResult, code, null);
     }
 
-    private void testTypeahead(String expectedResult, String code, boolean isPartial)
+    private void testTypeahead(String expectedResult, String code, String baseQueryCode)
     {
         try
         {
-            Assert.assertEquals(expectedResult, objectMapper.writeValueAsString(DataCubeHelpers.getCodeTypeahead(code, isPartial, pureModelContextData, completerExtensions, legendInterface)));
+            Assert.assertEquals(expectedResult, objectMapper.writeValueAsString(DataCubeHelpers.getCodeTypeahead(code, baseQueryCode, pureModelContextData, completerExtensions, legendInterface)));
         }
         catch (IOException e)
         {