@finos/legend-query-builder 4.19.8 → 4.19.10

package/src/stores/QueryBuilderState.ts

@@ -304,6 +304,7 @@ export abstract class QueryBuilderState implements CommandRegistrar {
       setLambdaWriteMode: action,
       setINTERNAL__enableInitializingDefaultSimpleExpressionValue: action,
       TEMPORARY_initializeExecContext: action,
+      reconcileExecutionContextState: action,
 
       resetQueryResult: action,
       resetQueryContent: action,
@@ -355,11 +356,9 @@ export abstract class QueryBuilderState implements CommandRegistrar {
     isTypedTDS: boolean,
   ): QueryBuilderExecutionContextState {
     if (isTypedTDS) {
-      const context = new QueryBuilderEmbeddedFromExecutionContextState(this);
-      this.setLambdaWriteMode(
-        QUERY_BUILDER_LAMBDA_WRITER_MODE.TYPED_FETCH_STRUCTURE,
-      );
-      return context;
+      this.lambdaWriteMode =
+        QUERY_BUILDER_LAMBDA_WRITER_MODE.TYPED_FETCH_STRUCTURE;
+      return new QueryBuilderEmbeddedFromExecutionContextState(this);
     }
     return new QueryBuilderExternalExecutionContextState(this);
   }
@@ -421,6 +420,12 @@ export abstract class QueryBuilderState implements CommandRegistrar {
     );
   }
 
+  get requiresEmbeddedExecutionContext(): boolean {
+    return (
+      this.isFetchStructureTyped && Boolean(this.executionContextState.mapping)
+    );
+  }
+
   get forceFromExpressionForExec(): boolean {
     return this.isFetchStructureTyped;
   }
