Update order-by field when the breakout changes (#17960)

This should fix #16770.

frontend/src/metabase/lib/query/query.js

@@ -28,8 +28,6 @@ import * as E from "./expression";
 import * as FIELD from "./field";
 import * as FIELD_REF from "./field_ref";
 
-import _ from "underscore";
-
 // AGGREGATION
 
 export const getAggregations = (query: SQ) =>
@@ -180,9 +178,23 @@ function setBreakoutClause(query: SQ, breakoutClause: ?BreakoutClause): SQ {
     .map(b => FIELD_REF.getFieldTargetId(b))
     .filter(id => id != null);
   for (const [index, sort] of getOrderBys(query).entries()) {
-    const sortId = FIELD_REF.getFieldTargetId(sort[1]);
-    if (sortId != null && !_.contains(breakoutIds, sortId)) {
-      query = removeOrderBy(query, index);
+    const sortField = sort[1];
+    const sortId = FIELD_REF.getFieldTargetId(sortField);
+    if (sortId != null) {
+      // Remove invalid field reference
+      if (!breakoutIds.includes(sortId)) {
+        query = removeOrderBy(query, index);
+      } else {
+        // Update the field, since it can change its binning, temporal unit, etc
+        const breakoutFields = B.getBreakouts(breakoutClause);
+        const breakoutField = breakoutFields.find(
+          field => FIELD_REF.getFieldTargetId(field) === sortId,
+        );
+        if (breakoutField) {
+          const direction = sort[0];
+          query = updateOrderBy(query, index, [direction, breakoutField]);
+        }
+      }
     }
   }
   // clear fields when changing breakouts

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

@@ -108,6 +108,37 @@ describe("Query", () => {
     });
   });
 
+  describe("updateBreakout", () => {
+    it("should update the field", () => {
+      expect(
+        Query.updateBreakout(
+          {
+            breakout: [["field", 1, null]],
+          },
+          0,
+          ["field", 2, null],
+        ),
+      ).toEqual({
+        breakout: [["field", 2, null]],
+      });
+    });
+    it("should update sort as well", () => {
+      expect(
+        Query.updateBreakout(
+          {
+            breakout: [["field", 3, { "temporal-unit": "month" }]],
+            "order-by": [["asc", ["field", 3, { "temporal-unit": "month" }]]],
+          },
+          0,
+          ["field", 3, { "temporal-unit": "year" }],
+        ),
+      ).toEqual({
+        breakout: [["field", 3, { "temporal-unit": "year" }]],
+        "order-by": [["asc", ["field", 3, { "temporal-unit": "year" }]]],
+      });
+    });
+  });
+
   describe("removeBreakout", () => {
     it("should remove sort as well", () => {
       expect(

frontend/test/metabase/scenarios/binning/binning-reproductions.cy.spec.js

@@ -99,7 +99,7 @@ describe("binning related reproductions", () => {
     cy.findByText("Month");
   });
 
-  it.skip("should be able to update the bucket size / granularity on a field that has sorting applied to it (metabase#16770)", () => {
+  it("should be able to update the bucket size / granularity on a field that has sorting applied to it (metabase#16770)", () => {
     cy.intercept("POST", "/api/dataset").as("dataset");
 
     visitQuestionAdhoc({