Use PMCD returned by mappinganalysis to build minimal graph for query

.changeset/clean-toes-learn.md

@@ -0,0 +1,4 @@
+---
+'@finos/legend-graph': patch
+'@finos/legend-application-studio': patch
+---

.changeset/kind-rabbits-draw.md

@@ -0,0 +1,6 @@
+---
+'@finos/legend-extension-dsl-data-space': minor
+'@finos/legend-application-query': minor
+---
+
+Use PMCD returned by mapping analysis to build minimal graph for query

packages/legend-application-query/src/components/__test-utils__/QueryEditorComponentTestUtils.tsx

@@ -22,7 +22,6 @@ import {
   Query,
   LightQuery,
   RawLambda,
-  PackageableElementExplicitReference,
   type RawMappingModelCoverageAnalysisResult,
 } from '@finos/legend-graph';
 import { DepotServerClient } from '@finos/legend-server-depot';
@@ -128,10 +127,8 @@ export const TEST__setUpQueryEditor = async (
   query.owner = lightQuery.owner;
   query.isCurrentUserQuery = lightQuery.isCurrentUserQuery;
   const _mapping = graphManagerState.graph.getMapping(mappingPath);
-  query.mapping = PackageableElementExplicitReference.create(_mapping);
-  query.runtime = PackageableElementExplicitReference.create(
-    graphManagerState.graph.getRuntime(runtimePath),
-  );
+  query.mapping = mappingPath;
+  query.runtime = runtimePath;
   query.content = 'some content';
 
   createSpy(

packages/legend-application-query/src/stores/LegendQueryApplicationPlugin.ts

@@ -15,15 +15,16 @@
  */
 
 import { LegendApplicationPlugin } from '@finos/legend-application';
-import type { Query } from '@finos/legend-graph';
 import type { QueryBuilderState } from '@finos/legend-query-builder';
+import type { GeneratorFn } from '@finos/legend-shared';
 import type React from 'react';
 import type { LegendQueryPluginManager } from '../application/LegendQueryPluginManager.js';
 import type {
   ExistingQueryEditorStore,
   QueryEditorStore,
 } from './QueryEditorStore.js';
 import type { QuerySetupLandingPageStore } from './QuerySetupStore.js';
+import type { Query } from '@finos/legend-graph';
 
 export enum QuerySetupActionTag {
   PRODUCTIONIZATION = 'Productionization',
@@ -56,6 +57,10 @@ export type ExistingQueryEditorStateBuilder = (
   editorStore: ExistingQueryEditorStore,
 ) => Promise<QueryBuilderState | undefined>;
 
+export type QueryGraphBuilderGetter = (
+  editorStore: QueryEditorStore,
+) => ((editorStore: QueryEditorStore) => GeneratorFn<void>) | undefined;
+
 export type QueryEditorActionConfiguration = {
   key: string;
   renderer: (
@@ -98,6 +103,11 @@ export abstract class LegendQueryApplicationPlugin extends LegendApplicationPlug
    */
   getExtraExistingQueryEditorStateBuilders?(): ExistingQueryEditorStateBuilder[];
 
+  /**
+   * Get the list of query graph builders
+   */
+  getExtraQueryGraphBuilderGetters?(): QueryGraphBuilderGetter[];
+
   /**
    * Get the list of query editor action renderer configurations.
    */

packages/legend-application-query/src/stores/QueryEditorStore.ts

@@ -293,6 +293,7 @@ export abstract class QueryEditorStore {
       setExistingQueryName: action,
       initialize: flow,
       buildGraph: flow,
+      buildFullGraph: flow,
       searchExistingQueryName: flow,
     });
 
@@ -381,6 +382,10 @@ export abstract class QueryEditorStore {
     // do nothing
   }
 
+  requiresGraphBuilding(): boolean {
+    return true;
+  }
+
   async buildQueryForPersistence(
     query: Query,
     rawLambda: RawLambda,
@@ -401,10 +406,9 @@ export abstract class QueryEditorStore {
         RuntimePointer,
         'Query runtime must be of type runtime pointer',
       );
-      query.mapping = PackageableElementExplicitReference.create(
-        this.queryBuilderState.executionContextState.mapping,
-      );
-      query.runtime = runtimeValue.packageableRuntime;
+      (query.mapping =
+        this.queryBuilderState.executionContextState.mapping.path),
+        (query.runtime = runtimeValue.packageableRuntime.value.path);
       query.content =
         await this.graphManagerState.graphManager.lambdaToPureCode(rawLambda);
       config?.decorator?.(query);
@@ -486,6 +490,7 @@ export abstract class QueryEditorStore {
 
       yield this.setUpEditorState();
       yield flowResult(this.buildGraph());
+
       this.queryBuilderState = (yield this.initializeQueryBuilderState(
         stopWatch,
       )) as QueryBuilderState;
@@ -537,7 +542,7 @@ export abstract class QueryEditorStore {
     );
   }
 
-  *buildGraph(): GeneratorFn<void> {
+  *buildFullGraph(): GeneratorFn<void> {
     const stopWatch = new StopWatch();
 
     const { groupId, artifactId, versionId } = this.getProjectInfo();
@@ -626,6 +631,20 @@ export abstract class QueryEditorStore {
       graphBuilderReportData,
     );
   }
+
+  *buildGraph(): GeneratorFn<void> {
+    const queryGraphBuilderGetters = this.applicationStore.pluginManager
+      .getApplicationPlugins()
+      .flatMap((plugin) => plugin.getExtraQueryGraphBuilderGetters?.() ?? []);
+    for (const getter of queryGraphBuilderGetters) {
+      const builderFunction = getter(this);
+      if (builderFunction) {
+        yield flowResult(builderFunction(this));
+        return;
+      }
+    }
+    yield flowResult(this.buildFullGraph());
+  }
 }
 
 export class MappingQueryCreatorStore extends QueryEditorStore {
@@ -986,23 +1005,22 @@ export class ExistingQueryEditorStore extends QueryEditorStore {
   }
 
   override async setUpEditorState(): Promise<void> {
-    this.setLightQuery(
-      await this.graphManagerState.graphManager.getLightQuery(this.queryId),
-    );
-  }
-
-  async initializeQueryBuilderState(
-    stopWatch: StopWatch,
-  ): Promise<QueryBuilderState> {
     const query = await this.graphManagerState.graphManager.getQuery(
       this.queryId,
       this.graphManagerState.graph,
     );
     this.setQuery(query);
+    this.setLightQuery(toLightQuery(query));
     LegendQueryUserDataHelper.addRecentlyViewedQuery(
       this.applicationStore.userDataService,
       query.id,
     );
+  }
+
+  async initializeQueryBuilderState(
+    stopWatch: StopWatch,
+  ): Promise<QueryBuilderState> {
+    const query = guaranteeNonNullable(this.query);
     let queryBuilderState: QueryBuilderState | undefined;
     const existingQueryEditorStateBuilders = this.applicationStore.pluginManager
       .getApplicationPlugins()
@@ -1025,11 +1043,11 @@ export class ExistingQueryEditorStore extends QueryEditorStore {
         this.applicationStore.config.options.queryBuilderConfig,
       );
 
-    queryBuilderState.executionContextState.setMapping(query.mapping.value);
+    const mapping = this.graphManagerState.graph.getMapping(query.mapping);
+    const runtime = this.graphManagerState.graph.getRuntime(query.runtime);
+    queryBuilderState.executionContextState.setMapping(mapping);
     queryBuilderState.executionContextState.setRuntimeValue(
-      new RuntimePointer(
-        PackageableElementExplicitReference.create(query.runtime.value),
-      ),
+      new RuntimePointer(PackageableElementExplicitReference.create(runtime)),
     );
 
     // leverage initialization of query builder state to ensure we handle unsupported queries

packages/legend-application-studio/src/components/editor/editor-group/FunctionEditor.tsx

@@ -1105,8 +1105,7 @@ export const FunctionEditor = observer(() => {
         );
         await flowResult(
           embeddedQueryBuilderState.setEmbeddedQueryBuilderConfiguration({
-            setupQueryBuilderState: (): QueryBuilderState =>
-              functionQueryBuilderState,
+            setupQueryBuilderState: async () => functionQueryBuilderState,
             actionConfigs: [
               {
                 key: 'save-query-btn',

packages/legend-application-studio/src/components/editor/editor-group/mapping-editor/DEPRECATED__MappingTestEditor.tsx

@@ -112,7 +112,7 @@ const MappingTestQueryEditor = observer(
         const embeddedQueryBuilderState = editorStore.embeddedQueryBuilderState;
         await flowResult(
           embeddedQueryBuilderState.setEmbeddedQueryBuilderConfiguration({
-            setupQueryBuilderState: (): QueryBuilderState => {
+            setupQueryBuilderState: async () => {
               const queryBuilderState = new MappingExecutionQueryBuilderState(
                 embeddedQueryBuilderState.editorStore.applicationStore,
                 embeddedQueryBuilderState.editorStore.graphManagerState,

packages/legend-application-studio/src/components/editor/editor-group/mapping-editor/MappingExecutionBuilder.tsx

@@ -205,7 +205,7 @@ const MappingExecutionQueryEditor = observer(
         const embeddedQueryBuilderState = editorStore.embeddedQueryBuilderState;
         await flowResult(
           embeddedQueryBuilderState.setEmbeddedQueryBuilderConfiguration({
-            setupQueryBuilderState: (): QueryBuilderState => {
+            setupQueryBuilderState: async () => {
               const queryBuilderState = new MappingExecutionQueryBuilderState(
                 embeddedQueryBuilderState.editorStore.applicationStore,
                 embeddedQueryBuilderState.editorStore.graphManagerState,

packages/legend-application-studio/src/components/editor/editor-group/mapping-editor/MappingTestableEditor.tsx

@@ -589,7 +589,7 @@ const MappingTestSuiteQueryEditor = observer(
         const embeddedQueryBuilderState = editorStore.embeddedQueryBuilderState;
         await flowResult(
           embeddedQueryBuilderState.setEmbeddedQueryBuilderConfiguration({
-            setupQueryBuilderState: (): QueryBuilderState => {
+            setupQueryBuilderState: async () => {
               const queryBuilderState = new MappingExecutionQueryBuilderState(
                 embeddedQueryBuilderState.editorStore.applicationStore,
                 embeddedQueryBuilderState.editorStore.graphManagerState,

packages/legend-application-studio/src/components/editor/editor-group/service-editor/ServiceExecutionQueryEditor.tsx

@@ -125,7 +125,7 @@ export const ServiceExecutionQueryEditor = observer(
           executionState.selectedExecutionContextState;
         await flowResult(
           embeddedQueryBuilderState.setEmbeddedQueryBuilderConfiguration({
-            setupQueryBuilderState: (): QueryBuilderState => {
+            setupQueryBuilderState: async () => {
               const queryBuilderState = new ServiceQueryBuilderState(
                 embeddedQueryBuilderState.editorStore.applicationStore,
                 embeddedQueryBuilderState.editorStore.graphManagerState,
@@ -459,7 +459,7 @@ export const queryService = async (
       : undefined;
   await flowResult(
     embeddedQueryBuilderState.setEmbeddedQueryBuilderConfiguration({
-      setupQueryBuilderState: (): QueryBuilderState => {
+      setupQueryBuilderState: async () => {
         const queryBuilderState = new ServiceQueryBuilderState(
           embeddedQueryBuilderState.editorStore.applicationStore,
           embeddedQueryBuilderState.editorStore.graphManagerState,

packages/legend-application-studio/src/components/editor/editor-group/uml-editor/ClassQueryBuilder.tsx

@@ -375,7 +375,7 @@ export const queryClass = async (
   const embeddedQueryBuilderState = editorStore.embeddedQueryBuilderState;
   await flowResult(
     embeddedQueryBuilderState.setEmbeddedQueryBuilderConfiguration({
-      setupQueryBuilderState: () => {
+      setupQueryBuilderState: async () => {
         const queryBuilderState = new ClassQueryBuilderState(
           embeddedQueryBuilderState.editorStore.applicationStore,
           embeddedQueryBuilderState.editorStore.graphManagerState,

packages/legend-application-studio/src/stores/editor/EmbeddedQueryBuilderState.ts

@@ -30,7 +30,7 @@ type EmbeddedQueryBuilderActionConfiguration = {
 };
 
 type EmbeddedQueryBuilderConfiguration = {
-  setupQueryBuilderState: () => QueryBuilderState;
+  setupQueryBuilderState: () => Promise<QueryBuilderState>;
   disableCompile?: boolean | undefined;
   actionConfigs: EmbeddedQueryBuilderActionConfiguration[];
 };
@@ -100,7 +100,8 @@ export class EmbeddedQueryBuilderState {
         }
       }
       if (!this.editorStore.graphState.error) {
-        this.queryBuilderState = config.setupQueryBuilderState();
+        this.queryBuilderState =
+          (yield config.setupQueryBuilderState()) as QueryBuilderState;
         this.actionConfigs = config.actionConfigs;
         this.editorStore.applicationStore.layoutService.setBackdropContainerElementID(
           QUERY_BUILDER_COMPONENT_ELEMENT_ID.BACKDROP_CONTAINER,