@finos/legend-query-builder 4.18.1 → 4.18.3

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

@@ -32,6 +32,7 @@ import {
   SimpleFunctionExpression,
   InstanceValue,
   Multiplicity,
+  type MappingModelCoverageAnalysisResult,
   extractElementNameFromPath,
   SUPPORTED_FUNCTIONS,
   PackageableElementExplicitReference,
@@ -53,6 +54,10 @@ import {
   GenericTypeExplicitReference,
   findLakehouseAccessPointGroup,
   type PureModel,
+  V1_RelationType,
+  V1_getGenericTypeFullPath,
+  LegendSDLC,
+  DataProductAccessType,
 } from '@finos/legend-graph';
 import { QueryBuilderState } from '../../QueryBuilderState.js';
 
@@ -74,27 +79,64 @@ import { action, computed, flow, makeObservable, observable } from 'mobx';
 import {
   DepotEntityWithOrigin,
   type QueryableSourceInfo,
+  type ProjectGAVCoordinates,
 } from '@finos/legend-storage';
 import { compareLabelFn } from '@finos/legend-art';
 import { QueryBuilderEmbeddedFromExecutionContextState } from '../../QueryBuilderExecutionContextState.js';
 import { buildLambdaFunction } from '../../QueryBuilderValueSpecificationBuilder.js';
 
-export const buildDataProductAccessor = (
-  relationMetadata: RelationTypeMetadata,
+export const resolveDataProductAccessor = (
   dataProduct: DataProduct,
   accessPoint: LakehouseAccessPoint,
   graph: PureModel,
+  artifact: V1_DataProductArtifact | undefined,
+  relationMetadata?: RelationTypeMetadata | undefined,
 ): DataProductAccessor => {
   const relationType = new RelationType(accessPoint.title ?? accessPoint.id);
-  relationType.columns = relationMetadata.columns.map(
-    (col) =>
-      new RelationColumn(
-        col.name,
-        GenericTypeExplicitReference.create(
-          new GenericType(graph.getType(col.type)),
-        ),
+  if (artifact) {
+    const artifactGroup = artifact.accessPointGroups.find((apg) =>
+      apg.accessPointImplementations.some(
+        (apImpl) => apImpl.id === accessPoint.id,
       ),
-  );
+    );
+    const artifactImpl = artifactGroup?.accessPointImplementations.find(
+      (apImpl) => apImpl.id === accessPoint.id,
+    );
+    const v1RelationType = artifactImpl?.lambdaGenericType?.typeArguments
+      .map((typeArg) => typeArg.rawType)
+      .find((rawType) => rawType instanceof V1_RelationType);
+    if (v1RelationType) {
+      relationType.columns = v1RelationType.columns.map(
+        (col) =>
+          new RelationColumn(
+            col.name,
+            GenericTypeExplicitReference.create(
+              new GenericType(
+                graph.getType(V1_getGenericTypeFullPath(col.genericType)),
+              ),
+            ),
+          ),
+      );
+      return new DataProductAccessor(
+        dataProduct.path,
+        artifactGroup?.id,
+        accessPoint.id,
+        relationType,
+        dataProduct,
+      );
+    }
+  }
+  if (relationMetadata) {
+    relationType.columns = relationMetadata.columns.map(
+      (col) =>
+        new RelationColumn(
+          col.name,
+          GenericTypeExplicitReference.create(
+            new GenericType(graph.getType(col.type)),
+          ),
+        ),
+    );
+  }
   const groupResult = findLakehouseAccessPointGroup(
     dataProduct,
     accessPoint.id,
@@ -268,7 +310,9 @@ export class LakehouseDataProductExecutionState extends DataProductExecutionStat
     queryBuilderState: DataProductQueryBuilderState,
   ) {
     super(executionState, queryBuilderState);
-    makeObservable(this, {});
+    makeObservable(this, {
+      withAdhocRuntime: observable,
+    });
     this.selectedRuntime = this.compatibleRuntimes[0];
   }
 
@@ -277,6 +321,11 @@ export class LakehouseDataProductExecutionState extends DataProductExecutionStat
     this.queryBuilderState.changeRuntime(val);
   }
 
+  withAdhocRuntime(): LakehouseDataProductExecutionState {
+    this.adhocRuntime = true;
+    return this;
+  }
+
   override get label(): string {
     return this.exectionValue.title ?? this.exectionValue.id;
   }
@@ -303,9 +352,6 @@ export class LakehouseDataProductExecutionState extends DataProductExecutionStat
 export class DataProductQueryBuilderState extends QueryBuilderState {
   readonly onClassChange?: ((val: Class) => void) | undefined;
   readonly onDataProductChange?: (val: DepotEntityWithOrigin) => Promise<void>;
-  readonly onExecutionContextChange?:
-    | ((val: NativeModelExecutionContext) => void)
-    | undefined;
 
   loadDataProductModelState = ActionState.create();
   dataProduct: DataProduct;
@@ -314,6 +360,10 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
     NativeModelExecutionContext | ModelAccessPointGroup | LakehouseAccessPoint
   >;
   entities: DepotEntityWithOrigin[] | undefined;
+  mappingToMappingCoverageResult?: Map<
+    string,
+    MappingModelCoverageAnalysisResult
+  >;
 
   prioritizeEntityFunc?: ((val: DepotEntityWithOrigin) => boolean) | undefined;
 
@@ -333,9 +383,6 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
       | LakehouseAccessPoint,
     prioritizeEntityFunc: ((val: DepotEntityWithOrigin) => boolean) | undefined,
     onDataProductChange: (val: DepotEntityWithOrigin) => Promise<void>,
-    onExecutionContextChange?:
-      | ((val: NativeModelExecutionContext) => void)
-      | undefined,
     onClassChange?: ((val: Class) => void) | undefined,
     config?: QueryBuilderConfig | undefined,
     sourceInfo?: QueryableSourceInfo | undefined,
@@ -357,6 +404,11 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
       showExecutionIdSelector: computed,
       executionIdOptions: computed,
       selectedExecutionIdOption: computed,
+      hasNativeModelAccess: computed,
+      hasBothAccessModes: computed,
+      showContextSelector: computed,
+      allContextOptions: computed,
+      selectedContextOption: computed,
       selectedExecOption: computed,
       selectedModelAccessPointGroupOption: computed,
       usableClasses: computed,
@@ -382,13 +434,44 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
           : new ModelAccessPointDataProductExecutionState(executionState, this);
     this.prioritizeEntityFunc = prioritizeEntityFunc;
     this.onDataProductChange = onDataProductChange;
-    this.onExecutionContextChange = onExecutionContextChange;
     this.onClassChange = onClassChange;
     // force from.
     this.executionContextState =
       new QueryBuilderEmbeddedFromExecutionContextState(this);
   }
 
+  async prepareAccessForExecution(): Promise<void> {
+    if (this.executionState instanceof NativeModelDataProductExecutionState) {
+      const runtime = this.executionState.exectionValue.runtime;
+      if (runtime) {
+        this.changeRuntime(new RuntimePointer(runtime));
+      }
+    } else if (
+      this.executionState instanceof
+        ModelAccessPointDataProductExecutionState &&
+      this.executionState.selectedRuntime
+    ) {
+      this.changeRuntime(
+        new RuntimePointer(
+          PackageableElementExplicitReference.create(
+            this.executionState.selectedRuntime,
+          ),
+        ),
+      );
+    } else if (
+      this.executionState instanceof LakehouseDataProductExecutionState &&
+      this.executionState.selectedRuntime
+    ) {
+      this.changeRuntime(
+        new RuntimePointer(
+          PackageableElementExplicitReference.create(
+            this.executionState.selectedRuntime,
+          ),
+        ),
+      );
+    }
+  }
+
   get isProductLinkable(): boolean {
     return false;
   }
@@ -417,19 +500,40 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
       this.modelAccessPointGroupOptions.length > 1
     );
   }
+  get hasNativeModelAccess(): boolean {
+    return this.dataProduct.nativeModelAccess !== undefined;
+  }
+
+  get hasBothAccessModes(): boolean {
+    return this.hasModelAccessPointGroups && this.hasNativeModelAccess;
+  }
+
+  // showContextSelector / allContextOptions / selectedContextOption are aliases
+  // over the upstream ExecutionIdOption API so existing UI code keeps working.
+  get showContextSelector(): boolean {
+    return this.showExecutionIdSelector;
+  }
+
+  get allContextOptions(): ExecutionIdOption[] {
+    return this.executionIdOptions;
+  }
+
+  get selectedContextOption(): ExecutionIdOption | undefined {
+    return this.selectedExecutionIdOption;
+  }
 
   get executionIdOptions(): ExecutionIdOption[] {
     const nativeOptions: ExecutionIdOption[] = (
       this.dataProduct.nativeModelAccess?.nativeModelExecutionContexts ?? []
     ).map((ctx) => ({
       label: ctx.key,
-      tag: 'Native',
+      tag: 'NATIVE',
       value: ctx,
     }));
     const modelOptions: ExecutionIdOption[] = this.modelAccessPointGroups.map(
       (group) => ({
         label: group.title ?? group.id,
-        tag: 'Model',
+        tag: 'MODEL',
         value: group,
       }),
     );
@@ -439,7 +543,7 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
         .filter(filterByType(LakehouseAccessPoint))
         .map((ap) => ({
           label: ap.title ?? ap.id,
-          tag: 'Lakehouse',
+          tag: 'LAKEHOUSE',
           value: ap,
         }));
     return [...modelOptions, ...lakehouseOptions, ...nativeOptions].sort(
@@ -452,19 +556,19 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
     if (state instanceof NativeModelDataProductExecutionState) {
       return {
         label: state.exectionValue.key,
-        tag: 'Native',
+        tag: 'NATIVE',
         value: state.exectionValue,
       };
     } else if (state instanceof ModelAccessPointDataProductExecutionState) {
       return {
         label: state.exectionValue.title ?? state.exectionValue.id,
-        tag: 'Model',
+        tag: 'MODEL',
         value: state.exectionValue,
       };
     } else if (state instanceof LakehouseDataProductExecutionState) {
       return {
         label: state.exectionValue.title ?? state.exectionValue.id,
-        tag: 'Lakehouse',
+        tag: 'LAKEHOUSE',
         value: state.exectionValue,
       };
     }
@@ -480,11 +584,20 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
     if (val === this.executionState.exectionValue) {
       return;
     }
+    const switchingToModel = val instanceof ModelAccessPointGroup;
+    const switchingToNative = val instanceof NativeModelExecutionContext;
+    const wasModeSwitch =
+      (switchingToModel && this.isNativeMode) ||
+      (switchingToNative && this.isModelAccessPointGroupMode);
+
+    if (wasModeSwitch) {
+      this.changeHistoryState.querySnapshotBuffer = [];
+      this.changeHistoryState.pointer = -1;
+      this.changeHistoryState.setCurrentQuery(undefined);
+    }
+
     await this.changeExecutionState(val);
     await this.propagateExecutionContextChange();
-    if (val instanceof NativeModelExecutionContext) {
-      this.onExecutionContextChange?.(val);
-    }
   }
 
   get selectedExecOption():
@@ -523,29 +636,6 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
       : [];
   }
 
-  changeNativeExecutionContext(val: NativeModelExecutionContext): void {
-    if (this.isNativeMode && val === this.executionState.exectionValue) {
-      return;
-    }
-    this.setExecutionState(val);
-    this.propagateExecutionContextChange()
-      .then(() => this.onExecutionContextChange?.(val))
-      .catch(this.applicationStore.alertUnhandledError);
-  }
-
-  changeModelAccessPointGroupValue(val: ModelAccessPointGroup): void {
-    if (
-      this.isModelAccessPointGroupMode &&
-      val === this.executionState.exectionValue
-    ) {
-      return;
-    }
-    this.setExecutionState(val);
-    this.propagateExecutionContextChange().catch(
-      this.applicationStore.alertUnhandledError,
-    );
-  }
-
   override buildQueryForPersistence(): RawLambda {
     if (!this.isQuerySupported) {
       return this.buildQuery();
@@ -580,7 +670,9 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
     ) {
       const execContext = new QueryDataProductLakehouseExecutionContext();
       execContext.dataProductPath = this.dataProduct.path;
+      execContext.accessGroupId = this.executionState.exectionValue.__owner.id;
       execContext.accessPointId = this.executionState.exectionValue.id;
+
       return execContext;
     }
     return super.getQueryExecutionContext();
@@ -632,12 +724,7 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
       const execValue =
         preResolvedState ?? resolveDataProductExecutionState(dataProduct);
       this.dataProduct = dataProduct;
-      this.executionState =
-        execValue instanceof NativeModelExecutionContext
-          ? new NativeModelDataProductExecutionState(execValue, this)
-          : execValue instanceof LakehouseAccessPoint
-            ? new LakehouseDataProductExecutionState(execValue, this)
-            : new ModelAccessPointDataProductExecutionState(execValue, this);
+      this.setExecutionState(execValue);
       const mapping = this.executionState.mapping;
       if (mapping) {
         this.changeMapping(mapping);
@@ -669,11 +756,18 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
         );
       }
       if (mapping) {
+        const coverageResult = this.mappingToMappingCoverageResult?.get(
+          mapping.path,
+        );
+        if (coverageResult) {
+          this.explorerState.mappingModelCoverageAnalysisResult =
+            coverageResult;
+        }
         const compatibleClasses = resolveUsableDataProductClasses(
           this.activeFeaturedElements,
           mapping,
           this.graphManagerState,
-          undefined,
+          coverageResult,
         );
         // if there is no chosen class or the chosen one is not compatible
         // with the mapping then pick a compatible class if possible
@@ -729,16 +823,18 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
   ): Promise<void> {
     this.setExecutionState(val);
     if (val instanceof LakehouseAccessPoint) {
-      const relationMetadata =
-        await this.graphManagerState.graphManager.getLambdaRelationType(
-          val.func,
-          this.graphManagerState.graph,
-        );
-      const accessor = buildDataProductAccessor(
-        relationMetadata,
+      const relationMetadata = !this.dataProductArtifact
+        ? await this.graphManagerState.graphManager.getLambdaRelationType(
+            val.func,
+            this.graphManagerState.graph,
+          )
+        : undefined;
+      const accessor = resolveDataProductAccessor(
         this.dataProduct,
         val,
         this.graphManagerState.graph,
+        this.dataProductArtifact,
+        relationMetadata,
       );
       this.setSourceElement(accessor);
     }
@@ -882,19 +978,92 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
   ): Promise<void> {
     const currentMapping = this.executionContextState.mapping;
     const newMapping = this.activeMapping;
-    if (newMapping && newMapping !== currentMapping) {
-      this.changeMapping(newMapping, {
-        keepQueryContent: true,
-      });
-      const classes = resolveUsableDataProductClasses(
-        this.activeFeaturedElements,
-        newMapping,
-        this.graphManagerState,
-        undefined,
-      );
-      if (this.sourceClass && !classes.includes(this.sourceClass)) {
-        this.setSourceElement(classes[0]);
+    if (!newMapping || newMapping === currentMapping) {
+      return;
+    }
+
+    const coverageResult = this.mappingToMappingCoverageResult?.get(
+      newMapping.path,
+    );
+
+    if (coverageResult && this.dataProductArtifact) {
+      const origin = this.graphManagerState.graph.origin;
+      if (origin instanceof LegendSDLC) {
+        const newGraph = this.graphManagerState.createNewGraph();
+        const projectInfo: ProjectGAVCoordinates = {
+          groupId: origin.groupId,
+          artifactId: origin.artifactId,
+          versionId: origin.versionId,
+        };
+        let accessPointId: string;
+        let dataProductAccessType: DataProductAccessType;
+        if (
+          this.executionState instanceof NativeModelDataProductExecutionState
+        ) {
+          accessPointId = this.executionState.exectionValue.key;
+          dataProductAccessType = DataProductAccessType.NATIVE;
+        } else if (
+          this.executionState instanceof LakehouseDataProductExecutionState
+        ) {
+          accessPointId = this.executionState.exectionValue.id;
+          dataProductAccessType = DataProductAccessType.LAKEHOUSE;
+        } else {
+          accessPointId = (
+            this.executionState as ModelAccessPointDataProductExecutionState
+          ).exectionValue.id;
+          dataProductAccessType = DataProductAccessType.MODEL;
+        }
+        const analysisResult =
+          await this.graphManagerState.graphManager.buildDataProductAnalysis(
+            this.dataProductArtifact,
+            this.dataProduct.path,
+            newGraph,
+            accessPointId,
+            dataProductAccessType,
+            projectInfo,
+          );
+        this.graphManagerState.graph = newGraph;
+
+        this.dataProduct = newGraph.getDataProduct(this.dataProduct.path);
+        this.setExecutionState(analysisResult.targetExecState);
+        if (analysisResult.dataProductAnalysis.mappingToMappingCoverageResult) {
+          this.mappingToMappingCoverageResult =
+            analysisResult.dataProductAnalysis.mappingToMappingCoverageResult;
+        }
+        const newCoverageResult = this.mappingToMappingCoverageResult?.get(
+          newMapping.path,
+        );
+        if (newCoverageResult) {
+          this.explorerState.mappingModelCoverageAnalysisResult =
+            newCoverageResult;
+        }
+      } else {
+        this.explorerState.mappingModelCoverageAnalysisResult = coverageResult;
+      }
+    } else if (coverageResult) {
+      this.explorerState.mappingModelCoverageAnalysisResult = coverageResult;
+    }
+
+    await this.prepareAccessForExecution();
+
+    this.changeMapping(newMapping, {
+      keepQueryContent: true,
+    });
+    const classes = resolveUsableDataProductClasses(
+      this.activeFeaturedElements,
+      newMapping,
+      this.graphManagerState,
+      this.explorerState.mappingModelCoverageAnalysisResult,
+    );
+    if (
+      !this.sourceClass ||
+      (!classes.includes(this.sourceClass) && classes.length)
+    ) {
+      const possibleNewClass = classes[0];
+      if (possibleNewClass) {
+        this.changeSourceElement(possibleNewClass);
       }
     }
+    this.explorerState.refreshTreeData();
   }
 }