@@ -431,6 +436,29 @@ export abstract class QueryBuilderState implements CommandRegistrar {
 
   setLambdaWriteMode(val: QUERY_BUILDER_LAMBDA_WRITER_MODE): void {
     this.lambdaWriteMode = val;
+    this.reconcileExecutionContextState();
+  }
+  reconcileExecutionContextState(options?: {
+    allowDowngrade?: boolean | undefined;
+  }): void {
+    const requiresEmbedded = this.requiresEmbeddedExecutionContext;
+    const isEmbedded =
+      this.executionContextState instanceof
+      QueryBuilderEmbeddedFromExecutionContextState;
+    if (requiresEmbedded === isEmbedded) {
+      return;
+    }
+    if (isEmbedded && !options?.allowDowngrade) {
+      return;
+    }
+    const preservedMapping = this.executionContextState.mapping;
+    const preservedRuntime = this.executionContextState.runtimeValue;
+    const next = requiresEmbedded
+      ? new QueryBuilderEmbeddedFromExecutionContextState(this)
+      : new QueryBuilderExternalExecutionContextState(this);
+    next.setMapping(preservedMapping);
+    next.setRuntimeValue(preservedRuntime);
+    this.setExecutionContextState(next);
   }
 
   getQueryExecutionContext(): QueryExecutionContext {
@@ -794,6 +822,19 @@ export abstract class QueryBuilderState implements CommandRegistrar {
     return this.buildQuery();
   }
 
+  protected buildQueryLambdaWithoutExecutionContext(): RawLambda {
+    if (!this.isQuerySupported) {
+      return this.buildQuery();
+    }
+    return buildRawLambdaFromLambdaFunction(
+      buildLambdaFunction(this, {
+        skipExecutionContext: true,
+        useTypedRelationFunctions: this.isFetchStructureTyped,
+      }),
+      this.graphManagerState,
+    );
+  }
+
   buildFromQuery(): RawLambda {
     assertTrue(
       this.isQuerySupported,

package/src/stores/QueryBuilderStateBuilder.ts

@@ -895,6 +895,19 @@ export class QueryBuilderValueSpecificationProcessor
         this.queryBuilderState,
       );
       return;
+    } else if (matchFunctionName(functionName, [SUPPORTED_FUNCTIONS.WITH])) {
+      const parameters = valueSpecification.parametersValues;
+      assertTrue(
+        parameters.length === 2,
+        'With function expects a body expression and a data product',
+      );
+      QueryBuilderValueSpecificationProcessor.processChild(
+        guaranteeNonNullable(parameters[0]),
+        valueSpecification,
+        this.parentLambda,
+        this.queryBuilderState,
+      );
+      return;
     } else if (
       matchFunctionName(
         functionName,

package/src/stores/data-cube/QueryBuilderDataCubeHelper.ts

@@ -37,6 +37,7 @@ export const createDataCubeViewerStateFromQueryBuilder = async (
   }
   const mappingPath = queryBuilderState.executionContextState.mapping?.path;
   const currentLambdaWriterMode = queryBuilderState.lambdaWriteMode;
+  const currentExecutionContextState = queryBuilderState.executionContextState;
   // ensure we write in new tds mode
   queryBuilderState.setLambdaWriteMode(
     QUERY_BUILDER_LAMBDA_WRITER_MODE.TYPED_FETCH_STRUCTURE,
@@ -55,6 +56,11 @@ export const createDataCubeViewerStateFromQueryBuilder = async (
     queryBuilderState,
   );
   queryBuilderState.setLambdaWriteMode(currentLambdaWriterMode);
+  if (
+    queryBuilderState.executionContextState !== currentExecutionContextState
+  ) {
+    queryBuilderState.setExecutionContextState(currentExecutionContextState);
+  }
   const query = await engine.generateInitialSpecification();
   return new QueryBuilderDataCubeViewerState(query, engine);
 };

package/src/stores/workflows/accessor/AccessorQueryBuilderState.ts

@@ -149,6 +149,10 @@ export class AccessorQueryBuilderState extends QueryBuilderState {
     return false;
   }
 
+  override get requiresEmbeddedExecutionContext(): boolean {
+    return true;
+  }
+
   /**
    * In accessor mode, configure both filter and post-filter panels to be visible.
    */

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

@@ -29,12 +29,8 @@ import {
   resolveUsableDataProductClasses,
   LakehouseRuntime,
   type LambdaFunction,
-  SimpleFunctionExpression,
-  InstanceValue,
-  Multiplicity,
+  attachWithFromQuery,
   type MappingModelCoverageAnalysisResult,
-  extractElementNameFromPath,
-  SUPPORTED_FUNCTIONS,
   PackageableElementExplicitReference,
   resolveDataProductExecutionState,
   RuntimePointer,
@@ -44,7 +40,6 @@ import {
   QueryDataProductLakehouseExecutionContext,
   LakehouseAccessPoint,
   type RawLambda,
-  buildRawLambdaFromLambdaFunction,
   type Accessor,
   type RelationTypeMetadata,
   DataProductAccessor,
@@ -82,7 +77,6 @@ import {
 } from '@finos/legend-storage';
 import { compareLabelFn } from '@finos/legend-art';
 import { QueryBuilderEmbeddedFromExecutionContextState } from '../../QueryBuilderExecutionContextState.js';
-import { buildLambdaFunction } from '../../QueryBuilderValueSpecificationBuilder.js';
 
 export const resolveDataProductAccessor = (
   dataProduct: DataProduct,
@@ -624,6 +618,10 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
     return true;
   }
 
+  override get requiresEmbeddedExecutionContext(): boolean {
+    return true;
+  }
+
   get selectedModelAccessPointGroupOption():
     | ModelAccessPointGroupOption
     | undefined {
@@ -646,19 +644,7 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
   }
 
   override buildQueryForPersistence(): RawLambda {
-    if (!this.isQuerySupported) {
-      return this.buildQuery();
-    }
-    // Build without embedding execution context in the lambda body.
-    // The execution context (dataProductPath, executionKey/accessPointGroupId)
-    // is stored separately in query.executionContext via getQueryExecutionContext().
-    return buildRawLambdaFromLambdaFunction(
-      buildLambdaFunction(this, {
-        skipExecutionContext: true,
-        useTypedRelationFunctions: this.isFetchStructureTyped,
-      }),
-      this.graphManagerState,
-    );
+    return this.buildQueryLambdaWithoutExecutionContext();
   }
 
   override getQueryExecutionContext(): QueryExecutionContext {
@@ -905,34 +891,9 @@ export class DataProductQueryBuilderState extends QueryBuilderState {
     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 attachWithFromQuery(lambdaFunction, this.dataProduct, runtime);
       }
       return lambdaFunction;
     }