Fix drillPK for composite keys (#18039)

frontend/src/metabase-lib/lib/Question.js

@@ -544,13 +544,38 @@ export default class Question {
 
   drillPK(field: Field, value: Value): ?Question {
     const query = this.query();
-    if (query instanceof StructuredQuery) {
-      return query
-        .reset()
-        .setTable(field.table)
-        .filter(["=", ["field", field.id, null], value])
-        .question();
+
+    if (!(query instanceof StructuredQuery)) {
+      return;
     }
+
+    const otherPKFilters = query
+      .filters()
+      ?.filter(filter => {
+        const filterField = filter?.field();
+        if (!filterField) {
+          return false;
+        }
+
+        const isNotSameField = filterField.id !== field.id;
+        const isPKEqualsFilter =
+          filterField.isPK() && filter.operatorName() === "=";
+        const isFromSameTable = filterField.table.id === field.table.id;
+
+        return isPKEqualsFilter && isNotSameField && isFromSameTable;
+      })
+      .map(filter => filter.raw());
+
+    const filtersToApply = [
+      ["=", ["field", field.id, null], value],
+      ...otherPKFilters,
+    ];
+
+    const resultedQuery = filtersToApply.reduce((query, filter) => {
+      return query.addFilter(filter);
+    }, query.reset().setTable(field.table));
+
+    return resultedQuery.question();
   }
 
   _syncStructuredQueryColumnsAndSettings(previousQuestion, previousQuery) {

frontend/test/metabase-lib/lib/Question.unit.spec.js

@@ -3,6 +3,7 @@ import {
   SAMPLE_DATASET,
   ORDERS,
   PRODUCTS,
+  createMetadata,
 } from "__support__/sample_dataset_fixture";
 
 import { assoc, dissoc } from "icepick";
@@ -623,6 +624,71 @@ describe("Question", () => {
           },
         });
       });
+
+      describe("with composite PK", () => {
+        // Making TOTAL a PK column in addition to ID
+        const metadata = createMetadata(state =>
+          state.assocIn(
+            ["entities", "fields", ORDERS.TOTAL.id, "semantic_type"],
+            "type/PK",
+          ),
+        );
+        let question;
+
+        beforeEach(() => {
+          question = new Question(orders_raw_card, metadata);
+        });
+
+        it("when drills to one column of a composite key returns equals filter by the column", () => {
+          const drilledQuestion = question.drillPK(ORDERS.ID, 1);
+
+          expect(drilledQuestion.canRun()).toBe(true);
+          expect(drilledQuestion._card.dataset_query).toEqual({
+            type: "query",
+            database: SAMPLE_DATASET.id,
+            query: {
+              "source-table": ORDERS.id,
+              filter: ["=", ["field", ORDERS.ID.id, null], 1],
+            },
+          });
+        });
+
+        it("when drills to both columns of a composite key returns query with equality filter by both PKs", () => {
+          const drilledQuestion = question
+            .drillPK(ORDERS.ID, 1)
+            .drillPK(ORDERS.TOTAL, 1);
+
+          expect(drilledQuestion.canRun()).toBe(true);
+          expect(drilledQuestion._card.dataset_query).toEqual({
+            type: "query",
+            database: SAMPLE_DATASET.id,
+            query: {
+              "source-table": ORDERS.id,
+              filter: [
+                "and",
+                ["=", ["field", ORDERS.TOTAL.id, null], 1],
+                ["=", ["field", ORDERS.ID.id, null], 1],
+              ],
+            },
+          });
+        });
+
+        it("when drills to other table PK removes the previous table PK filters", () => {
+          const drilledQuestion = question
+            .drillPK(ORDERS.ID, 1)
+            .drillPK(PRODUCTS.ID, 1);
+
+          expect(drilledQuestion.canRun()).toBe(true);
+          expect(drilledQuestion._card.dataset_query).toEqual({
+            type: "query",
+            database: SAMPLE_DATASET.id,
+            query: {
+              "source-table": PRODUCTS.id,
+              filter: ["=", ["field", PRODUCTS.ID.id, null], 1],
+            },
+          });
+        });
+      });
     });
   });