@finos/legend-query-builder 4.17.100 → 4.18.0

package/src/stores/QueryBuilderState.ts

@@ -55,10 +55,12 @@ import {
   type PackageableElement,
   type ValueSpecification,
   type Type,
+  type RelationType,
   type QueryGridConfig,
   type QueryExecutionContext,
   type FunctionAnalysisInfo,
   type GraphData,
+  Accessor,
   GRAPH_MANAGER_EVENT,
   CompilationError,
   extractSourceInformationCoordinates,
@@ -226,7 +228,7 @@ export abstract class QueryBuilderState implements CommandRegistrar {
 
   lambdaWriteMode = QUERY_BUILDER_LAMBDA_WRITER_MODE.STANDARD;
 
-  class?: Class | undefined;
+  sourceElement?: Class | Accessor | undefined;
   getAllFunction: QUERY_BUILDER_SUPPORTED_GET_ALL_FUNCTIONS =
     QUERY_BUILDER_SUPPORTED_GET_ALL_FUNCTIONS.GET_ALL;
   executionContextState: QueryBuilderExecutionContextState;
@@ -270,7 +272,7 @@ export abstract class QueryBuilderState implements CommandRegistrar {
       changeDetectionState: observable,
       changeHistoryState: observable,
       executionContextState: observable,
-      class: observable,
+      sourceElement: observable,
       queryChatState: observable,
       isQueryChatOpened: observable,
       isLocalModeEnabled: observable,
@@ -280,9 +282,13 @@ export abstract class QueryBuilderState implements CommandRegistrar {
       INTERNAL__enableInitializingDefaultSimpleExpressionValue: observable,
 
       sideBarClassName: computed,
+      sourceClass: computed,
+      sourceAccessor: computed,
+      sourceRelationType: computed,
       isQuerySupported: computed,
       allValidationIssues: computed,
       canBuildQuery: computed,
+      useRelation: computed,
 
       setShowFunctionsExplorerPanel: action,
       setShowParametersPanel: action,
@@ -291,7 +297,7 @@ export abstract class QueryBuilderState implements CommandRegistrar {
       setDataCubeViewerState: action,
       openDataCubeEngine: action,
       setIsCheckingEntitlments: action,
-      setClass: action,
+      setSourceElement: action,
       setIsQueryChatOpened: action,
       setIsLocalModeEnabled: action,
       setGetAllFunction: action,
@@ -301,7 +307,7 @@ export abstract class QueryBuilderState implements CommandRegistrar {
 
       resetQueryResult: action,
       resetQueryContent: action,
-      changeClass: action,
+      changeSourceElement: action,
       changeMapping: action,
       changeRuntime: action,
       setExecutionContextState: action,
@@ -358,6 +364,10 @@ export abstract class QueryBuilderState implements CommandRegistrar {
     return new QueryBuilderExternalExecutionContextState(this);
   }
 
+  get useRelation(): boolean {
+    return this.sourceElement instanceof Accessor || this.isFetchStructureTyped;
+  }
+
   get isMappingReadOnly(): boolean {
     return false;
   }
@@ -415,6 +425,10 @@ export abstract class QueryBuilderState implements CommandRegistrar {
     return this.isFetchStructureTyped;
   }
 
+  get requiresMappingForExecution(): boolean {
+    return true;
+  }
+
   setLambdaWriteMode(val: QUERY_BUILDER_LAMBDA_WRITER_MODE): void {
     this.lambdaWriteMode = val;
   }
@@ -463,7 +477,7 @@ export abstract class QueryBuilderState implements CommandRegistrar {
    */
   getStateInfo(): QueryableClassMappingRuntimeInfo | undefined {
     if (this.sourceInfo) {
-      const classPath = this.class?.path;
+      const classPath = this.sourceClass?.path;
       const mappingPath = this.executionContextState.mapping?.path;
       const runtimePath =
         this.executionContextState.runtimeValue instanceof RuntimePointer
@@ -528,8 +542,29 @@ export abstract class QueryBuilderState implements CommandRegistrar {
     this.isCalendarEnabled = val;
   }
 
-  setClass(val: Class | undefined): void {
-    this.class = val;
+  /**
+   * Convenience getter that returns the class only when the source
+   * type is a Class (not a RelationType). Use this in code paths
+   * that are Class-specific (milestoning, explorer tree, graph-fetch, etc.).
+   */
+  get sourceClass(): Class | undefined {
+    return this.sourceElement instanceof Accessor
+      ? undefined
+      : this.sourceElement;
+  }
+
+  get sourceAccessor(): Accessor | undefined {
+    return this.sourceElement instanceof Accessor
+      ? this.sourceElement
+      : undefined;
+  }
+
+  get sourceRelationType(): RelationType | undefined {
+    return this.sourceAccessor?.relationType;
+  }
+
+  setSourceElement(val: Class | Accessor | undefined): void {
+    this.sourceElement = val;
   }
 
   setExecutionContextState(val: QueryBuilderExecutionContextState): void {
@@ -651,13 +686,13 @@ export abstract class QueryBuilderState implements CommandRegistrar {
     }
   }
 
-  changeClass(val: Class): void {
+  changeSourceElement(val: Class | Accessor): void {
     this.resetQueryResult();
     this.resetQueryContent();
     this.setGetAllFunction(QUERY_BUILDER_SUPPORTED_GET_ALL_FUNCTIONS.GET_ALL);
-    this.setClass(val);
+    this.setSourceElement(val);
     this.explorerState.refreshTreeData();
-    this.fetchStructureState.implementation.onClassChange(val);
+    this.fetchStructureState.implementation.onClassChange();
     this.milestoningState.updateMilestoningConfiguration();
     this.changeHistoryState.cacheNewQuery(this.buildQuery());
   }
@@ -878,7 +913,7 @@ export abstract class QueryBuilderState implements CommandRegistrar {
       this.resetQueryContent();
       this.unsupportedQueryState.setLambdaError(error);
       this.unsupportedQueryState.setRawLambda(query);
-      this.setClass(undefined);
+      this.setSourceElement(undefined);
       const parameters = buildLambdaVariableExpressions(
         query,
         this.graphManagerState,
@@ -1019,7 +1054,7 @@ export abstract class QueryBuilderState implements CommandRegistrar {
       this.workflowState,
       undefined,
     );
-    basicState.class = this.class;
+    basicState.sourceElement = this.sourceElement;
     basicState.executionContextState.mapping =
       this.executionContextState.mapping;
     basicState.executionContextState.runtimeValue =

package/src/stores/QueryBuilderStateBuilder.ts

@@ -27,8 +27,8 @@ import {
 import type { QueryBuilderState } from './QueryBuilderState.js';
 import {
   AbstractPropertyExpression,
+  FunctionExpression,
   type EnumValueInstanceValue,
-  type FunctionExpression,
   type GraphFetchTreeInstanceValue,
   type ValueSpecificationVisitor,
   type LambdaFunction,
@@ -55,6 +55,8 @@ import {
   PackageableElementExplicitReference,
   MILESTONING_STEREOTYPE,
   type ColSpecInstanceValue,
+  type AccessorInstanceValue,
+  RelationColumn,
 } from '@finos/legend-graph';
 import { processTDSPostFilterExpression } from './fetch-structure/tds/post-filter/QueryBuilderPostFilterStateBuilder.js';
 import { processFilterExpression } from './filter/QueryBuilderFilterStateBuilder.js';
@@ -76,6 +78,7 @@ import {
   processTDSProjectExpression,
   processTDSProjectionColumnPropertyExpression,
   processTDSProjectionDerivationExpression,
+  processTDSRelationColumn,
   processTDSSliceExpression,
   processTDSSortDirectionExpression,
   processTDSSortExpression,
@@ -87,7 +90,10 @@ import {
   QUERY_BUILDER_SUPPORTED_GET_ALL_FUNCTIONS,
 } from '../graph/QueryBuilderMetaModelConst.js';
 import { LambdaParameterState } from './shared/LambdaParameterState.js';
-import { processTDS_OLAPGroupByExpression } from './fetch-structure/tds/window/QueryBuilderWindowStateBuilder.js';
+import {
+  processTDS_ExtendExpression,
+  processTDS_OLAPGroupByExpression,
+} from './fetch-structure/tds/window/QueryBuilderWindowStateBuilder.js';
 import { processWatermarkExpression } from './watermark/QueryBuilderWatermarkStateBuilder.js';
 import {
   type QueryBuilderConstantExpressionState,
@@ -117,7 +123,7 @@ const processGetAllExpression = (
     Class,
     `Can't process getAll() expression: getAll() return type is missing`,
   );
-  queryBuilderState.setClass(_class);
+  queryBuilderState.setSourceElement(_class);
   queryBuilderState.milestoningState.clearMilestoningDates();
   queryBuilderState.explorerState.refreshTreeData();
 
@@ -149,7 +155,7 @@ const processGetAllVersionsExpression = (
     Class,
     `Can't process getAllVersions() expression: getAllVersions() return type is missing`,
   );
-  queryBuilderState.setClass(_class);
+  queryBuilderState.setSourceElement(_class);
   queryBuilderState.milestoningState.clearMilestoningDates();
   queryBuilderState.explorerState.refreshTreeData();
 
@@ -183,7 +189,7 @@ const processGetAllVersionsInRangeExpression = (
     Class,
     `Can't process getAllVersionsInRange() expression: getAllVersionsInRange() return type is missing`,
   );
-  queryBuilderState.setClass(_class);
+  queryBuilderState.setSourceElement(_class);
   queryBuilderState.milestoningState.clearMilestoningDates();
   queryBuilderState.explorerState.refreshTreeData();
 
@@ -309,24 +315,34 @@ const processFromFunction = (
   expression: SimpleFunctionExpression,
   queryBuilderState: QueryBuilderState,
 ): void => {
-  // mapping
-  const mappingInstanceExpression = guaranteeType(
-    expression.parametersValues[1],
-    InstanceValue,
-    `Can't process from() expression: only support from() with 1st parameter as instance value`,
-  );
-  const mapping = guaranteeType(
-    guaranteeType(
-      mappingInstanceExpression.values[0],
-      PackageableElementReference,
-      `Can't process from() expression: only support from() with 1st parameter as packagableElement value`,
-    ).value,
-    Mapping,
-    `Can't process from() expression: only support from() with 1st parameter as mapping value`,
-  );
+  const mappingParameter =
+    expression.parametersValues[2] !== undefined
+      ? expression.parametersValues[1]
+      : undefined;
+  const runtimeParameter = mappingParameter
+    ? expression.parametersValues[2]
+    : expression.parametersValues[1];
+  let mapping: Mapping | undefined;
+  if (mappingParameter) {
+    // mapping
+    const mappingInstanceExpression = guaranteeType(
+      mappingParameter,
+      InstanceValue,
+      `Can't process from() expression: only support from() with 1st parameter as instance value`,
+    );
+    mapping = guaranteeType(
+      guaranteeType(
+        mappingInstanceExpression.values[0],
+        PackageableElementReference,
+        `Can't process from() expression: only support from() with 1st parameter as packagableElement value`,
+      ).value,
+      Mapping,
+      `Can't process from() expression: only support from() with 1st parameter as mapping value`,
+    );
+  }
   // runtime
   const runtimeInstanceExpression = guaranteeType(
-    expression.parametersValues[2],
+    runtimeParameter,
     InstanceValue,
     `Can't process from() expression: only support from() with 2nd parameter as instance value`,
   );
@@ -516,6 +532,16 @@ export class QueryBuilderValueSpecificationProcessor
     throw new UnsupportedOperationError();
   }
 
+  visit_AccessorInstanceValue(
+    valueAccessorInstanceValue: AccessorInstanceValue,
+  ): void {
+    const value = guaranteeNonNullable(
+      valueAccessorInstanceValue.values[0],
+      `Accessor instance value must have a value`,
+    );
+    this.queryBuilderState.setSourceElement(value);
+  }
+
   visit_SimpleFunctionExpression(
     valueSpecification: SimpleFunctionExpression,
   ): void {
@@ -832,8 +858,8 @@ export class QueryBuilderValueSpecificationProcessor
     } else if (matchFunctionName(functionName, [SUPPORTED_FUNCTIONS.FROM])) {
       const parameters = valueSpecification.parametersValues;
       assertTrue(
-        parameters.length === 3,
-        'From function expects 2 parameters (mapping and runtime)',
+        parameters.length === 3 || parameters.length === 2,
+        'From function expects an optional mapping with required runtime',
       );
       processFromFunction(valueSpecification, this.queryBuilderState);
       QueryBuilderValueSpecificationProcessor.processChild(
@@ -874,6 +900,18 @@ export class QueryBuilderValueSpecificationProcessor
         this.parentLambda,
       );
       return;
+    } else if (
+      matchFunctionName(
+        functionName,
+        QUERY_BUILDER_SUPPORTED_FUNCTIONS.RELATION_EXTEND,
+      )
+    ) {
+      processTDS_ExtendExpression(
+        valueSpecification,
+        this.queryBuilderState,
+        this.parentLambda,
+      );
+      return;
     }
     throw new UnsupportedOperationError(
       `Can't process expression of function ${functionName}()`,
@@ -983,38 +1021,64 @@ export class QueryBuilderValueSpecificationProcessor
         `Can't process col spec array instance: value expected to be of size 1`,
       );
       guaranteeNonNullable(spec[0]).colSpecs.forEach((col) => {
-        const _function1 = guaranteeType(
-          col.function1,
-          LambdaFunctionInstanceValue,
-          `Can't process col spec: function1 not a lambda function instance value`,
-        );
-        assertTrue(_function1.values.length === 1);
-        const lambdaVal = guaranteeNonNullable(_function1.values[0]);
-        assertTrue(lambdaVal.expressionSequence.length === 1);
-        const expression = guaranteeNonNullable(
-          lambdaVal.expressionSequence[0],
-        );
-
-        if (expression instanceof AbstractPropertyExpression) {
-          processTDSProjectionColumnPropertyExpression(
-            expression,
-            col.name,
-            this.queryBuilderState,
+        const _function1 = col.function1;
+        if (_function1 instanceof LambdaFunctionInstanceValue) {
+          assertTrue(_function1.values.length === 1);
+          const lambdaVal = guaranteeNonNullable(_function1.values[0]);
+          assertTrue(lambdaVal.expressionSequence.length === 1);
+          const expression = guaranteeNonNullable(
+            lambdaVal.expressionSequence[0],
           );
-        } else if (expression instanceof INTERNAL__UnknownValueSpecification) {
+          if (expression instanceof AbstractPropertyExpression) {
+            processTDSProjectionColumnPropertyExpression(
+              expression,
+              col.name,
+              this.queryBuilderState,
+            );
+          } else if (expression instanceof FunctionExpression) {
+            const func = expression.func;
+            const relationCol = guaranteeType(
+              func,
+              RelationColumn,
+              'Can`t process col spec: function1 lambda function does not contain a relation column',
+            );
+            processTDSRelationColumn(
+              expression,
+              col.name,
+              relationCol,
+              this.queryBuilderState,
+            );
+          } else if (
+            expression instanceof INTERNAL__UnknownValueSpecification
+          ) {
+            assertNonNullable(
+              this.parentExpression,
+              `Can't process unknown value: parent expression cannot be retrieved`,
+            );
+            processTDSProjectionDerivationExpression(
+              expression,
+              col.name,
+              this.parentExpression,
+              this.queryBuilderState,
+            );
+          }
+        } else if (_function1 instanceof INTERNAL__UnknownValueSpecification) {
           assertNonNullable(
             this.parentExpression,
             `Can't process unknown value: parent expression cannot be retrieved`,
           );
           processTDSProjectionDerivationExpression(
-            expression,
+            _function1,
             col.name,
             this.parentExpression,
             this.queryBuilderState,
           );
+        } else {
+          throw new UnsupportedOperationError(
+            `Can't process col spec: only support col spec with function lambda value`,
+          );
         }
       });
-
       return;
     } else if (
       matchFunctionName(this.parentExpression.functionName, [

package/src/stores/QueryBuilderStateHashUtils.ts

@@ -36,6 +36,7 @@ export enum QUERY_BUILDER_STATE_HASH_STRUCTURE {
   AGGREGATE_OPERATOR_PERCENTILE = 'AGGREGATE_OPERATOR_PERCENTILE',
   AGGREGATION_STATE = 'AGGREGATION_STATE',
   SIMPLE_PROJECTION_COLUMN_STATE = 'SIMPLE_PROJECTION_COLUMN_STATE',
+  RELATION_COLUMN_PROJECTION_COLUMN_STATE = 'RELATION_COLUMN_PROJECTION_COLUMN_STATE',
   PROPERTY_EXPRESSION_STATE = 'PROPERTY_EXPRESSION_STATE',
   DERIVATION_PROJECTION_COLUMN_STATE = 'DERIVATION_PROJECTION_COLUMN_STATE',
   RESULT_SET_MODIFIER_STATE = 'RESULT_SET_MODIFIER_STATE',

package/src/stores/QueryBuilderTypeaheadHelper.ts

@@ -180,9 +180,9 @@ const buildRawLambda = (
   );
 
   // build getAll function
-  const _class = guaranteeNonNullable(queryBuilderState.class);
+  const _class = guaranteeNonNullable(queryBuilderState.sourceClass);
   const getAllFn = buildGetAllFunction(
-    guaranteeNonNullable(queryBuilderState.class),
+    guaranteeNonNullable(queryBuilderState.sourceClass),
     Multiplicity.ONE,
   );
   const milestoningStereotype = getMilestoneTemporalStereotype(

package/src/stores/QueryBuilderValueSpecificationBuilder.ts

@@ -20,7 +20,7 @@ import {
   isNonNullable,
 } from '@finos/legend-shared';
 import {
-  type Class,
+  Class,
   Multiplicity,
   getMilestoneTemporalStereotype,
   extractElementNameFromPath,
@@ -34,6 +34,7 @@ import {
   SimpleFunctionExpression,
   SUPPORTED_FUNCTIONS,
   RuntimePointer,
+  AccessorInstanceValue,
 } from '@finos/legend-graph';
 import type { QueryBuilderState } from './QueryBuilderState.js';
 import { buildFilterExpression } from './filter/QueryBuilderFilterValueSpecificationBuilder.js';
@@ -154,10 +155,6 @@ export const buildLambdaFunction = (
   queryBuilderState: QueryBuilderState,
   options?: LambdaFunctionBuilderOption,
 ): LambdaFunction => {
-  const _class = guaranteeNonNullable(
-    queryBuilderState.class,
-    'Class is required to build query',
-  );
   const lambdaFunction = new LambdaFunction(
     new FunctionType(
       PackageableElementExplicitReference.create(
@@ -167,70 +164,86 @@ export const buildLambdaFunction = (
     ),
   );
 
-  const milestoningStereotype = getMilestoneTemporalStereotype(
-    _class,
-    queryBuilderState.graphManagerState.graph,
+  const sourceElement = guaranteeNonNullable(
+    queryBuilderState.sourceElement,
+    'Source element of type class or accessor is required to build query',
   );
+  if (sourceElement instanceof Class) {
+    const _class = guaranteeNonNullable(
+      queryBuilderState.sourceClass,
+      'Class is required to build query',
+    );
 
-  if (milestoningStereotype && options?.useAllVersionsForMilestoning) {
-    // build getAllVersions() when we preview data for milestoned classes
-    // because if we use getAll() we need to pass in data to execute the query
-    // but we don't give user that option in this flow.
-    const getAllVersionsFunction = buildGetAllVersionsFunction(
+    const milestoningStereotype = getMilestoneTemporalStereotype(
       _class,
-      Multiplicity.ONE,
+      queryBuilderState.graphManagerState.graph,
     );
-    lambdaFunction.expressionSequence[0] = getAllVersionsFunction;
-  } else {
-    switch (queryBuilderState.getAllFunction) {
-      case QUERY_BUILDER_SUPPORTED_GET_ALL_FUNCTIONS.GET_ALL_VERSIONS: {
-        if (milestoningStereotype) {
-          const getAllVersionsFunction = buildGetAllVersionsFunction(
-            _class,
-            Multiplicity.ONE,
-          );
-          lambdaFunction.expressionSequence[0] = getAllVersionsFunction;
-        } else {
-          throw new UnsupportedOperationError(
-            `Unable to build query lamdba: getAllVersions() expects source class to be milestoned`,
-          );
+
+    if (milestoningStereotype && options?.useAllVersionsForMilestoning) {
+      // build getAllVersions() when we preview data for milestoned classes
+      // because if we use getAll() we need to pass in data to execute the query
+      // but we don't give user that option in this flow.
+      const getAllVersionsFunction = buildGetAllVersionsFunction(
+        _class,
+        Multiplicity.ONE,
+      );
+      lambdaFunction.expressionSequence[0] = getAllVersionsFunction;
+    } else {
+      switch (queryBuilderState.getAllFunction) {
+        case QUERY_BUILDER_SUPPORTED_GET_ALL_FUNCTIONS.GET_ALL_VERSIONS: {
+          if (milestoningStereotype) {
+            const getAllVersionsFunction = buildGetAllVersionsFunction(
+              _class,
+              Multiplicity.ONE,
+            );
+            lambdaFunction.expressionSequence[0] = getAllVersionsFunction;
+          } else {
+            throw new UnsupportedOperationError(
+              `Unable to build query lamdba: getAllVersions() expects source class to be milestoned`,
+            );
+          }
+          break;
         }
-        break;
-      }
-      case QUERY_BUILDER_SUPPORTED_GET_ALL_FUNCTIONS.GET_ALL_VERSIONS_IN_RANGE: {
-        if (milestoningStereotype) {
-          const getAllVersionsInRangeFunction =
-            buildGetAllVersionsInRangeFunction(_class, Multiplicity.ONE);
-          queryBuilderState.milestoningState
-            .getMilestoningImplementation(milestoningStereotype)
-            .buildGetAllVersionsInRangeParameters(
-              getAllVersionsInRangeFunction,
+        case QUERY_BUILDER_SUPPORTED_GET_ALL_FUNCTIONS.GET_ALL_VERSIONS_IN_RANGE: {
+          if (milestoningStereotype) {
+            const getAllVersionsInRangeFunction =
+              buildGetAllVersionsInRangeFunction(_class, Multiplicity.ONE);
+            queryBuilderState.milestoningState
+              .getMilestoningImplementation(milestoningStereotype)
+              .buildGetAllVersionsInRangeParameters(
+                getAllVersionsInRangeFunction,
+              );
+            lambdaFunction.expressionSequence[0] =
+              getAllVersionsInRangeFunction;
+          } else {
+            throw new UnsupportedOperationError(
+              `Unable to build query lamdba: getAllVersionsInRange() expects source class to be milestoned`,
             );
-          lambdaFunction.expressionSequence[0] = getAllVersionsInRangeFunction;
-        } else {
-          throw new UnsupportedOperationError(
-            `Unable to build query lamdba: getAllVersionsInRange() expects source class to be milestoned`,
-          );
+          }
+          break;
         }
-        break;
-      }
-      case QUERY_BUILDER_SUPPORTED_GET_ALL_FUNCTIONS.GET_ALL: {
-        // build getAll()
-        const getAllFunction = buildGetAllFunction(_class, Multiplicity.ONE);
-        if (milestoningStereotype) {
-          // build milestoning parameter(s) for getAll()
-          queryBuilderState.milestoningState
-            .getMilestoningImplementation(milestoningStereotype)
-            .buildGetAllParameters(getAllFunction);
+        case QUERY_BUILDER_SUPPORTED_GET_ALL_FUNCTIONS.GET_ALL: {
+          // build getAll()
+          const getAllFunction = buildGetAllFunction(_class, Multiplicity.ONE);
+          if (milestoningStereotype) {
+            // build milestoning parameter(s) for getAll()
+            queryBuilderState.milestoningState
+              .getMilestoningImplementation(milestoningStereotype)
+              .buildGetAllParameters(getAllFunction);
+          }
+          lambdaFunction.expressionSequence[0] = getAllFunction;
+          break;
         }
-        lambdaFunction.expressionSequence[0] = getAllFunction;
-        break;
+        default:
+          throw new UnsupportedOperationError(
+            `Unable to build query lambda: unknown ${queryBuilderState.getAllFunction} function`,
+          );
       }
-      default:
-        throw new UnsupportedOperationError(
-          `Unable to build query lambda: unknown ${queryBuilderState.getAllFunction} function`,
-        );
     }
+  } else {
+    const accessorInstanceValue = new AccessorInstanceValue();
+    accessorInstanceValue.values = [sourceElement];
+    lambdaFunction.expressionSequence[0] = accessorInstanceValue;
   }
 
   // build watermark

package/src/stores/__test-utils__/QueryBuilderStateTestUtils.ts

@@ -140,7 +140,7 @@ export const TEST__setUpQueryBuilderState = async (
   }
   if (executionContext) {
     const graph = queryBuilderState.graphManagerState.graph;
-    queryBuilderState.class = graph.getClass(executionContext._class);
+    queryBuilderState.sourceElement = graph.getClass(executionContext._class);
     queryBuilderState.executionContextState.mapping = graph.getMapping(
       executionContext.mapping,
     );

package/src/stores/explorer/QueryBuilderExplorerState.ts

@@ -60,6 +60,7 @@ import {
   getAllSubclasses,
   PropertyExplicitReference,
   reportGraphAnalytics,
+  type RelationColumn,
 } from '@finos/legend-graph';
 import type { QueryBuilderState } from '../QueryBuilderState.js';
 import { action, flow, flowResult, makeObservable, observable } from 'mobx';
@@ -80,6 +81,7 @@ export enum QUERY_BUILDER_EXPLORER_TREE_DND_TYPE {
   CLASS_PROPERTY = 'CLASS_PROPERTY',
   ENUM_PROPERTY = 'ENUM_PROPERTY',
   PRIMITIVE_PROPERTY = 'PRIMITIVE_PROPERTY',
+  RELATION_COLUMN = 'RELATION_COLUMN',
 }
 
 export const generateExplorerTreePropertyNodeID = (
@@ -96,6 +98,10 @@ export interface QueryBuilderExplorerTreeDragSource {
   node: QueryBuilderExplorerTreePropertyNodeData;
 }
 
+export interface QueryBuilderExplorerTreeRelationColumnDragSource {
+  node: QueryBuilderExplorerTreeRelationColumnNodeData;
+}
+
 export abstract class QueryBuilderExplorerTreeNodeData implements TreeNodeData {
   isSelected?: boolean | undefined;
   isOpen?: boolean | undefined;
@@ -221,6 +227,24 @@ export class QueryBuilderExplorerTreeSubTypeNodeData extends QueryBuilderExplore
   }
 }
 
+export class QueryBuilderExplorerTreeRelationRootNodeData extends QueryBuilderExplorerTreeNodeData {}
+
+export class QueryBuilderExplorerTreeRelationColumnNodeData extends QueryBuilderExplorerTreeNodeData {
+  column: RelationColumn;
+
+  constructor(
+    id: string,
+    label: string,
+    dndText: string,
+    column: RelationColumn,
+    type: Type,
+    mappingData: QueryBuilderExplorerTreeNodeMappingData,
+  ) {
+    super(id, label, dndText, false, type, mappingData);
+    this.column = column;
+  }
+}
+
 export const buildPropertyExpressionFromExplorerTreeNodeData = (
   node: QueryBuilderExplorerTreePropertyNodeData,
   explorerState: QueryBuilderExplorerState,
@@ -807,7 +831,7 @@ export class QueryBuilderExplorerState {
   }
 
   refreshTreeData(): void {
-    const _class = this.queryBuilderState.class;
+    const _class = this.queryBuilderState.sourceClass;
     const _mapping = this.queryBuilderState.executionContextState.mapping;
     this.setTreeData(
       _class && _mapping && this.mappingModelCoverageAnalysisResult
@@ -984,7 +1008,7 @@ export class QueryBuilderExplorerState {
     }
     if (
       !node.mappingData.mapped ||
-      !this.queryBuilderState.class ||
+      !this.queryBuilderState.sourceClass ||
       !this.queryBuilderState.executionContextState.mapping
     ) {
       return;