Move query builder command handling to FunctionExecutionSupport inter…

…face (#262) * refactor: Move query builder command handling to FunctionExecutionSupport * refactor: Fix .equals param order * refactor: PR feedback * refactor: whitespace change * refactor: Make collectResults take FunctionExecutionSupport * refactor: whitespace change
finos · Nov 11, 2024 · 83eacca · 83eacca
1 parent 919fddd
commit 83eacca
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 41 deletions.
diff --git a/...rc/main/java/org/finos/legend/engine/ide/lsp/extension/core/FunctionExecutionSupport.java b/...rc/main/java/org/finos/legend/engine/ide/lsp/extension/core/FunctionExecutionSupport.java
@@ -38,10 +38,7 @@
 import org.finos.legend.engine.entitlement.services.EntitlementModelObjectMapperFactory;
 import org.finos.legend.engine.entitlement.services.EntitlementServiceExtension;
 import org.finos.legend.engine.entitlement.services.EntitlementServiceExtensionLoader;
-import org.finos.legend.engine.ide.lsp.extension.AbstractLSPGrammarExtension;
-import org.finos.legend.engine.ide.lsp.extension.CommandConsumer;
-import org.finos.legend.engine.ide.lsp.extension.CompileResult;
-import org.finos.legend.engine.ide.lsp.extension.Constants;
+import org.finos.legend.engine.ide.lsp.extension.*;
 import org.finos.legend.engine.ide.lsp.extension.execution.LegendCommandType;
 import org.finos.legend.engine.ide.lsp.extension.execution.LegendExecutionResult;
 import org.finos.legend.engine.ide.lsp.extension.execution.LegendInputParameter;
@@ -124,6 +121,49 @@ default String getExecutionKey(PackageableElement element, Map<String, Object> a
         return "";
     }
 
+    static Iterable<? extends LegendExecutionResult> execute(FunctionExecutionSupport executionSupport, SectionState section, String entityPath, String commandId, Map<String, String> executableArgs, Map<String, Object> inputParameters)
+    {
+        switch (commandId)
+        {
+            case FunctionExecutionSupport.EXECUTE_COMMAND_ID:
+            {
+                return FunctionExecutionSupport.executeFunction(executionSupport, section, entityPath, inputParameters);
+            }
+            case FunctionExecutionSupport.EXECUTE_QUERY_ID:
+            {
+                return FunctionExecutionSupport.executeQuery(executionSupport, section, entityPath, executableArgs, inputParameters);
+            }
+            case FunctionExecutionSupport.GENERATE_EXECUTION_PLAN_ID:
+            {
+                return FunctionExecutionSupport.generateExecutionPlan(executionSupport, section, entityPath, executableArgs, inputParameters);
+            }
+            case FunctionExecutionSupport.GRAMMAR_TO_JSON_LAMBDA_ID:
+            {
+                return FunctionExecutionSupport.convertGrammarToLambdaJson(executionSupport, section, entityPath, executableArgs, inputParameters);
+            }
+            case FunctionExecutionSupport.JSON_TO_GRAMMAR_LAMBDA_BATCH_ID:
+            {
+                return FunctionExecutionSupport.convertLambdaJsonToGrammarBatch(executionSupport, section, entityPath, executableArgs, inputParameters);
+            }
+            case FunctionExecutionSupport.GET_LAMBDA_RETURN_TYPE_ID:
+            {
+                return FunctionExecutionSupport.getLambdaReturnType(executionSupport, section, entityPath, executableArgs, inputParameters);
+            }
+            case FunctionExecutionSupport.SURVEY_DATASETS_ID:
+            {
+                return FunctionExecutionSupport.generateDatasetSpecifications(executionSupport, section, entityPath, executableArgs, inputParameters);
+            }
+            case FunctionExecutionSupport.CHECK_DATASET_ENTITLEMENTS_ID:
+            {
+                return FunctionExecutionSupport.generateEntitlementReports(executionSupport, section, entityPath, executableArgs, inputParameters);
+            }
+            default:
+            {
+                throw new UnsupportedOperationException("Unsupported command: " + commandId);
+            }
+        }
+    }
+
     SingleExecutionPlan getExecutionPlan(PackageableElement element, Lambda lambda, PureModel pureModel, Map<String, Object> args, String clientVersion);
 
     static SingleExecutionPlan getExecutionPlan(Lambda lambda, String mappingPath, Runtime runtime, ExecutionContext context, PureModel pureModel, String clientVersion)

diff --git a/...src/main/java/org/finos/legend/engine/ide/lsp/extension/core/PureLSPGrammarExtension.java b/...src/main/java/org/finos/legend/engine/ide/lsp/extension/core/PureLSPGrammarExtension.java
@@ -316,10 +316,7 @@ private Stream<Optional<LegendReferenceResolver>> toReferences(Constraint constr
     @Override
     public Iterable<? extends LegendExecutionResult> execute(SectionState section, String entityPath, String commandId, Map<String, String> executableArgs, Map<String, Object> inputParameters)
     {
-        return FunctionExecutionSupport.EXECUTE_COMMAND_ID.equals(commandId) ?
-                FunctionExecutionSupport.executeFunction(this, section, entityPath, inputParameters)
-                :
-                super.execute(section, entityPath, commandId, executableArgs, inputParameters);
+        return FunctionExecutionSupport.execute(this, section, entityPath, commandId, executableArgs, inputParameters);
     }
 
     @Override

diff --git a/...in/java/org/finos/legend/engine/ide/lsp/extension/service/ServiceLSPGrammarExtension.java b/...in/java/org/finos/legend/engine/ide/lsp/extension/service/ServiceLSPGrammarExtension.java
@@ -301,41 +301,9 @@ public Iterable<? extends LegendExecutionResult> execute(SectionState section, S
             {
                 return registerService(section, entityPath);
             }
-            case FunctionExecutionSupport.EXECUTE_COMMAND_ID:
-            {
-                return FunctionExecutionSupport.executeFunction(this, section, entityPath, inputParams);
-            }
-            case FunctionExecutionSupport.EXECUTE_QUERY_ID:
-            {
-                return FunctionExecutionSupport.executeQuery(this, section, entityPath, executableArgs, inputParams);
-            }
-            case FunctionExecutionSupport.GENERATE_EXECUTION_PLAN_ID:
-            {
-                return FunctionExecutionSupport.generateExecutionPlan(this, section, entityPath, executableArgs, inputParams);
-            }
-            case FunctionExecutionSupport.GRAMMAR_TO_JSON_LAMBDA_ID:
-            {
-                return FunctionExecutionSupport.convertGrammarToLambdaJson(this, section, entityPath, executableArgs, inputParams);
-            }
-            case FunctionExecutionSupport.JSON_TO_GRAMMAR_LAMBDA_BATCH_ID:
-            {
-                return FunctionExecutionSupport.convertLambdaJsonToGrammarBatch(this, section, entityPath, executableArgs, inputParams);
-            }
-            case FunctionExecutionSupport.GET_LAMBDA_RETURN_TYPE_ID:
-            {
-                return FunctionExecutionSupport.getLambdaReturnType(this, section, entityPath, executableArgs, inputParams);
-            }
-            case FunctionExecutionSupport.SURVEY_DATASETS_ID:
-            {
-                return FunctionExecutionSupport.generateDatasetSpecifications(this, section, entityPath, executableArgs, inputParams);
-            }
-            case FunctionExecutionSupport.CHECK_DATASET_ENTITLEMENTS_ID:
-            {
-                return FunctionExecutionSupport.generateEntitlementReports(this, section, entityPath, executableArgs, inputParams);
-            }
             default:
             {
-                return super.execute(section, entityPath, commandId, executableArgs, Map.of());
+                return FunctionExecutionSupport.execute(this, section, entityPath, commandId, executableArgs, inputParams);
             }
         }
     }