@finos/legend-application-studio 28.3.4 → 28.3.5

package/src/stores/editor/editor-state/element-editor-state/service/testable/ServiceTestDataState.ts

@@ -25,6 +25,13 @@ import {
   type Mapping,
   type TestDataGenerationResult,
   type GraphManagerState,
+  type Column,
+  type ValueSpecification,
+  type Table,
+  type PrimitiveInstanceValue,
+  type ObserverContext,
+  type PureModel,
+  type RelationalDataType,
   ConnectionTestData,
   PureSingleExecution,
   PureMultiExecution,
@@ -41,6 +48,10 @@ import {
   Database,
   RuntimePointer,
   buildRawLambdaFromLambdaFunction,
+  TableRowIdentifiers,
+  RowIdentifier,
+  ColumnValuePair,
+  TablePtr,
 } from '@finos/legend-graph';
 import {
   type GeneratorFn,
@@ -53,6 +64,7 @@ import {
   getNullableFirstEntry,
   assertType,
   assertTrue,
+  uuid,
 } from '@finos/legend-shared';
 import { action, flow, flowResult, makeObservable, observable } from 'mobx';
 import type { EditorStore } from '../../../../EditorStore.js';
@@ -83,7 +95,9 @@ import {
   buildExecutionParameterValues,
   buildParametersLetLambdaFunc,
   getExecutionQueryFromRawLambda,
+  createMockPrimitiveValueSpecification,
 } from '@finos/legend-query-builder';
+import { getPrimitiveTypeFromRelationalType } from '../../../../utils/MockDataUtils.js';
 
 export const createConnectionTestData = (
   val: IdentifiedConnection,
@@ -97,6 +111,24 @@ export const createConnectionTestData = (
   return connectionTestData;
 };
 
+export const createMockPrimitiveValueSpecificationFromRelationalDataType = (
+  relationalDataType: RelationalDataType,
+  graph: PureModel,
+  observerContext: ObserverContext,
+): ValueSpecification | undefined => {
+  const primitiveTypeFromRelational =
+    getPrimitiveTypeFromRelationalType(relationalDataType);
+
+  if (primitiveTypeFromRelational) {
+    return createMockPrimitiveValueSpecification(
+      primitiveTypeFromRelational,
+      graph,
+      observerContext,
+    );
+  }
+  return undefined;
+};
+
 export class ServiceTestDataParametersState extends LambdaParametersState {
   connectionTestDataState: ConnectionTestDataState;
 
@@ -111,24 +143,38 @@ export class ServiceTestDataParametersState extends LambdaParametersState {
     this.connectionTestDataState = connectionTestDataState;
   }
 
-  openModal(serviceExecutionParameters: {
-    query: RawLambda;
-    mapping: Mapping;
-    runtime: Runtime;
-  }): void {
+  openModal(
+    serviceExecutionParameters: {
+      query: RawLambda;
+      mapping: Mapping;
+      runtime: Runtime;
+    },
+    option?: {
+      generateWithSeedData: boolean;
+    },
+  ): void {
     this.parameterStates = this.build(serviceExecutionParameters.query);
-    this.parameterValuesEditorState.open(
-      (): Promise<void> =>
-        flowResult(
+    this.parameterValuesEditorState.open((): Promise<void> => {
+      if (option?.generateWithSeedData) {
+        return flowResult(
+          this.connectionTestDataState.generateTestDataWithSeedDataForDatabaseConnection(
+            serviceExecutionParameters,
+          ),
+        ).catch(
+          this.connectionTestDataState.editorStore.applicationStore
+            .alertUnhandledError,
+        );
+      } else {
+        return flowResult(
           this.connectionTestDataState.generateTestDataForDatabaseConnection(
             serviceExecutionParameters,
           ),
         ).catch(
           this.connectionTestDataState.editorStore.applicationStore
             .alertUnhandledError,
-        ),
-      PARAMETER_SUBMIT_ACTION.RUN,
-    );
+        );
+      }
+    }, PARAMETER_SUBMIT_ACTION.RUN);
   }
 
   build(query: RawLambda): LambdaParameterState[] {
@@ -157,14 +203,107 @@ export class ServiceTestDataParametersState extends LambdaParametersState {
   }
 }
 
+export class RowIdentifierState {
+  readonly connectionTestDataState: ConnectionTestDataState;
+  readonly _UUID = uuid();
+  column: Column;
+  value: ValueSpecification;
+
+  constructor(
+    connectionTestDataState: ConnectionTestDataState,
+    column: Column,
+    value: ValueSpecification,
+  ) {
+    makeObservable(this, {
+      column: observable,
+      value: observable,
+      updateRowIdentifierColumn: action,
+      updateRowIdentifierValue: action,
+    });
+    this.connectionTestDataState = connectionTestDataState;
+    this.column = column;
+    this.value = value;
+  }
+
+  updateRowIdentifierColumn(col: Column): void {
+    this.column = col;
+  }
+
+  updateRowIdentifierValue(val: ValueSpecification): void {
+    this.value = observe_ValueSpecification(
+      val,
+      this.connectionTestDataState.editorStore.changeDetectionState
+        .observerContext,
+    );
+  }
+}
+
+export class TableRowIdentifierState {
+  readonly connectionTestDataState: ConnectionTestDataState;
+  readonly _UUID = uuid();
+  table: Table;
+  rowIdentifierStates: RowIdentifierState[] = [];
+
+  constructor(
+    connectionTestDataState: ConnectionTestDataState,
+    table: Table,
+    rowIdentifierStates: RowIdentifierState[],
+  ) {
+    makeObservable(this, {
+      table: observable,
+      rowIdentifierStates: observable,
+      updateTable: action,
+      setNewRowIdentifierState: action,
+      addNewRowIdentifierState: action,
+      removeRowIdentifierState: action,
+    });
+    this.connectionTestDataState = connectionTestDataState;
+    this.table = table;
+    this.rowIdentifierStates = rowIdentifierStates;
+  }
+
+  updateTable(table: Table): void {
+    this.table = table;
+  }
+
+  setNewRowIdentifierState(rowIdentifierStates: RowIdentifierState[]): void {
+    this.rowIdentifierStates = rowIdentifierStates;
+  }
+
+  addNewRowIdentifierState(column: Column): void {
+    const valueSpec =
+      createMockPrimitiveValueSpecificationFromRelationalDataType(
+        guaranteeNonNullable(column.type),
+        this.connectionTestDataState.editorStore.graphManagerState.graph,
+        this.connectionTestDataState.editorStore.changeDetectionState
+          .observerContext,
+      );
+    if (valueSpec) {
+      const rowIdentifierState = new RowIdentifierState(
+        this.connectionTestDataState,
+        column,
+        valueSpec,
+      );
+      this.rowIdentifierStates.push(rowIdentifierState);
+    }
+  }
+
+  removeRowIdentifierState(rowIdentifierState: RowIdentifierState): void {
+    deleteEntry(this.rowIdentifierStates, rowIdentifierState);
+  }
+}
+
 export class ConnectionTestDataState {
   readonly editorStore: EditorStore;
   readonly testDataState: ServiceTestDataState;
   readonly connectionData: ConnectionTestData;
   readonly parametersState: ServiceTestDataParametersState;
   readonly generatingTestDataState = ActionState.create();
+  readonly generatingTestDataWithSeedDataState = ActionState.create();
   readonly generateSchemaQueryState = ActionState.create();
   useSharedModal = false;
+  useSeedDataInputModal = false;
+  tableRowIdentifierStates: TableRowIdentifierState[] = [];
 
   embeddedEditorState: EmbeddedDataEditorState;
   anonymizeGeneratedData = true;
@@ -174,17 +313,26 @@ export class ConnectionTestDataState {
     connectionData: ConnectionTestData,
   ) {
     makeObservable(this, {
+      tableRowIdentifierStates: observable,
       generatingTestDataState: observable,
+      generatingTestDataWithSeedDataState: observable,
       generateSchemaQueryState: observable,
       embeddedEditorState: observable,
       useSharedModal: observable,
+      useSeedDataInputModal: observable,
       anonymizeGeneratedData: observable,
       setAnonymizeGeneratedData: action,
+      setUseSeedDataInputModal: action,
+      setNewTableIdentifierState: action,
+      addNewTableIdentifierState: action,
+      removeTableIdentifierState: action,
       changeEmbeddedData: action,
       buildEmbeddedEditorState: action,
       createExecutableQuery: action,
       generateTestData: flow,
       generateTestDataForDatabaseConnection: flow,
+      generateTestDataWithSeedData: flow,
+      generateTestDataWithSeedDataForDatabaseConnection: flow,
       generateQuerySchemas: flow,
     });
     this.testDataState = testDataState;
@@ -214,10 +362,50 @@ export class ConnectionTestDataState {
     }
   }
 
+  getAvailableTables(): Table[] {
+    const runtime =
+      this.testDataState.testSuiteState.testableState.serviceEditorState
+        .executionState.serviceExecutionParameters?.runtime;
+    if (runtime instanceof RuntimePointer) {
+      const db =
+        runtime.packageableRuntime.value.runtimeValue.connections[0]?.store
+          .value;
+      if (db instanceof Database) {
+        return db.schemas.flatMap((schema) => schema.tables);
+      }
+    }
+    return [];
+  }
+
   setUseSharedModal(val: boolean): void {
     this.useSharedModal = val;
   }
 
+  setUseSeedDataInputModal(val: boolean): void {
+    this.useSeedDataInputModal = val;
+  }
+
+  setNewTableIdentifierState(
+    tableRowIdentifierStates: TableRowIdentifierState[],
+  ): void {
+    this.tableRowIdentifierStates = tableRowIdentifierStates;
+  }
+
+  addNewTableIdentifierState(table: Table): void {
+    const tableRowIdentifierState = new TableRowIdentifierState(
+      this,
+      table,
+      [],
+    );
+    this.tableRowIdentifierStates.push(tableRowIdentifierState);
+  }
+
+  removeTableIdentifierState(
+    tableRowIdentifierState: TableRowIdentifierState,
+  ): void {
+    deleteEntry(this.tableRowIdentifierStates, tableRowIdentifierState);
+  }
+
   setAnonymizeGeneratedData(val: boolean): void {
     this.anonymizeGeneratedData = val;
   }
@@ -337,6 +525,144 @@ export class ConnectionTestDataState {
     }
   }
 
+  *generateTestDataWithSeedDataForDatabaseConnection(serviceExecutionParameters: {
+    query: RawLambda;
+    mapping: Mapping;
+    runtime: Runtime;
+  }): GeneratorFn<void> {
+    try {
+      this.generatingTestDataWithSeedDataState.inProgress();
+      // NOTE: since we don't have a generic mechanism for test-data generation
+      // we will only report metrics around API usage, when we genericize, we will
+      // move this out
+      LegendStudioTelemetryHelper.logEvent_TestDataGenerationLaunched(
+        this.testDataState.editorStore.applicationStore.telemetryService,
+      );
+      const report = reportGraphAnalytics(
+        this.editorStore.graphManagerState.graph,
+      );
+      const tableRowIdentifiers = this.tableRowIdentifierStates.map(
+        (tableRowIdentifierState) => {
+          const rowIdentifiersArray =
+            tableRowIdentifierState.rowIdentifierStates.map(
+              (rowIdentifierState) =>
+                new RowIdentifier([
+                  new ColumnValuePair(
+                    rowIdentifierState.column.name,
+                    (rowIdentifierState.value as PrimitiveInstanceValue)
+                      .values[0] as object, // to do : change once engine code is in place
+                  ),
+                ]),
+            );
+          const tablePtr = new TablePtr();
+          tablePtr.database = tableRowIdentifierState.table.schema._OWNER.path;
+          tablePtr.mainTableDb = tablePtr.database;
+          tablePtr.schema = tableRowIdentifierState.table.schema.name;
+          tablePtr.table = tableRowIdentifierState.table.name;
+          return new TableRowIdentifiers(tablePtr, rowIdentifiersArray);
+        },
+      );
+      const value =
+        (yield this.editorStore.graphManagerState.graphManager.generateExecuteTestDataWithSeedData(
+          getExecutionQueryFromRawLambda(
+            serviceExecutionParameters.query,
+            this.parametersState.parameterStates,
+            this.editorStore.graphManagerState,
+          ),
+          tableRowIdentifiers,
+          serviceExecutionParameters.mapping,
+          serviceExecutionParameters.runtime,
+          this.editorStore.graphManagerState.graph,
+          {
+            anonymizeGeneratedData: this.anonymizeGeneratedData,
+            parameterValues: buildExecutionParameterValues(
+              this.parametersState.parameterStates,
+              this.editorStore.graphManagerState,
+            ),
+          },
+          report,
+        )) as string;
+
+      // NOTE: since we don't have a generic mechanism for test-data generation
+      // we will only report metrics around API usage, when we genericize, we will
+      // move this out
+      LegendStudioTelemetryHelper.logEvent_TestDataGenerationSucceeded(
+        this.editorStore.applicationStore.telemetryService,
+        report,
+      );
+      service_setConnectionTestDataEmbeddedData(
+        this.connectionData,
+        TEMPORARY__createRelationalDataFromCSV(value),
+        this.editorStore.changeDetectionState.observerContext,
+      );
+      this.embeddedEditorState = new EmbeddedDataEditorState(
+        this.testDataState.editorStore,
+        this.connectionData.testData,
+      );
+      this.generatingTestDataWithSeedDataState.pass();
+    } catch (error) {
+      assertErrorThrown(error);
+      this.editorStore.applicationStore.notificationService.notifyError(
+        `Unable to generate test data: ${error.message}`,
+      );
+      this.generatingTestDataWithSeedDataState.fail();
+    } finally {
+      this.generatingTestDataWithSeedDataState.complete();
+    }
+  }
+
+  *generateTestDataWithSeedData(): GeneratorFn<void> {
+    try {
+      this.generatingTestDataWithSeedDataState.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,
+        );
+        const parameters = (serviceExecutionParameters.query.parameters ??
+          []) as object[];
+        if (parameters.length > 0) {
+          this.parametersState.openModal(serviceExecutionParameters, {
+            generateWithSeedData: true,
+          });
+          return;
+        } else {
+          yield flowResult(
+            this.generateTestDataWithSeedDataForDatabaseConnection(
+              serviceExecutionParameters,
+            ),
+          );
+        }
+      } else {
+        // TODO: delete this once the backend code is in place
+        service_setConnectionTestDataEmbeddedData(
+          this.connectionData,
+          connection.accept_ConnectionVisitor(
+            new TEMPORARY__EmbeddedDataConnectionVisitor(this.editorStore),
+          ),
+          this.editorStore.changeDetectionState.observerContext,
+        );
+      }
+      this.embeddedEditorState = new EmbeddedDataEditorState(
+        this.testDataState.editorStore,
+        this.connectionData.testData,
+      );
+      this.generatingTestDataWithSeedDataState.pass();
+    } catch (error) {
+      assertErrorThrown(error);
+      this.editorStore.applicationStore.notificationService.notifyError(
+        `Unable to generate test data with seed data: ${error.message}`,
+      );
+      this.generatingTestDataWithSeedDataState.fail();
+    } finally {
+      this.generatingTestDataWithSeedDataState.complete();
+    }
+  }
+
   *generateQuerySchemas(): GeneratorFn<void> {
     try {
       this.generateSchemaQueryState.inProgress();