Minor fixes and improvements for LSP/Notebooks (finos#3183)

goldmansachs · Oct 16, 2024 · 982545d · 982545d
1 parent 7fe40bf
commit 982545d
Show file tree

Hide file tree

Showing 8 changed files with 362 additions and 160 deletions.
diff --git a/...ava/org/finos/legend/engine/repl/dataCube/commands/DataCube__DEV__runDuckDBSelectSQL.java b/...ava/org/finos/legend/engine/repl/dataCube/commands/DataCube__DEV__runDuckDBSelectSQL.java
@@ -29,7 +29,7 @@
 import java.sql.Statement;
 import java.util.List;
 
-import static org.finos.legend.engine.repl.relational.shared.ResultHelper.prettyGridPrint;
+import static org.finos.legend.engine.plan.execution.stores.relational.result.RelationalResultGridPrintUtility.prettyGridPrint;
 
 public class DataCube__DEV__runDuckDBSelectSQL implements Command
 {

diff --git a/...tional/src/main/java/org/finos/legend/engine/repl/relational/RelationalReplExtension.java b/...tional/src/main/java/org/finos/legend/engine/repl/relational/RelationalReplExtension.java
@@ -16,9 +16,9 @@
 
 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.plan.execution.result.Result;
 import org.finos.legend.engine.plan.execution.stores.relational.AlloyH2Server;
+import org.finos.legend.engine.plan.execution.stores.relational.result.RealizedRelationalResult;
 import org.finos.legend.engine.plan.execution.stores.relational.result.RelationalResult;
 import org.finos.legend.engine.repl.client.Client;
 import org.finos.legend.engine.repl.core.Command;
@@ -30,10 +30,9 @@
 import org.finos.legend.engine.repl.relational.local.LocalConnectionManagement;
 import org.finos.legend.engine.repl.relational.local.LocalConnectionType;
 
-import java.sql.ResultSet;
 import java.sql.SQLException;
 
-import static org.finos.legend.engine.repl.relational.shared.ResultHelper.prettyGridPrint;
+import static org.finos.legend.engine.plan.execution.stores.relational.result.RelationalResultGridPrintUtility.prettyGridPrint;
 
 public class RelationalReplExtension implements ReplExtension
 {
@@ -91,20 +90,21 @@ public MutableList<Command> getExtraCommands()
     @Override
     public boolean supports(Result res)
     {
-        return res instanceof RelationalResult;
+        return res instanceof RelationalResult || res instanceof RealizedRelationalResult;
     }
 
     @Override
     public String print(Result res)
     {
-        RelationalResult relationalResult = (RelationalResult) res;
-        try (ResultSet rs = relationalResult.resultSet)
+        if (res instanceof RelationalResult)
         {
-            return prettyGridPrint(rs, relationalResult.sqlColumns, ListIterate.collect(relationalResult.getSQLResultColumns(), col -> col.dataType), maxRowSize, 60);
+            RelationalResult relationalResult = (RelationalResult) res;
+            return prettyGridPrint(relationalResult, maxRowSize, 60);
         }
-        catch (Exception e)
+        else
         {
-            throw new RuntimeException(e);
+            RealizedRelationalResult relationalResult = (RealizedRelationalResult) res;
+            return prettyGridPrint(relationalResult, maxRowSize, 60);
         }
     }
 

diff --git a/...relational/src/main/java/org/finos/legend/engine/repl/relational/shared/ResultHelper.java b/...relational/src/main/java/org/finos/legend/engine/repl/relational/shared/ResultHelper.java
diff --git a/...tionPlan-execution/src/main/java/org/finos/legend/engine/plan/execution/PlanExecutor.java b/...tionPlan-execution/src/main/java/org/finos/legend/engine/plan/execution/PlanExecutor.java
@@ -368,7 +368,7 @@ public ExecutionState buildDefaultExecutionState(SingleExecutionPlan executionPl
         return buildDefaultExecutionState(executionPlan, vars, null);
     }
 
-    private ExecutionState buildDefaultExecutionState(SingleExecutionPlan executionPlan, Map<String, Result> vars, PlanExecutionContext planExecutionContext)
+    public ExecutionState buildDefaultExecutionState(SingleExecutionPlan executionPlan, Map<String, Result> vars, PlanExecutionContext planExecutionContext)
     {
         ExecutionState executionState = new ExecutionState(vars, executionPlan.templateFunctions, this.extraExecutors.collect(StoreExecutor::buildStoreExecutionState), this.isJavaCompilationAllowed, null, this.graphFetchExecutionConfiguration, this.logSQLWithParamValues);
 

diff --git a/...src/main/java/org/finos/legend/engine/language/pure/grammar/from/domain/DomainParser.java b/...src/main/java/org/finos/legend/engine/language/pure/grammar/from/domain/DomainParser.java
@@ -88,7 +88,7 @@ public Lambda parseLambda(String code, PureGrammarParserContext parserContext, S
     {
         ParseTreeWalkerSourceInformation lambdaWalkerSourceInformation = new ParseTreeWalkerSourceInformation.Builder(sourceId, lineOffset, columnOffset).withReturnSourceInfo(returnSourceInfo).build();
         String prefix = "function go():Any[*]{";
-        String fullCode = prefix + code + "}";
+        String fullCode = prefix + code + "\n}";
         ParseTreeWalkerSourceInformation walkerSourceInformation = new ParseTreeWalkerSourceInformation.Builder(lambdaWalkerSourceInformation)
                 // NOTE: as we prepend the lambda with this prefix, we need to subtract this prefix length from the column offset
                 .withColumnOffset(lambdaWalkerSourceInformation.getColumnOffset() - prefix.length()).build();

diff --git a/...ava/org/finos/legend/engine/language/pure/grammar/test/roundtrip/TestLambdaRoundtrip.java b/...ava/org/finos/legend/engine/language/pure/grammar/test/roundtrip/TestLambdaRoundtrip.java
@@ -435,6 +435,12 @@ public void testRenderingFunctionExpressionWithSinglePrimitiveArgument()
         testLambda("|%9999-12-30T19:00:00.0000->someDateFn()", "|someDateFn(%9999-12-30T19:00:00.0000)");
     }
 
+    @Test
+    public void testLambdaEndingOnComment()
+    {
+        testLambda("|1 + 1// comment here", "|1 + 1");
+    }
+
     static void testLambda(String text)
     {
         testLambda(text, text);