Migrate syncStructuredQueryColumnsAndSettings to MLv2 (#37693)

* Migrate syncStructuredQueryColumnsAndSettings to MLv2

* Migrate syncStructuredQueryColumnsAndSettings to MLv2

* Fix types

* Support converting aggregation references to legacy references

Fixes #37698.

* Update TS wrapper

* Pass query and stageIndex to legacyQuery

* Pass query and stageIndex to legacyQuery

* Remove no longer used columnDimensionWithName

* Update assertion in failing test

* Replace non-existing Question.prototype.isNative & Question.prototype.isStructured calls

---------

Co-authored-by: Tamás Benkő <tamas@metabase.com>
Co-authored-by: Kamil Mielnik <kamil@kamilmielnik.com>

frontend/src/metabase-lib/Question.ts

@@ -18,7 +18,7 @@ import type BaseQuery from "metabase-lib/queries/Query";
 import Metadata from "metabase-lib/metadata/Metadata";
 import type Database from "metabase-lib/metadata/Database";
 import type Table from "metabase-lib/metadata/Table";
-import { AggregationDimension, FieldDimension } from "metabase-lib/Dimension";
+import { FieldDimension } from "metabase-lib/Dimension";
 import { isFK } from "metabase-lib/types/utils/isa";
 import { sortObject } from "metabase-lib/utils";
 
@@ -34,6 +34,7 @@ import type {
   ParameterValues,
   ParameterId,
   VisualizationSettings,
+  Dataset,
 } from "metabase-types/api";
 
 import * as AGGREGATION from "metabase-lib/queries/utils/aggregation";
@@ -57,7 +58,6 @@ import {
   ALERT_TYPE_ROWS,
   ALERT_TYPE_TIMESERIES_GOAL,
 } from "metabase-lib/Alert";
-import { getBaseDimensionReference } from "metabase-lib/references";
 
 import type { Query } from "./types";
 
@@ -602,11 +602,12 @@ class Question {
     });
   }
 
