Repository: finos/legend-studio (ts)

Issue:
Bug: Query parameters getting lost while switching between graphFetch and projection mode in QueryBuilder
### Similar issues

- [X] I have searched and found no existing similar issues

### How are you using Studio?

Legend Query

### Current and expected behavior

Current Behaviour:
 When you toggle between graphFetch and projection mode in query builder you notice that the query parameters are getting lost.

Expected Behaviour:
When you toggle between graphFetch and projection mode in query builder query parameters should not get changed.

### Steps to reproduce

Copy and paste the below model data and open the query builder from service editor

### Model data

```text
###Service
Service my::service
{
  pattern: '/e70d2c08-4169-44f1-89d4-0fa9971c3995';
  documentation: '';
  autoActivateUpdates: true;
  execution: Single
  {
    query: businessDate: Date[1]|my::Person.all();
    mapping: my::map;
    runtime:
    #{
      mappings:
      [
        my::map
      ];
    }#;
  }
  test: Single
  {
    data: '';
    asserts:
    [
    ];
  }
}


###Relational
Database my::db
(
  Table PersonTable
  (
    NAME CHAR(200),
    FIRMID INTEGER
  )
  Table FirmTable
  (
    FIRMNAME CHAR(200),
    ID INTEGER
  )

  Join Firm_Person(PersonTable.FIRMID = FirmTable.ID)
)


###Pure
Class my::Firm
{
  legalName: String[1];
  id: Integer[1];
}

Class my::Person
{
  name: String[1];
  firmID: Integer[1];
}


###Mapping
Mapping my::map
(
  *my::Person: Pure
  {
  }
)

```


### Environment

```text
Windows 10
```


### Possible solution and workaround

_No response_

### Contribution

- [X] I would like to work on the fix for this issue

