feat: Get lambda relation type through lambda return type endpoint

src/LegendLanguageClient.ts

@@ -43,7 +43,6 @@ import {
   JSON_TO_GRAMMAR_LAMBDA_BATCH_COMMAND_ID,
   CHECK_DATASET_ENTITLEMENTS_COMMAND_ID,
   SURVEY_DATASETS_COMMAND_ID,
-  GET_LAMBDA_RELATION_TYPE_COMMAND_ID,
   GRAMMAR_TO_JSON_VALUE_SPECIFICATION_BATCH_ID,
 } from './utils/Const';
 import type { PlainObject } from './utils/SerializationUtils';
@@ -56,8 +55,8 @@ import type { ExecuteTestRequest } from './model/ExecuteTestRequest';
 import type { LegendTestExecutionResult } from './model/LegendTestExecutionResult';
 import { LegendEntity } from './model/LegendEntity';
 import {
-  V1_GrammarParserBatchInputEntry,
   type EXECUTION_SERIALIZATION_FORMAT,
+  type V1_GrammarParserBatchInputEntry,
   type V1_ParameterValue,
   type V1_RawExecutionContext,
   type V1_RawLambda,
@@ -389,20 +388,6 @@ export class LegendLanguageClient extends LanguageClient {
     );
   }
 
