@finos/legend-query-builder 4.17.85 → 4.17.87

package/src/stores/workflows/dataProduct/DataProductQueryBuilderState.ts

@@ -20,12 +20,26 @@ import {
   DataProduct,
   type V1_DataProductArtifact,
   ModelAccessPointGroup,
-  type NativeModelExecutionContext,
+  NativeModelExecutionContext,
   CORE_PURE_PATH,
   type NativeModelAccess,
+  type Mapping,
+  type DataProductElementScope,
+  PackageableRuntime,
   resolveUsableDataProductClasses,
+  LakehouseRuntime,
+  type LambdaFunction,
+  SimpleFunctionExpression,
+  InstanceValue,
+  Multiplicity,
+  extractElementNameFromPath,
+  SUPPORTED_FUNCTIONS,
+  PackageableElementExplicitReference,
+  resolveDataProductExecutionState,
+  RuntimePointer,
 } from '@finos/legend-graph';
 import { QueryBuilderState } from '../../QueryBuilderState.js';
+
 import { type GenericLegendApplicationStore } from '@finos/legend-application';
 import type {
   QueryBuilderActionConfig,
@@ -35,7 +49,9 @@ import type { QueryBuilderConfig } from '../../../graph-manager/QueryBuilderConf
 import { renderDataProductQueryBuilderSetupPanelContent } from '../../../components/workflows/DataProductQueryBuilder.js';
 import {
   ActionState,
+  assertErrorThrown,
   filterByType,
+  guaranteeNonNullable,
   type GeneratorFn,
 } from '@finos/legend-shared';
 import { action, computed, flow, makeObservable, observable } from 'mobx';
@@ -44,6 +60,7 @@ import {
   type QueryableSourceInfo,
 } from '@finos/legend-storage';
 import { compareLabelFn } from '@finos/legend-art';
+import { QueryBuilderEmbeddedFromExecutionContextState } from '../../QueryBuilderExecutionContextState.js';
 
 export type DataProductOption = {
   label: string;
@@ -69,6 +86,113 @@ export const buildExecOptions = (
   };
 };
 
+export type ModelAccessPointGroupOption = {
+  label: string;
+  value: ModelAccessPointGroup;
+};
+
+export const buildModelAccessPointGroupOption = (
+  value: ModelAccessPointGroup,
+): ModelAccessPointGroupOption => ({
+  label: value.title ?? value.id,
+  value,
+});
+
+export abstract class DataProductExecutionState<T> {
+  readonly queryBuilderState: DataProductQueryBuilderState;
+  exectionValue: T;
+
+  constructor(
+    executionState: T,
+    queryBuilderState: DataProductQueryBuilderState,
+  ) {
+    this.queryBuilderState = queryBuilderState;
+    this.exectionValue = executionState;
+  }
+
+  abstract get label(): string;
+
+  abstract get mapping(): Mapping;
+
+  abstract get featuredElements(): DataProductElementScope[] | undefined;
+
+  get selectedOption(): { label: string; value: T } {
+    return {
+      label: this.label,
+      value: this.exectionValue,
+    };
+  }
+}
+
+export class NativeModelDataProductExecutionState extends DataProductExecutionState<NativeModelExecutionContext> {
+  nativeModelAccess: NativeModelAccess;
+
+  constructor(
+    executionState: NativeModelExecutionContext,
+    queryBuilderState: DataProductQueryBuilderState,
+  ) {
+    super(executionState, queryBuilderState);
+    this.nativeModelAccess = executionState.__owner;
+  }
+
+  override get label(): string {
+    return this.exectionValue.key;
+  }
+
+  get mapping(): Mapping {
+    return this.exectionValue.mapping.value;
+  }
+
+  get featuredElements(): DataProductElementScope[] | undefined {
+    return this.nativeModelAccess.featuredElements;
+  }
+}
+
+export class ModelAccessPointDataProductExecutionState extends DataProductExecutionState<ModelAccessPointGroup> {
+  selectedRuntime: PackageableRuntime | undefined;
+
+  constructor(
+    executionState: ModelAccessPointGroup,
+    queryBuilderState: DataProductQueryBuilderState,
+  ) {
+    super(executionState, queryBuilderState);
+    makeObservable(this, {
+      selectedRuntime: observable,
+      compatibleRuntimes: computed,
+      showRuntimeOptions: computed,
+      changeSelectedRuntime: action,
+    });
+    this.selectedRuntime = this.compatibleRuntimes[0];
+  }
+
+  changeSelectedRuntime(val: PackageableRuntime): void {
+    this.selectedRuntime = val;
+    this.queryBuilderState.changeRuntime(val);
+  }
+
+  override get label(): string {
+    return this.exectionValue.title ?? this.exectionValue.id;
+  }
+
+  get mapping(): Mapping {
+    return this.exectionValue.mapping.value;
+  }
+
+  get featuredElements(): DataProductElementScope[] | undefined {
+    return this.exectionValue.featuredElements;
+  }
+
+  get showRuntimeOptions(): boolean {
+    return this.compatibleRuntimes.length > 1;
+  }
+
+  get compatibleRuntimes(): PackageableRuntime[] {
+    return this.queryBuilderState.graphManagerState.usableRuntimes.filter(
+      (runtime) => runtime.runtimeValue instanceof LakehouseRuntime,
+    );
+  }
+}
+
 export class DataProductQueryBuilderState extends QueryBuilderState {
   readonly onClassChange?: ((val: Class) => void) | undefined;
   readonly onDataProductChange?: (val: DepotEntityWithOrigin) => Promise<void>;
@@ -77,10 +201,11 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
     | undefined;
 
   loadDataProductModelState = ActionState.create();
-  nativeModelAccess: NativeModelAccess;
   dataProduct: DataProduct;
   dataProductArtifact: V1_DataProductArtifact | undefined;
-  selectedExecContext: NativeModelExecutionContext;
+  executionState: DataProductExecutionState<
+    NativeModelExecutionContext | ModelAccessPointGroup
+  >;
   entities: DepotEntityWithOrigin[] | undefined;
 
   prioritizeEntityFunc?: ((val: DepotEntityWithOrigin) => boolean) | undefined;
@@ -95,8 +220,7 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
     dataProduct: DataProduct,
     artifact: V1_DataProductArtifact | undefined,
     actionConfig: QueryBuilderActionConfig,
-    nativeNativeModelAccess: NativeModelAccess,
-    nativeModelExecContext: NativeModelExecutionContext,
+    executionState: NativeModelExecutionContext | ModelAccessPointGroup,
     prioritizeEntityFunc: ((val: DepotEntityWithOrigin) => boolean) | undefined,
     onDataProductChange: (val: DepotEntityWithOrigin) => Promise<void>,
     onExecutionContextChange?:
@@ -108,19 +232,31 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
   ) {
     super(applicationStore, graphManagerState, workflow, config, sourceInfo);
     makeObservable(this, {
-      selectedExecContext: observable,
       dataProduct: observable,
-      setExecOptions: action,
+      executionState: observable,
+      initWithDataProduct: action,
+      handleDataProductChange: action,
+      setExecutionState: action,
       selectedDataProductOption: computed,
       isProductLinkable: computed,
+      modelAccessPointGroups: computed,
+      hasModelAccessPointGroups: computed,
+      modelAccessPointGroupOptions: computed,
+      activeMapping: computed,
+      activeFeaturedElements: computed,
       loadEntities: flow,
       entities: observable,
     });
+    // for queries against data product, we always want to default to from
+    this.executionContextState =
+      new QueryBuilderEmbeddedFromExecutionContextState(this);
     this.workflowState.updateActionConfig(actionConfig);
     this.dataProduct = dataProduct;
     this.dataProductArtifact = artifact;
-    this.nativeModelAccess = nativeNativeModelAccess;
-    this.selectedExecContext = nativeModelExecContext;
+    this.executionState =
+      executionState instanceof NativeModelExecutionContext
+        ? new NativeModelDataProductExecutionState(executionState, this)
+        : new ModelAccessPointDataProductExecutionState(executionState, this);
     this.prioritizeEntityFunc = prioritizeEntityFunc;
     this.onDataProductChange = onDataProductChange;
     this.onExecutionContextChange = onExecutionContextChange;
@@ -131,6 +267,63 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
     return false;
   }
 
+  handleDataProductChange(val: DepotEntityWithOrigin): void {
+    try {
+      this.loadDataProductModelState.inProgress();
+      const dataProduct = this.graphManagerState.graph.getDataProduct(val.path);
+      this.initWithDataProduct(dataProduct);
+    } catch (error) {
+      assertErrorThrown(error);
+      this.applicationStore.notificationService.notifyError(
+        `Failed to change Data Product: ${error.message}`,
+      );
+    }
+  }
+
+  initWithDataProduct(dataProduct: DataProduct): void {
+    try {
+      const execValue = resolveDataProductExecutionState(dataProduct);
+      this.dataProduct = dataProduct;
+      this.executionState =
+        execValue instanceof NativeModelExecutionContext
+          ? new NativeModelDataProductExecutionState(execValue, this)
+          : new ModelAccessPointDataProductExecutionState(execValue, this);
+      this.changeMapping(this.executionState.mapping);
+      if (this.executionState instanceof NativeModelDataProductExecutionState) {
+        const runtime = guaranteeNonNullable(
+          this.executionState.exectionValue.runtime,
+          'runtime unable to be resolved',
+        );
+        this.changeRuntime(new RuntimePointer(runtime));
+      } else if (
+        this.executionState instanceof
+          ModelAccessPointDataProductExecutionState &&
+        this.executionState.selectedRuntime instanceof PackageableRuntime
+      ) {
+        this.changeRuntime(this.executionState.selectedRuntime);
+      }
+      const compatibleClasses = resolveUsableDataProductClasses(
+        this.activeFeaturedElements,
+        this.executionState.mapping,
+        this.graphManagerState,
+        undefined,
+      );
+      // if there is no chosen class or the chosen one is not compatible
+      // with the mapping then pick a compatible class if possible
+      if (!this.class || !compatibleClasses.includes(this.class)) {
+        const possibleNewClass = compatibleClasses[0];
+        if (possibleNewClass) {
+          this.changeClass(possibleNewClass);
+        }
+      }
+    } catch (error) {
+      assertErrorThrown(error);
+      this.applicationStore.notificationService.notifyError(
+        `Failed to load Data Product: ${error.message}`,
+      );
+    }
+  }
+
   copyDataProductLinkToClipBoard(): void {
     if (!this.isProductLinkable) {
       this.applicationStore.notificationService.notifyError(
@@ -143,8 +336,77 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
     return DataProduct;
   }
 
-  setExecOptions(exec: NativeModelExecutionContext): void {
-    this.selectedExecContext = exec;
+  setExecutionState(
+    val: NativeModelExecutionContext | ModelAccessPointGroup,
+  ): void {
+    this.executionState =
+      val instanceof NativeModelExecutionContext
+        ? new NativeModelDataProductExecutionState(val, this)
+        : new ModelAccessPointDataProductExecutionState(val, this);
+  }
+
+  get modelAccessPointGroups(): ModelAccessPointGroup[] {
+    return this.dataProduct.accessPointGroups.filter(
+      filterByType(ModelAccessPointGroup),
+    );
+  }
+
+  get hasModelAccessPointGroups(): boolean {
+    return this.modelAccessPointGroups.length > 0;
+  }
+
+  get modelAccessPointGroupOptions(): ModelAccessPointGroupOption[] {
+    return this.modelAccessPointGroups
+      .map(buildModelAccessPointGroupOption)
+      .sort(compareLabelFn);
+  }
+
+  get activeMapping(): Mapping | undefined {
+    return this.executionState.mapping;
+  }
+
+  get activeFeaturedElements(): DataProductElementScope[] | undefined {
+    return this.executionState.featuredElements;
+  }
+
+  override buildExecutionContextExpression(
+    lambdaFunction: LambdaFunction,
+  ): LambdaFunction {
+    if (
+      this.executionState instanceof ModelAccessPointDataProductExecutionState
+    ) {
+      // for model access point group with attached ->with(dataProduct)->from(runtime) to execute query
+      let precedingExpression = guaranteeNonNullable(
+        lambdaFunction.expressionSequence[0],
+        `Can't build from() expression: preceding expression is not defined`,
+      );
+      const withFunc = new SimpleFunctionExpression(
+        extractElementNameFromPath(SUPPORTED_FUNCTIONS.WITH),
+      );
+      const dataProductInstance = new InstanceValue(
+        Multiplicity.ONE,
+        undefined,
+      );
+      dataProductInstance.values = [
+        PackageableElementExplicitReference.create(this.dataProduct),
+      ];
+      withFunc.parametersValues = [precedingExpression, dataProductInstance];
+      precedingExpression = withFunc;
+      const fromFunc = new SimpleFunctionExpression(
+        extractElementNameFromPath(SUPPORTED_FUNCTIONS.FROM),
+      );
+      const runtime = this.executionState.selectedRuntime;
+      if (runtime) {
+        const runtimeInstance = new InstanceValue(Multiplicity.ONE, undefined);
+        runtimeInstance.values = [
+          PackageableElementExplicitReference.create(runtime),
+        ];
+        fromFunc.parametersValues = [precedingExpression, runtimeInstance];
+        lambdaFunction.expressionSequence[0] = fromFunc;
+      }
+      return lambdaFunction;
+    }
+    return super.buildExecutionContextExpression(lambdaFunction);
   }
 
   *loadEntities(): GeneratorFn<void> {
@@ -216,14 +478,14 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
     requireReBuildingGraph?: boolean | undefined,
   ): Promise<void> {
     const currentMapping = this.executionContextState.mapping;
-    const execMapping = this.selectedExecContext.mapping.value;
-    if (execMapping !== currentMapping) {
-      this.changeMapping(execMapping, {
+    const newMapping = this.activeMapping;
+    if (newMapping && newMapping !== currentMapping) {
+      this.changeMapping(newMapping, {
         keepQueryContent: true,
       });
       const classes = resolveUsableDataProductClasses(
-        this.nativeModelAccess.featuredElements,
-        this.selectedExecContext.mapping.value,
+        this.activeFeaturedElements,
+        newMapping,
         this.graphManagerState,
         undefined,
       );