Pull request:
Refactor reset states of QueryBuilder
<!--
  Thanks for submitting a pull request!
  We appreciate you spending the time to work on these changes. Please provide enough information so that others can review your pull request.

  Before submitting a pull request, please make sure the following is done:
  1. Fork [the repository](https://github.com/finos/legend-studio) and create your branch from `master`.
  2. Run `yarn` in the repository root.
  3. If you've fixed a bug or added code that should be tested, add tests!
  4. Add a changeset to summarize your change in the changelogs.
  5. If you haven't already, complete the CLA.

  Learn more about contributing: https://github.com/finos/legend-studio/blob/master/CONTRIBUTING.md
-->

## Summary
This PR makes the resetting behaviour in QueryBuilder more consistent. It also fig the bug where the query parameters are getting lost while switching between graphFetch and projection in Query builder.
Closes #704 
<!--
  Explain the **motivation** for making this change. What existing problem(s) does this PR solve?
  Also, link the issue(s) to this PR: e.g. Fixes #1, Resolves #2
  If you also added documentation, link the PR from https://github.com/finos/legend
-->

## How did you test this change?

- [ ] Test(s) added
- [ ] Manual testing (please provide screenshots/recordings)
- [x] No testing (please provide an explanation)
NA

<!--
  Demonstrate the code is solid. Screenshots/videos if the pull request changes the user interface.
  How exactly did you verify that your PR solves the issue you wanted to solve?
  If you leave this empty, your PR will very likely be closed.
-->

<!--
  More in-depth docs below:
  - Git workflow: https://github.com/finos/legend-studio/blob/master/docs/workflow/working-with-github.md
  - Test: https://github.com/finos/legend-studio/blob/master/docs/technical/test-strategy.md
  - Changeset: https://github.com/finos/legend-studio/blob/master/CONTRIBUTING.md#changeset
  - Dependency: https://github.com/finos/legend-studio/blob/master/docs/workflow/dependencies.md
  - UX/UI: https://github.com/finos/legend-studio/tree/master/docs/ux
-->


Patch:
diff --git a/.changeset/little-boxes-laugh.md b/.changeset/little-boxes-laugh.md
new file mode 100644
index 000000000..da4f92dcc
--- /dev/null
+++ b/.changeset/little-boxes-laugh.md
@@ -0,0 +1,3 @@
+---
+"@finos/legend-studio-extension-query-builder": patch
+---
diff --git a/.changeset/spotty-cows-rush.md b/.changeset/spotty-cows-rush.md
new file mode 100644
index 000000000..c797225a5
--- /dev/null
+++ b/.changeset/spotty-cows-rush.md
@@ -0,0 +1,5 @@
+---
+"@finos/legend-query": patch
+---
+
+Make resetting behaviour in query builder more consistent and systematic.
diff --git a/packages/legend-query/src/components/QueryBuilderFetchStructurePanel.tsx b/packages/legend-query/src/components/QueryBuilderFetchStructurePanel.tsx
index a7f9474ed..62d77db60 100644
--- a/packages/legend-query/src/components/QueryBuilderFetchStructurePanel.tsx
+++ b/packages/legend-query/src/components/QueryBuilderFetchStructurePanel.tsx
@@ -79,7 +79,7 @@ export const QueryBuilderFetchStructurePanel = observer(
         if (fetchStructureState.fetchStructureMode !== fetchMode) {
           fetchStructureState.setFetchStructureMode(fetchMode);
           // TODO: might want to add alert modal to alert user changing fetch structure resets state
-          queryBuilderState.resetData();
+          queryBuilderState.changeFetchStructure();
         }
       };
 
diff --git a/packages/legend-query/src/components/QueryBuilderSetupPanel.tsx b/packages/legend-query/src/components/QueryBuilderSetupPanel.tsx
index c3edad481..12d139576 100644
--- a/packages/legend-query/src/components/QueryBuilderSetupPanel.tsx
+++ b/packages/legend-query/src/components/QueryBuilderSetupPanel.tsx
@@ -50,8 +50,7 @@ export const QueryBuilderSetupPanel = observer(
       ? { value: querySetupState._class, label: querySetupState._class.name }
       : null;
     const changeClass = (val: PackageableElementOption<Class>): void => {
-      querySetupState.setClass(val.value);
-      queryBuilderState.resetData();
+      queryBuilderState.changeClass(val.value);
     };
     // mapping
     const mappingOptions = querySetupState.possibleMappings.map((mapping) => ({
@@ -81,7 +80,8 @@ export const QueryBuilderSetupPanel = observer(
     const changeMapping = (val: PackageableElementOption<Mapping>): void => {
       if (querySetupState._class && !querySetupState.mappingIsReadOnly) {
         querySetupState.setMapping(val.value);
-        queryBuilderState.resetData();
+        queryBuilderState.resetQueryBuilder();
+        queryBuilderState.resetQuerySetup();
       }
     };
     // runtime
diff --git a/packages/legend-query/src/stores/LegendQueryStore.ts b/packages/legend-query/src/stores/LegendQueryStore.ts
index 377c1e774..1445304a7 100644
--- a/packages/legend-query/src/stores/LegendQueryStore.ts
+++ b/packages/legend-query/src/stores/LegendQueryStore.ts
@@ -609,7 +609,8 @@ export class LegendQueryStore {
           this.queryBuilderState.querySetupState._class = possibleTargets[0];
         }
       }
-      this.queryBuilderState.resetData();
+      this.queryBuilderState.resetQueryBuilder();
+      this.queryBuilderState.resetQuerySetup();
       this.setOnSaveQuery(async (lambda: RawLambda) => {
         this.setQueryExportState(
           new QueryExportState(this, lambda, false, undefined),
diff --git a/packages/legend-query/src/stores/QueryBuilderState.ts b/packages/legend-query/src/stores/QueryBuilderState.ts
index 34e9a69f6..e6ccf44dc 100644
--- a/packages/legend-query/src/stores/QueryBuilderState.ts
+++ b/packages/legend-query/src/stores/QueryBuilderState.ts
@@ -163,10 +163,13 @@ export class QueryBuilderState {
       runtimeOptions: computed,
       serviceOptions: computed,
       setMode: action,
-      resetData: action,
+      resetQueryBuilder: action,
+      resetQuerySetup: action,
       buildStateFromRawLambda: action,
       saveQuery: action,
       setBackdrop: action,
+      changeClass: action,
+      changeFetchStructure: action,
       compileQuery: flow,
     });
 
@@ -203,8 +206,17 @@ export class QueryBuilderState {
       : guaranteeNonNullable(this.queryUnsupportedState.rawLambda);
   }
 
-  resetData(): void {
+  resetQueryBuilder(): void {
+    const resultState = new QueryBuilderResultState(this);
+    resultState.setPreviewLimit(this.resultState.previewLimit);
+    this.resultState = resultState;
+    this.queryTextEditorState = new QueryTextEditorState(this);
+    this.queryUnsupportedState = new QueryBuilderUnsupportedState(this);
+  }
+
+  resetQuerySetup(): void {
     this.explorerState = new QueryBuilderExplorerState(this);
+    this.explorerState.refreshTreeData();
     this.queryParametersState = new QueryParametersState(this);
     const fetchStructureState = new QueryBuilderFetchStructureState(this);
     fetchStructureState.setFetchStructureMode(
@@ -213,12 +225,6 @@ export class QueryBuilderState {
     this.fetchStructureState = fetchStructureState;
     this.filterState = new QueryBuilderFilterState(this, this.filterOperators);
     this.resultSetModifierState = new QueryResultSetModifierState(this);
-    const resultState = new QueryBuilderResultState(this);
-    resultState.setPreviewLimit(this.resultState.previewLimit);
-    this.resultState = resultState;
-    this.queryTextEditorState = new QueryTextEditorState(this);
-    this.queryUnsupportedState = new QueryBuilderUnsupportedState(this);
-    this.explorerState.refreshTreeData();
     this.fetchStructureState.graphFetchTreeState.initialize();
   }
 
@@ -231,8 +237,7 @@ export class QueryBuilderState {
       this.buildStateFromRawLambda(rawLambda);
     } catch (error) {
       assertErrorThrown(error);
-      this.querySetupState.setClass(undefined, true);
-      this.resetData();
+      this.changeClass(undefined, true);
       if (options?.notifyError) {
         this.applicationStore.notifyError(
           `Unable to initialize query builder: ${error.message}`,
@@ -250,7 +255,8 @@ export class QueryBuilderState {
    * consumers of function should handle the errors.
    */
   buildStateFromRawLambda(rawLambda: RawLambda): void {
-    this.resetData();
+    this.resetQueryBuilder();
+    this.resetQuerySetup();
     if (!rawLambda.isStub) {
       const valueSpec =
         this.graphManagerState.graphManager.buildValueSpecification(
@@ -393,6 +399,18 @@ export class QueryBuilderState {
     }
   }
 
+  changeClass(val: Class | undefined, isRebuildingState?: boolean): void {
+    this.resetQueryBuilder();
+    this.resetQuerySetup();
+    this.querySetupState.setClass(val, isRebuildingState);
+    this.explorerState.refreshTreeData();
+  }
+
+  changeFetchStructure(): void {
+    this.resultSetModifierState = new QueryResultSetModifierState(this);
+    this.fetchStructureState.graphFetchTreeState.initialize();
+  }
+
   get classOptions(): PackageableElementOption<Class>[] {
     return this.graphManagerState.graph.ownClasses
       .concat(
diff --git a/packages/legend-studio-extension-query-builder/src/components/QueryBuilder_LegendStudioPlugin.tsx b/packages/legend-studio-extension-query-builder/src/components/QueryBuilder_LegendStudioPlugin.tsx
index 4f2c8679f..59ab7eaf0 100644
--- a/packages/legend-studio-extension-query-builder/src/components/QueryBuilder_LegendStudioPlugin.tsx
+++ b/packages/legend-studio-extension-query-builder/src/components/QueryBuilder_LegendStudioPlugin.tsx
@@ -222,10 +222,7 @@ export class QueryBuilder_LegendStudioPlugin
                 }),
               );
               if (queryBuilderExtension.mode) {
-                queryBuilderExtension.queryBuilderState.querySetupState.setClass(
-                  element,
-                );
-                queryBuilderExtension.queryBuilderState.resetData();
+                queryBuilderExtension.queryBuilderState.changeClass(element);
               }
             };
             return (
@@ -319,10 +316,9 @@ export class QueryBuilder_LegendStudioPlugin
                 }),
               );
               if (queryBuilderExtension.mode) {
-                queryBuilderExtension.queryBuilderState.querySetupState.setClass(
+                queryBuilderExtension.queryBuilderState.changeClass(
                   classView.class.value,
                 );
-                queryBuilderExtension.queryBuilderState.resetData();
               }
             };
             return (


Test patch:
diff --git a/packages/legend-query/src/components/__tests__/QueryBuilder_FetchStructure.test.tsx b/packages/legend-query/src/components/__tests__/QueryBuilder_FetchStructure.test.tsx
index ccea9d44e..4f8dac3e3 100644
--- a/packages/legend-query/src/components/__tests__/QueryBuilder_FetchStructure.test.tsx
+++ b/packages/legend-query/src/components/__tests__/QueryBuilder_FetchStructure.test.tsx
@@ -84,8 +84,7 @@ test(
       'model::relational::tests::simpleRelationalMapping',
     );
 
-    queryBuilderState.querySetupState.setClass(_personClass);
-    queryBuilderState.resetData();
+    queryBuilderState.changeClass(_personClass);
     const queryBuilderSetup = await waitFor(() =>
       renderResult.getByTestId(QUERY_BUILDER_TEST_ID.QUERY_BUILDER_SETUP),
     );
@@ -281,7 +280,8 @@ test(
     ).toBe(0);
 
     // filter with group condition
-    queryBuilderState.resetData();
+    queryBuilderState.resetQueryBuilder();
+    queryBuilderState.resetQuerySetup();
     await waitFor(() => renderResult.getByText('Add a filter condition'));
     queryBuilderState.initialize(
       getRawLambda(TEST_DATA__getAllWithGroupedFilter),
@@ -313,7 +313,8 @@ test(
     ).toBe(0);
 
     // projection column with derived property
-    queryBuilderState.resetData();
+    queryBuilderState.resetQueryBuilder();
+    queryBuilderState.resetQuerySetup();
     await waitFor(() => renderResult.getByText('Add a filter condition'));
     queryBuilderState.initialize(
       getRawLambda(TEST_DATA__projectWithDerivedProperty),
@@ -369,8 +370,7 @@ test(
       'model::target::NFirm',
     );
 
-    queryBuilderState.querySetupState.setClass(_personClass);
-    queryBuilderState.resetData();
+    queryBuilderState.changeClass(_personClass);
     const queryBuilderSetup = await waitFor(() =>
       renderResult.getByTestId(QUERY_BUILDER_TEST_ID.QUERY_BUILDER_SETUP),
     );