-  async getLambdaRelationType(
-    entityTextLocation: TextLocation,
-    lambda: V1_RawLambda,
-  ): Promise<string> {
-    return commands.executeCommand(
-      LEGEND_COMMAND,
-      entityTextLocation,
-      GET_LAMBDA_RELATION_TYPE_COMMAND_ID,
-      {
-        lambda: JSON.stringify(lambda),
-      },
-    );
-  }
-
   async generateDatasetSpecifications(
     entityTextLocation: TextLocation,
     mapping: string,

src/graph/V1_LSPEngine.ts

@@ -124,8 +124,6 @@ import {
   GET_CLASSIFIER_PATH_MAP_RESPONSE,
   GET_CURRENT_USER_ID_REQUEST_ID,
   GET_CURRENT_USER_ID_RESPONSE,
-  GET_LAMBDA_RELATION_TYPE_COMMAND_ID,
-  GET_LAMBDA_RELATION_TYPE_RESPONSE,
   GET_LAMBDA_RETURN_TYPE_COMMAND_ID,
   GET_LAMBDA_RETURN_TYPE_RESPONSE,
   GET_SUBTYPE_INFO_REQUEST_ID,
@@ -147,6 +145,7 @@ import {
   executeInputToLSPInput,
   surveyDatasetsInputToLSPInput,
 } from '../utils/GraphUtils';
+import { type V1_LSPLambdaReturnTypeResult } from '../model/engine/LambdaReturnTypeResult';
 
 class V1_LSPEngine_Config extends TEMPORARY__AbstractEngineConfig {}
 
@@ -410,10 +409,10 @@ export class V1_LSPEngine implements V1_GraphManagerEngine {
   ): Promise<RelationTypeMetadata> {
     const response = await postAndWaitForMessage<LegendExecutionResult[]>(
       {
-        command: GET_LAMBDA_RELATION_TYPE_COMMAND_ID,
+        command: GET_LAMBDA_RETURN_TYPE_COMMAND_ID,
         msg: V1_LambdaReturnTypeInput.serialization.toJson(rawInput),
       },
-      GET_LAMBDA_RELATION_TYPE_RESPONSE,
+      GET_LAMBDA_RETURN_TYPE_RESPONSE,
     );
     if (response?.[0]?.type === LegendExecutionResultType.ERROR) {
       const sourceInformation = response[0].location
@@ -427,7 +426,14 @@ export class V1_LSPEngine implements V1_GraphManagerEngine {
       );
     }
     const v1_relationType = V1_RelationType.serialization.fromJson(
-      JSON.parse(guaranteeNonNullable(response?.[0]?.message)),
+      guaranteeNonNullable(
+        (
+          JSON.parse(
+            guaranteeNonNullable(response?.[0]?.message),
+          ) as V1_LSPLambdaReturnTypeResult
+        ).relationType,
+        'Lambda return type response does not contain relationType',
+      ),
     );
     const result = new RelationTypeMetadata();
     result.columns = v1_relationType.columns.map(

src/model/engine/LambdaReturnTypeResult.ts

@@ -0,0 +1,26 @@
+/**
+ * Copyright (c) 2020-present, 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.
+ */
+
+import {
+  type PlainObject,
+  type V1_LambdaReturnTypeResult,
+  type V1_RelationType,
+} from '@finos/legend-vscode-extension-dependencies';
+
+export interface V1_LSPLambdaReturnTypeResult
+  extends V1_LambdaReturnTypeResult {
+  relationType?: PlainObject<V1_RelationType>;
+}

src/utils/Const.ts

@@ -65,7 +65,6 @@ export const JSON_TO_GRAMMAR_LAMBDA_BATCH_COMMAND_ID =
 export const GRAMMAR_TO_JSON_VALUE_SPECIFICATION_BATCH_ID =
   'legend.grammarToJson.valueSpecification.batch';
 export const GET_LAMBDA_RETURN_TYPE_COMMAND_ID = 'legend.lambda.returnType';
-export const GET_LAMBDA_RELATION_TYPE_COMMAND_ID = 'legend.lambda.relationType';
 export const SURVEY_DATASETS_COMMAND_ID = 'legend.entitlements.surveyDatasets';
 export const CHECK_DATASET_ENTITLEMENTS_COMMAND_ID =
   'legend.entitlements.checkDatasetEntitlements';
@@ -103,7 +102,6 @@ export const GRAMMAR_TO_JSON_LAMBDA_RESPONSE = `${GRAMMAR_TO_JSON_LAMBDA_COMMAND
 export const JSON_TO_GRAMMAR_LAMBDA_BATCH_RESPONSE = `${JSON_TO_GRAMMAR_LAMBDA_BATCH_COMMAND_ID}/response`;
 export const GRAMMAR_TO_JSON_VALUE_SPECIFICATION_BATCH_RESPONSE = `${GRAMMAR_TO_JSON_VALUE_SPECIFICATION_BATCH_ID}/response`;
 export const GET_LAMBDA_RETURN_TYPE_RESPONSE = `${GET_LAMBDA_RETURN_TYPE_COMMAND_ID}/response`;
-export const GET_LAMBDA_RELATION_TYPE_RESPONSE = `${GET_LAMBDA_RELATION_TYPE_COMMAND_ID}/response`;
 export const EXPORT_DATA_RESPONSE = `${EXPORT_DATA_COMMAND_ID}/response`;
 export const SURVEY_DATASETS_RESPONSE = `${SURVEY_DATASETS_COMMAND_ID}/response`;
 export const CHECK_DATASET_ENTITLEMENTS_RESPONSE = `${CHECK_DATASET_ENTITLEMENTS_COMMAND_ID}/response`;

src/webviews/utils.ts

@@ -41,8 +41,6 @@ import {
   GET_CLASSIFIER_PATH_MAP_RESPONSE,
   GET_CURRENT_USER_ID_REQUEST_ID,
   GET_CURRENT_USER_ID_RESPONSE,
-  GET_LAMBDA_RELATION_TYPE_COMMAND_ID,
-  GET_LAMBDA_RELATION_TYPE_RESPONSE,
   GET_LAMBDA_RETURN_TYPE_COMMAND_ID,
   GET_LAMBDA_RETURN_TYPE_RESPONSE,
   GET_PROJECT_ENTITIES,
@@ -336,18 +334,6 @@ export const handleV1LSPEngineMessage = async (
       });
       return true;
     }
-    case GET_LAMBDA_RELATION_TYPE_COMMAND_ID: {
-      const { lambda } = message.msg;
-      const result = await client.getLambdaRelationType(
-        entityTextLocation,
-        lambda,
-      );
-      webview.postMessage({
-        command: GET_LAMBDA_RELATION_TYPE_RESPONSE,
-        result,
-      });
-      return true;
-    }
     case SURVEY_DATASETS_COMMAND_ID: {
       const { mapping, runtime, lambda } = message.msg;
       const result = await client.generateDatasetSpecifications(