From 246fa969e5a540d66939f81a557a3cc9a9ab98c6 Mon Sep 17 00:00:00 2001
From: Yannan <73408381+YannanGao-gs@users.noreply.github.com>
Date: Thu, 10 Aug 2023 14:09:27 -0400
Subject: [PATCH] support generating schemas/tables from query without
 authentication (#2487)

---
 .changeset/brave-kiwis-yell.md                |  5 ++
 .changeset/red-carpets-taste.md               |  3 +
 .../testable/ServiceTestDataEditor.tsx        | 73 ++++++++++++++-----
 .../service/testable/ServiceTestDataState.ts  | 61 ++++++++++++++++
 .../graph-manager/AbstractPureGraphManager.ts | 10 +++
 .../protocol/pure/v1/V1_PureGraphManager.ts   | 34 +++++++++
 .../protocol/pure/v1/engine/V1_Engine.ts      | 19 +++++
 .../pure/v1/engine/V1_EngineServerClient.ts   | 17 +++++
 .../service/V1_TestDataGenerationInput.ts     | 46 ++++++++++++
 .../service/V1_TestDataGenerationResult.ts    | 50 +++++++++++++
 .../service/TestGenerationResult.ts           | 25 +++++++
 packages/legend-graph/src/index.ts            |  1 +
 12 files changed, 326 insertions(+), 18 deletions(-)
 create mode 100644 .changeset/brave-kiwis-yell.md
 create mode 100644 .changeset/red-carpets-taste.md
 create mode 100644 packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/service/V1_TestDataGenerationInput.ts
 create mode 100644 packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/service/V1_TestDataGenerationResult.ts
 create mode 100644 packages/legend-graph/src/graph/metamodel/pure/packageableElements/service/TestGenerationResult.ts

diff --git a/.changeset/brave-kiwis-yell.md b/.changeset/brave-kiwis-yell.md
new file mode 100644
index 0000000000..622c1e5bea
--- /dev/null
+++ b/.changeset/brave-kiwis-yell.md
@@ -0,0 +1,5 @@
+---
+'@finos/legend-application-studio': patch
+---
+
+support generating schemas/tables from query without authentication
diff --git a/.changeset/red-carpets-taste.md b/.changeset/red-carpets-taste.md
new file mode 100644
index 0000000000..d78bbd6f64
--- /dev/null
+++ b/.changeset/red-carpets-taste.md
@@ -0,0 +1,3 @@
+---
+'@finos/legend-graph': patch
+---
diff --git a/packages/legend-application-studio/src/components/editor/editor-group/service-editor/testable/ServiceTestDataEditor.tsx b/packages/legend-application-studio/src/components/editor/editor-group/service-editor/testable/ServiceTestDataEditor.tsx
index 0f252d9344..82eb5e3372 100644
--- a/packages/legend-application-studio/src/components/editor/editor-group/service-editor/testable/ServiceTestDataEditor.tsx
+++ b/packages/legend-application-studio/src/components/editor/editor-group/service-editor/testable/ServiceTestDataEditor.tsx
@@ -16,10 +16,12 @@
 
 import {
   BlankPanelPlaceholder,
+  CaretDownIcon,
   clsx,
   ContextMenu,
   CustomSelectorInput,
   Dialog,
+  DropdownMenu,
   InfoCircleIcon,
   MaskIcon,
   MenuContent,
@@ -199,6 +201,12 @@ export const ConnectionTestDataEditor = observer(
       }
     };
 
+    const generateQuerySchemas = (): void => {
+      flowResult(connectionTestDataState.generateQuerySchemas()).catch(
+        applicationStore.alertUnhandledError,
+      );
+    };
+
     return (
       <div className="service-test-data-editor">
         <div className="service-test-suite-editor__header">
@@ -231,22 +239,45 @@ export const ConnectionTestDataEditor = observer(
                 <MaskIcon />
               </button>
             </div>
-            <button
-              className="panel__header__action service-execution-editor__test-data__generate-btn"
-              onClick={generateTestData}
-              title="Generate test data if possible"
-              disabled={
-                connectionTestDataState.generatingTestDataState.isInProgress
-              }
-              tabIndex={-1}
-            >
-              <div className="service-execution-editor__test-data__generate-btn__label">
-                <RefreshIcon className="service-execution-editor__test-data__generate-btn__label__icon" />
-                <div className="service-execution-editor__test-data__generate-btn__label__title">
+            <div className="btn__dropdown-combo btn__dropdown-combo--primary">
+              <button
+                className="btn__dropdown-combo__label"
+                onClick={generateTestData}
+                title="Generate test data if possible"
+                disabled={
+                  connectionTestDataState.generatingTestDataState.isInProgress
+                }
+                tabIndex={-1}
+              >
+                <RefreshIcon className="btn__dropdown-combo__label__icon" />
+                <div className="btn__dropdown-combo__label__title">
                   Generate
                 </div>
-              </div>
-            </button>
+              </button>
+              <DropdownMenu
+                className="btn__dropdown-combo__dropdown-btn"
+                content={
+                  <MenuContent>
+                    <MenuContentItem
+                      className="btn__dropdown-combo__option"
+                      onClick={generateQuerySchemas}
+                      disabled={
+                        connectionTestDataState.generateSchemaQueryState
+                          .isInProgress
+                      }
+                    >
+                      Generate Query Schemas
+                    </MenuContentItem>
+                  </MenuContent>
+                }
+                menuProps={{
+                  anchorOrigin: { vertical: 'bottom', horizontal: 'right' },
+                  transformOrigin: { vertical: 'top', horizontal: 'right' },
+                }}
+              >
+                <CaretDownIcon />
+              </DropdownMenu>
+            </div>
             <button
               className="panel__header__action service-execution-editor__test-data__generate-btn"
               onClick={openShared}
@@ -631,10 +662,16 @@ export const ServiceTestDataEditor = observer(
               </ResizablePanelSplitter>
               <ResizablePanel minSize={600}>
                 <PanelLoadingIndicator
-                  isLoading={Boolean(
-                    testDataState.selectedDataState?.generatingTestDataState
-                      .isInProgress,
-                  )}
+                  isLoading={
+                    Boolean(
+                      testDataState.selectedDataState?.generatingTestDataState
+                        .isInProgress,
+                    ) ||
+                    Boolean(
+                      testDataState.selectedDataState?.generateSchemaQueryState
+                        .isInProgress,
+                    )
+                  }
                 />
                 {testDataState.selectedDataState && (
                   <ConnectionTestDataEditor
diff --git a/packages/legend-application-studio/src/stores/editor/editor-state/element-editor-state/service/testable/ServiceTestDataState.ts b/packages/legend-application-studio/src/stores/editor/editor-state/element-editor-state/service/testable/ServiceTestDataState.ts
index 333c0606fb..dbf111b3cd 100644
--- a/packages/legend-application-studio/src/stores/editor/editor-state/element-editor-state/service/testable/ServiceTestDataState.ts
+++ b/packages/legend-application-studio/src/stores/editor/editor-state/element-editor-state/service/testable/ServiceTestDataState.ts
@@ -23,6 +23,7 @@ import {
   type RawLambda,
   type DataElement,
   type Mapping,
+  type TestDataGenerationResult,
   ConnectionTestData,
   PureSingleExecution,
   PureMultiExecution,
@@ -37,6 +38,7 @@ import {
   getAllIdentifiedConnectionsFromRuntime,
   getAllIdentifiedServiceConnections,
   Database,
+  RuntimePointer,
 } from '@finos/legend-graph';
 import {
   type GeneratorFn,
@@ -47,6 +49,8 @@ import {
   guaranteeNonNullable,
   returnUndefOnError,
   getNullableFirstEntry,
+  assertType,
+  assertTrue,
 } from '@finos/legend-shared';
 import { action, flow, flowResult, makeObservable, observable } from 'mobx';
 import type { EditorStore } from '../../../../EditorStore.js';
@@ -156,6 +160,7 @@ export class ConnectionTestDataState {
   readonly connectionData: ConnectionTestData;
   readonly parametersState: ServiceTestDataParametersState;
   readonly generatingTestDataState = ActionState.create();
+  readonly generateSchemaQueryState = ActionState.create();
   useSharedModal = false;
 
   embeddedEditorState: EmbeddedDataEditorState;
@@ -167,6 +172,7 @@ export class ConnectionTestDataState {
   ) {
     makeObservable(this, {
       generatingTestDataState: observable,
+      generateSchemaQueryState: observable,
       embeddedEditorState: observable,
       useSharedModal: observable,
       anonymizeGeneratedData: observable,
@@ -175,6 +181,7 @@ export class ConnectionTestDataState {
       buildEmbeddedEditorState: action,
       generateTestData: flow,
       generateTestDataForDatabaseConnection: flow,
+      generateQuerySchemas: flow,
     });
     this.testDataState = testDataState;
     this.editorStore = testDataState.editorStore;
@@ -326,6 +333,60 @@ export class ConnectionTestDataState {
     }
   }
 
+  *generateQuerySchemas(): GeneratorFn<void> {
+    try {
+      this.generateSchemaQueryState.inProgress();
+      const connection = guaranteeNonNullable(
+        this.resolveConnectionValue(this.connectionData.connectionId),
+        `Unable to resolve connection ID '${this.connectionData.connectionId}`,
+      );
+      if (connection instanceof DatabaseConnection) {
+        const serviceExecutionParameters = guaranteeNonNullable(
+          this.testDataState.testSuiteState.testableState.serviceEditorState
+            .executionState.serviceExecutionParameters,
+        );
+        assertType(
+          serviceExecutionParameters.runtime,
+          RuntimePointer,
+          'Expected runtime type to be RuntimePointer',
+        );
+        const testDataGenerationResult =
+          (yield this.editorStore.graphManagerState.graphManager.generateTestData(
+            getExecutionQueryFromRawLambda(
+              serviceExecutionParameters.query,
+              this.parametersState.parameterStates,
+              this.editorStore.graphManagerState,
+            ),
+            serviceExecutionParameters.mapping.path,
+            serviceExecutionParameters.runtime.packageableRuntime.value.path,
+            this.editorStore.graphManagerState.graph,
+          )) as TestDataGenerationResult;
+        assertTrue(
+          testDataGenerationResult.data.length >= 1,
+          'Expected generated data to at least have one data',
+        );
+        service_setConnectionTestDataEmbeddedData(
+          this.connectionData,
+          guaranteeNonNullable(testDataGenerationResult.data[0]),
+          this.editorStore.changeDetectionState.observerContext,
+        );
+        this.embeddedEditorState = new EmbeddedDataEditorState(
+          this.testDataState.editorStore,
+          this.connectionData.testData,
+        );
+      }
+      this.generateSchemaQueryState.pass();
+    } catch (error) {
+      assertErrorThrown(error);
+      this.editorStore.applicationStore.notificationService.notifyError(
+        `Unable to generate query schemas: ${error.message}`,
+      );
+      this.generateSchemaQueryState.fail();
+    } finally {
+      this.generateSchemaQueryState.complete();
+    }
+  }
+
   changeEmbeddedData(val: EmbeddedData): void {
     service_setConnectionTestDataEmbeddedData(
       this.connectionData,
diff --git a/packages/legend-graph/src/graph-manager/AbstractPureGraphManager.ts b/packages/legend-graph/src/graph-manager/AbstractPureGraphManager.ts
index d4f9329e7d..b3dd7b9c89 100644
--- a/packages/legend-graph/src/graph-manager/AbstractPureGraphManager.ts
+++ b/packages/legend-graph/src/graph-manager/AbstractPureGraphManager.ts
@@ -101,6 +101,7 @@ import type { FunctionActivatorConfiguration } from './action/functionActivator/
 import type { FunctionActivator } from '../graph/metamodel/pure/packageableElements/function/FunctionActivator.js';
 import type { RelationalDatabaseConnection } from '../STO_Relational_Exports.js';
 import type { ArtifactGenerationExtensionResult } from './action/generation/ArtifactGenerationExtensionResult.js';
+import type { TestDataGenerationResult } from '../graph/metamodel/pure/packageableElements/service/TestGenerationResult.js';
 
 export interface TEMPORARY__EngineSetupConfig {
   env: string;
@@ -412,6 +413,15 @@ export abstract class AbstractPureGraphManager {
     graph: PureModel,
   ): Promise<Entity[]>;
 
+  // --------------------------------------------- Test Data Generation ---------------------------------------------
+
+  abstract generateTestData(
+    query: RawLambda,
+    mapping: string,
+    runtime: string,
+    graph: PureModel,
+  ): Promise<TestDataGenerationResult>;
+
   // ------------------------------------------- External Format ----------------------------------
 
   abstract getAvailableExternalFormatsDescriptions(): Promise<
diff --git a/packages/legend-graph/src/graph-manager/protocol/pure/v1/V1_PureGraphManager.ts b/packages/legend-graph/src/graph-manager/protocol/pure/v1/V1_PureGraphManager.ts
index df46ff80ae..e5840683b7 100644
--- a/packages/legend-graph/src/graph-manager/protocol/pure/v1/V1_PureGraphManager.ts
+++ b/packages/legend-graph/src/graph-manager/protocol/pure/v1/V1_PureGraphManager.ts
@@ -304,6 +304,9 @@ import {
   V1_buildArtifactsByExtensionElement,
 } from './engine/generation/V1_ArtifactGenerationExtensionApi.js';
 import type { V1_RawValueSpecification } from './model/rawValueSpecification/V1_RawValueSpecification.js';
+import { V1_TestDataGenerationInput } from './engine/service/V1_TestDataGenerationInput.js';
+import type { TestDataGenerationResult } from '../../../../graph/metamodel/pure/packageableElements/service/TestGenerationResult.js';
+import { V1_buildTestDataGenerationResult } from './engine/service/V1_TestDataGenerationResult.js';
 
 class V1_PureModelContextDataIndex {
   elements: V1_PackageableElement[] = [];
@@ -2060,6 +2063,37 @@ export class V1_PureGraphManager extends AbstractPureGraphManager {
     return this.pureModelContextDataToEntities(generatedModel);
   }
 
+  // --------------------------------------------- Test Data Generation ---------------------------------------------
+
+  async generateTestData(
+    query: RawLambda,
+    mapping: string,
+    runtime: string,
+    graph: PureModel,
+  ): Promise<TestDataGenerationResult> {
+    const testDataGenerationInput = new V1_TestDataGenerationInput();
+    testDataGenerationInput.query = V1_transformRawLambda(
+      query,
+      new V1_GraphTransformerContextBuilder(
+        this.pluginManager.getPureProtocolProcessorPlugins(),
+      ).build(),
+    );
+    testDataGenerationInput.mapping = mapping;
+    testDataGenerationInput.runtime = runtime;
+    const graphData = this.getFullGraphModelContext(
+      graph,
+      V1_PureGraphManager.DEV_PROTOCOL_VERSION,
+    );
+    testDataGenerationInput.model = graphData;
+    return V1_buildTestDataGenerationResult(
+      await this.engine.generateTestData(
+        testDataGenerationInput,
+        this.pluginManager.getPureProtocolProcessorPlugins(),
+      ),
+      this.getBuilderContext(graph, graph, new V1_Mapping()), // use a dummy V1_PackageableElement to generate V1_GraphBuilderContext which is used in V1_buildEmbeddedData()
+    );
+  }
+
   // ------------------------------------------- Test  -------------------------------------------
 
   async runTests(
diff --git a/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/V1_Engine.ts b/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/V1_Engine.ts
index 9a4d9b00c4..ff6f1bae32 100644
--- a/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/V1_Engine.ts
+++ b/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/V1_Engine.ts
@@ -136,6 +136,11 @@ import {
   V1_ArtifactGenerationExtensionInput,
 } from './generation/V1_ArtifactGenerationExtensionApi.js';
 import { V1_DatabaseToModelGenerationInput } from './relational/V1_DatabaseToModelGenerationInput.js';
+import { V1_TestDataGenerationInput } from './service/V1_TestDataGenerationInput.js';
+import {
+  type V1_TestDataGenerationResult,
+  V1_testDataGenerationResultModelSchema,
+} from './service/V1_TestDataGenerationResult.js';
 
 class V1_EngineConfig extends TEMPORARY__AbstractEngineConfig {
   private engine: V1_Engine;
@@ -719,6 +724,20 @@ export class V1_Engine {
     );
   }
 
+  // --------------------------------------------- Test Data Generation ---------------------------------------------
+
+  async generateTestData(
+    input: V1_TestDataGenerationInput,
+    plugins: PureProtocolProcessorPlugin[],
+  ): Promise<V1_TestDataGenerationResult> {
+    return deserialize(
+      V1_testDataGenerationResultModelSchema(plugins),
+      await this.engineServerClient.generateTestData(
+        V1_TestDataGenerationInput.serialization.toJson(input),
+      ),
+    );
+  }
+
   // ------------------------------------------- File Generation -------------------------------------------
 
   async getAvailableGenerationConfigurationDescriptions(): Promise<
diff --git a/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/V1_EngineServerClient.ts b/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/V1_EngineServerClient.ts
index b52523df69..0c8b147389 100644
--- a/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/V1_EngineServerClient.ts
+++ b/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/V1_EngineServerClient.ts
@@ -78,12 +78,15 @@ import type {
   V1_ArtifactGenerationExtensionOutput,
 } from './generation/V1_ArtifactGenerationExtensionApi.js';
 import type { V1_DatabaseToModelGenerationInput } from './relational/V1_DatabaseToModelGenerationInput.js';
+import type { V1_TestDataGenerationInput } from './service/V1_TestDataGenerationInput.js';
+import type { V1_TestDataGenerationResult } from './service/V1_TestDataGenerationResult.js';
 
 enum CORE_ENGINE_ACTIVITY_TRACE {
   GRAMMAR_TO_JSON = 'transform Pure code to protocol',
   JSON_TO_GRAMMAR = 'transform protocol to Pure code',
 
   DATABASE_TO_MODELS = 'generate models from database',
+  TEST_DATA_GENERATION = 'generate test data',
 
   EXTERNAL_FORMAT_TO_PROTOCOL = 'transform external format code to protocol',
   GENERATE_FILE = 'generate file',
@@ -502,6 +505,20 @@ export class V1_EngineServerClient extends AbstractServerClient {
     PlainObject<V1_GenerationConfigurationDescription>[]
   > => this.get(`${this._pure()}/schemaGeneration/availableGenerations`);
 
+  // --------------------------------------------- Test Data Generation ---------------------------------------------
+
+  _testDataGeneration = (): string => `${this._pure()}/testData/generation`;
+
+  generateTestData(
+    input: PlainObject<V1_TestDataGenerationInput>,
+  ): Promise<PlainObject<V1_TestDataGenerationResult>> {
+    return this.postWithTracing(
+      this.getTraceData(CORE_ENGINE_ACTIVITY_TRACE.TEST_DATA_GENERATION),
+      `${this._testDataGeneration()}/DONOTUSE_generateTestData`,
+      this.debugPayload(input, CORE_ENGINE_ACTIVITY_TRACE.TEST_DATA_GENERATION),
+    );
+  }
+
   // ------------------------------------------- Compile -------------------------------------------
 
   compile = (
diff --git a/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/service/V1_TestDataGenerationInput.ts b/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/service/V1_TestDataGenerationInput.ts
new file mode 100644
index 0000000000..92ce1d60f9
--- /dev/null
+++ b/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/service/V1_TestDataGenerationInput.ts
@@ -0,0 +1,46 @@
+/**
+ * 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 { SerializationFactory, usingModelSchema } from '@finos/legend-shared';
+import { createModelSchema, optional, primitive } from 'serializr';
+import type { V1_PureModelContext } from '../../model/context/V1_PureModelContext.js';
+import type { V1_RawLambda } from '../../model/rawValueSpecification/V1_RawLambda.js';
+import { V1_pureModelContextPropSchema } from '../../transformation/pureProtocol/V1_PureProtocolSerialization.js';
+import { V1_rawLambdaModelSchema } from '../../transformation/pureProtocol/serializationHelpers/V1_RawValueSpecificationSerializationHelper.js';
+
+export class V1_TestDataGenerationInput {
+  clientVersion: string | undefined;
+  /**
+   * Studio does not process value specification, they are left in raw JSON form
+   *
+   * Should be optional. For now, it's mandatory.
+   * @discrepancy model
+   */
+  query!: V1_RawLambda;
+  mapping!: string;
+  runtime!: string;
+  model!: V1_PureModelContext;
+
+  static readonly serialization = new SerializationFactory(
+    createModelSchema(V1_TestDataGenerationInput, {
+      clientVersion: optional(primitive()),
+      query: usingModelSchema(V1_rawLambdaModelSchema),
+      mapping: primitive(),
+      model: V1_pureModelContextPropSchema,
+      runtime: primitive(),
+    }),
+  );
+}
diff --git a/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/service/V1_TestDataGenerationResult.ts b/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/service/V1_TestDataGenerationResult.ts
new file mode 100644
index 0000000000..54c9decfae
--- /dev/null
+++ b/packages/legend-graph/src/graph-manager/protocol/pure/v1/engine/service/V1_TestDataGenerationResult.ts
@@ -0,0 +1,50 @@
+/**
+ * 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 { createModelSchema, list, custom, type ModelSchema } from 'serializr';
+import type { V1_EmbeddedData } from '../../model/data/V1_EmbeddedData.js';
+import {
+  V1_deserializeEmbeddedDataType,
+  V1_serializeEmbeddedDataType,
+} from '../../transformation/pureProtocol/serializationHelpers/V1_DataElementSerializationHelper.js';
+import type { PureProtocolProcessorPlugin } from '../../../PureProtocolProcessorPlugin.js';
+import type { V1_GraphBuilderContext } from '../../transformation/pureGraph/to/V1_GraphBuilderContext.js';
+import { TestDataGenerationResult } from '../../../../../../graph/metamodel/pure/packageableElements/service/TestGenerationResult.js';
+import { V1_buildEmbeddedData } from '../../transformation/pureGraph/to/helpers/V1_DataElementBuilderHelper.js';
+
+export class V1_TestDataGenerationResult {
+  data: V1_EmbeddedData[] = [];
+}
+
+export const V1_testDataGenerationResultModelSchema = (
+  plugins: PureProtocolProcessorPlugin[],
+): ModelSchema<V1_TestDataGenerationResult> =>
+  createModelSchema(V1_TestDataGenerationResult, {
+    data: list(
+      custom(
+        (val) => V1_serializeEmbeddedDataType(val, plugins),
+        (val) => V1_deserializeEmbeddedDataType(val, plugins),
+      ),
+    ),
+  });
+
+export const V1_buildTestDataGenerationResult = (
+  protocol: V1_TestDataGenerationResult,
+  context: V1_GraphBuilderContext,
+): TestDataGenerationResult =>
+  new TestDataGenerationResult(
+    protocol.data.map((p) => V1_buildEmbeddedData(p, context)),
+  );
diff --git a/packages/legend-graph/src/graph/metamodel/pure/packageableElements/service/TestGenerationResult.ts b/packages/legend-graph/src/graph/metamodel/pure/packageableElements/service/TestGenerationResult.ts
new file mode 100644
index 0000000000..9987babf11
--- /dev/null
+++ b/packages/legend-graph/src/graph/metamodel/pure/packageableElements/service/TestGenerationResult.ts
@@ -0,0 +1,25 @@
+/**
+ * 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 { EmbeddedData } from '../../data/EmbeddedData.js';
+
+export class TestDataGenerationResult {
+  data: EmbeddedData[] = [];
+
+  constructor(data: EmbeddedData[]) {
+    this.data = data;
+  }
+}
diff --git a/packages/legend-graph/src/index.ts b/packages/legend-graph/src/index.ts
index 110c5c0fc2..cd85d579f9 100644
--- a/packages/legend-graph/src/index.ts
+++ b/packages/legend-graph/src/index.ts
@@ -70,6 +70,7 @@ export {
   TagExplicitReference,
 } from './graph/metamodel/pure/packageableElements/domain/TagReference.js';
 export * from './graph/metamodel/pure/packageableElements/service/ExecutionEnvironmentInstance.js';
+export { TestDataGenerationResult } from './graph/metamodel/pure/packageableElements/service/TestGenerationResult.js';
 export { FunctionActivator } from './graph/metamodel/pure/packageableElements/function/FunctionActivator.js';
 export { INTERNAL__UnknownFunctionActivator } from './graph/metamodel/pure/packageableElements/function/INTERNAL__UnknownFunctionActivator.js';