-  private _syncStructuredQueryColumnsAndSettings(
-    previousQuestion,
-    previousQuery,
-  ) {
-    const query = this.legacyQuery({ useStructuredQuery: true });
+  private _syncStructuredQueryColumnsAndSettings(previousQuestion: Question) {
+    const query = this.query();
+    const previousQuery = previousQuestion.query();
+    const stageIndex = -1;
+    const columns = Lib.returnedColumns(query, stageIndex);
+    const previousColumns = Lib.returnedColumns(previousQuery, stageIndex);
 
     if (
       !_.isEqual(
@@ -617,32 +618,37 @@ class Question {
       return this;
     }
 
-    const addedColumnNames = _.difference(
-      query.columnNames(),
-      previousQuery.columnNames(),
-    );
-
-    const removedColumnNames = _.difference(
-      previousQuery.columnNames(),
-      query.columnNames(),
-    );
-
+    const addedColumns = columns
+      .filter(
+        column =>
+          !Lib.findMatchingColumn(query, stageIndex, column, previousColumns),
+      )
+      .map(column => ({
+        column,
+        columnInfo: Lib.displayInfo(query, stageIndex, column),
+      }));
+    const removedColumns = previousColumns
+      .filter(
+        column =>
+          !Lib.findMatchingColumn(previousQuery, stageIndex, column, columns),
+      )
+      .map(column => ({
+        column,
+        columnInfo: Lib.displayInfo(previousQuery, stageIndex, column),
+      }));
     const graphMetrics = this.setting("graph.metrics");
 
     if (
       graphMetrics &&
-      (addedColumnNames.length > 0 || removedColumnNames.length > 0)
+      (addedColumns.length > 0 || removedColumns.length > 0)
     ) {
-      const addedMetricColumnNames = addedColumnNames.filter(
-        name =>
-          query.columnDimensionWithName(name) instanceof AggregationDimension,
-      );
+      const addedMetricColumnNames = addedColumns
+        .filter(({ columnInfo }) => columnInfo.isAggregation)
+        .map(({ columnInfo }) => columnInfo.name);
 
-      const removedMetricColumnNames = removedColumnNames.filter(
-        name =>
-          previousQuery.columnDimensionWithName(name) instanceof
-          AggregationDimension,
-      );
+      const removedMetricColumnNames = removedColumns
+        .filter(({ columnInfo }) => columnInfo.isAggregation)
+        .map(({ columnInfo }) => columnInfo.name);
 
       if (
         addedMetricColumnNames.length > 0 ||
@@ -660,20 +666,23 @@ class Question {
     const tableColumns = this.setting("table.columns");
     if (
       tableColumns &&
-      (addedColumnNames.length > 0 || removedColumnNames.length > 0)
+      (addedColumns.length > 0 || removedColumns.length > 0)
     ) {
       return this.updateSettings({
         "table.columns": [
           ...tableColumns.filter(
             column =>
-              !removedColumnNames.includes(column.name) &&
-              !addedColumnNames.includes(column.name),
+              !addedColumns.some(
+                ({ columnInfo }) => column.name === columnInfo.name,
+              ) &&
+              !removedColumns.some(
+                ({ columnInfo }) => column.name === columnInfo.name,
+              ),
           ),
-          ...addedColumnNames.map(name => {
-            const dimension = query.columnDimensionWithName(name);
+          ...addedColumns.map(({ column, columnInfo }) => {
             return {
-              name: name,
-              fieldRef: getBaseDimensionReference(dimension.mbql()),
+              name: columnInfo.name,
+              fieldRef: Lib.legacyRef(query, stageIndex, column),
               enabled: true,
             };
           }),
@@ -721,25 +730,22 @@ class Question {
     });
   }
 
-  syncColumnsAndSettings(previous, queryResults) {
-    const query = this.legacyQuery({ useStructuredQuery: true });
-    const isQueryResultValid = queryResults && !queryResults.error;
+  syncColumnsAndSettings(previousQuestion?: Question, queryResults?: Dataset) {
+    const query = this.query();
+    const { isNative } = Lib.queryDisplayInfo(query);
 
-    if (query instanceof NativeQuery && isQueryResultValid) {
+    if (isNative && queryResults && !queryResults.error) {
       return this._syncNativeQuerySettings(queryResults);
     }
 
-    const previousQuery =
-      previous && previous.legacyQuery({ useStructuredQuery: true });
+    if (previousQuestion) {
+      const previousQuery = previousQuestion.query();
+      const { isNative: isPreviousQuestionNative } =
+        Lib.queryDisplayInfo(previousQuery);
 
-    if (
-      query instanceof StructuredQuery &&
-      previousQuery instanceof StructuredQuery
-    ) {
-      return this._syncStructuredQueryColumnsAndSettings(
-        previous,
-        previousQuery,
-      );
+      if (!isNative && !isPreviousQuestionNative) {
+        return this._syncStructuredQueryColumnsAndSettings(previousQuestion);
+      }
     }
 
     return this;

frontend/src/metabase-lib/queries/StructuredQuery.ts

@@ -1240,14 +1240,6 @@ class StructuredQuery extends AtomicQuery {
     }));
   });
 
-  columnDimensionWithName(columnName) {
-    const index = this.columnNames().findIndex(n => n === columnName);
-
-    if (index >= 0) {
-      return this.columnDimensions()[index];
-    }
-  }
-
   setDatasetQuery(datasetQuery: DatasetQuery): StructuredQuery {
     return new StructuredQuery(this._originalQuestion, datasetQuery);
   }

frontend/test/metabase/lib/dataset.unit.spec.js

@@ -171,7 +171,7 @@ describe("metabase/util/dataset", () => {
         ...prevQuestion.setting("table.columns"),
         createMockTableColumnOrderSetting({
           name: "VENDOR",
-          fieldRef: ["field", PRODUCTS.VENDOR, null],
+          fieldRef: ["field", PRODUCTS.VENDOR, { "base-type": "type/Text" }],
           enabled: true,
         }),
       ]);