minor changes to the query load process

.changeset/clean-toes-learn.md

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

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

@@ -16,13 +16,15 @@
 
 import { LegendApplicationPlugin } from '@finos/legend-application';
 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',
@@ -51,12 +53,13 @@ export type QuerySetupActionConfiguration = {
 };
 
 export type ExistingQueryEditorStateBuilder = (
+  query: Query,
   editorStore: ExistingQueryEditorStore,
 ) => Promise<QueryBuilderState | undefined>;
 
-export type ExistingQueryGraphBuilderChecker = (
-  editorStore: ExistingQueryEditorStore,
-) => boolean | undefined;
+export type QueryGraphBuilderGetter = (
+  editorStore: QueryEditorStore,
+) => ((editorStore: QueryEditorStore) => GeneratorFn<void>) | undefined;
 
 export type QueryEditorActionConfiguration = {
   key: string;
@@ -101,9 +104,9 @@ export abstract class LegendQueryApplicationPlugin extends LegendApplicationPlug
   getExtraExistingQueryEditorStateBuilders?(): ExistingQueryEditorStateBuilder[];
 
   /**
-   * Get the list of existing query graph builder checkers
+   * Get the list of query graph builders
    */
-  getExtraExistingQueryGraphBuilderCheckers?(): ExistingQueryGraphBuilderChecker[];
+  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,
     });
 
@@ -405,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);
@@ -489,9 +489,8 @@ export abstract class QueryEditorStore {
       );
 
       yield this.setUpEditorState();
-      if (this.requiresGraphBuilding()) {
-        yield flowResult(this.buildGraph());
-      }
+      yield flowResult(this.buildGraph());
+
       this.queryBuilderState = (yield this.initializeQueryBuilderState(
         stopWatch,
       )) as QueryBuilderState;
@@ -543,7 +542,7 @@ export abstract class QueryEditorStore {
     );
   }
 
-  *buildGraph(): GeneratorFn<void> {
+  *buildFullGraph(): GeneratorFn<void> {
     const stopWatch = new StopWatch();
 
     const { groupId, artifactId, versionId } = this.getProjectInfo();
@@ -632,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 {
@@ -992,34 +1005,12 @@ export class ExistingQueryEditorStore extends QueryEditorStore {
   }
 
   override async setUpEditorState(): Promise<void> {
-    this.setLightQuery(
-      await this.graphManagerState.graphManager.getLightQuery(this.queryId),
-    );
-  }
-
-  override requiresGraphBuilding(): boolean {
-    const existingQueryGraphBuilderCheckers =
-      this.applicationStore.pluginManager
-        .getApplicationPlugins()
-        .flatMap(
-          (plugin) =>
-            plugin.getExtraExistingQueryGraphBuilderCheckers?.() ?? [],
-        );
-    for (const checker of existingQueryGraphBuilderCheckers) {
-      const isGraphBuildingRequired = checker(this);
-      if (isGraphBuildingRequired !== undefined) {
-        return isGraphBuildingRequired;
-      }
-    }
-    return true;
-  }
-
-  async setupQuery(): Promise<void> {
     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,
@@ -1029,22 +1020,20 @@ export class ExistingQueryEditorStore extends QueryEditorStore {
   async initializeQueryBuilderState(
     stopWatch: StopWatch,
   ): Promise<QueryBuilderState> {
+    const query = guaranteeNonNullable(this.query);
     let queryBuilderState: QueryBuilderState | undefined;
     const existingQueryEditorStateBuilders = this.applicationStore.pluginManager
       .getApplicationPlugins()
       .flatMap(
         (plugin) => plugin.getExtraExistingQueryEditorStateBuilders?.() ?? [],
       );
     for (const builder of existingQueryEditorStateBuilders) {
-      queryBuilderState = await builder(this);
+      queryBuilderState = await builder(query, this);
       if (queryBuilderState) {
         break;
       }
     }
 
-    await this.setupQuery();
-    const query = guaranteeNonNullable(this.query);
-
     // if no extension found, fall back to basic `class -> mapping -> runtime` mode
     queryBuilderState =
       queryBuilderState ??
@@ -1054,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,