fix: join query segments on dimensions instead of primary keys

This ensures that metrics are sliced correctly
VeryBigThings · Apr 14, 2024 · 6fbf1ee · 6fbf1ee
1 parent 0eff7a4
commit 6fbf1ee
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 29 deletions.
diff --git a/src/__tests__/index.test.ts b/src/__tests__/index.test.ts
@@ -84,6 +84,11 @@ await describe("semantic layer", async () => {
           sql`${model.dimension("first_name")} || ' ' || ${model.dimension(
             "last_name",
           )}`,
+      })
+      .withMetric("count", {
+        type: "string",
+        aggregateWith: "count",
+        sql: ({ model }) => model.column("CustomerId"),
       });
 
     const invoicesModel = semanticLayer
@@ -304,6 +309,48 @@ await describe("semantic layer", async () => {
       ]);
     });
 
+    await it("can query a metric and slice it correctly by a non primary key dimension", async () => {
+      const query = queryBuilder.buildQuery({
+        dimensions: ["customers.country"],
+        metrics: ["customers.count"],
+        order: {
+          "customers.country": "asc",
+        },
+      });
+
+      const result = await client.query<InferSqlQueryResultType<typeof query>>(
+        query.sql,
+        query.bindings,
+      );
+
+      assert.deepEqual(result.rows, [
+        { customers___country: "Argentina", customers___count: "1" },
+        { customers___country: "Australia", customers___count: "1" },
+        { customers___country: "Austria", customers___count: "1" },
+        { customers___country: "Belgium", customers___count: "1" },
+        { customers___country: "Brazil", customers___count: "5" },
+        { customers___country: "Canada", customers___count: "8" },
+        { customers___country: "Chile", customers___count: "1" },
+        { customers___country: "Czech Republic", customers___count: "2" },
+        { customers___country: "Denmark", customers___count: "1" },
+        { customers___country: "Finland", customers___count: "1" },
+        { customers___country: "France", customers___count: "5" },
+        { customers___country: "Germany", customers___count: "4" },
+        { customers___country: "Hungary", customers___count: "1" },
+        { customers___country: "India", customers___count: "2" },
+        { customers___country: "Ireland", customers___count: "1" },
+        { customers___country: "Italy", customers___count: "1" },
+        { customers___country: "Netherlands", customers___count: "1" },
+        { customers___country: "Norway", customers___count: "1" },
+        { customers___country: "Poland", customers___count: "1" },
+        { customers___country: "Portugal", customers___count: "2" },
+        { customers___country: "Spain", customers___count: "1" },
+        { customers___country: "Sweden", customers___count: "1" },
+        { customers___country: "USA", customers___count: "13" },
+        { customers___country: "United Kingdom", customers___count: "3" },
+      ]);
+    });
+
     await it("will correctly load distinct dimensions when no metrics are loaded", async () => {
       const query = queryBuilder.buildQuery({
         dimensions: ["customers.country"],

diff --git a/src/lib/query-builder/build-query.ts b/src/lib/query-builder/build-query.ts
@@ -253,9 +253,14 @@ export function buildQuery(
 
   invariant(initialSqlQuerySegment, "No initial sql query segment found");
 
-  const joinOnDimensions = referencedModels.dimensions.flatMap((modelName) =>
+  /*const joinOnDimensions = referencedModels.dimensions.flatMap((modelName) =>
     repository.getModel(modelName).getPrimaryKeyDimensions(),
-  );
+  );*/
+
+  const joinOnDimensions = query.dimensions?.map((dimensionName) => {
+    return repository.getDimension(dimensionName);
+  });
+
   const rootAlias = getAlias(0);
   const rootSqlQuery = knex(initialSqlQuerySegment.sqlQuery.as(rootAlias));
 
@@ -289,7 +294,7 @@ export function buildQuery(
     const segment = restSqlQuerySegments[i]!;
     const alias = getAlias(i + 1);
     const joinOn =
-      (query.dimensions?.length ?? 0) > 0
+      joinOnDimensions && joinOnDimensions.length > 0
         ? joinOnDimensions
             .map((dimension) => {
               return `${dialect.asIdentifier(rootAlias)}.${dimension.getAlias(

diff --git a/src/lib/query-builder/process-query-and-expand-to-segments.ts b/src/lib/query-builder/process-query-and-expand-to-segments.ts
@@ -102,11 +102,8 @@ interface PreparedQuery {
 
 // biome-ignore lint/complexity/noExcessiveCognitiveComplexity: <explanation>
 function getQuerySegment(
-  repository: AnyRepository,
   queryAnalysis: ReturnType<typeof analyzeQuery>,
   metricModel: string | null,
-  index: number,
-  hasMultipleSegments: boolean,
 ): QuerySegment {
   const queries: {
     query: PreparedQuery;
@@ -138,32 +135,15 @@ function getQuerySegment(
     for (const [modelName, dimensions] of Object.entries(
       queryAnalysis.projectedDimensionsByModel,
     )) {
-      const model = repository.getModel(modelName);
       referencedModels.all.add(modelName);
       referencedModels.dimensions.add(modelName);
 
-      const primaryKeyDimensionNames = hasMultipleSegments
-        ? model.getPrimaryKeyDimensions().map((d) => d.getPath())
-        : [];
-
-      if (index === 0) {
-        for (const dimension of dimensions) {
-          queries[q].dimensions.add(dimension);
-        }
-      }
-
-      if (q === "query") {
-        for (const dimension of primaryKeyDimensionNames) {
-          queries[q].dimensions.add(dimension);
-        }
+      for (const dimension of dimensions) {
+        queries[q].dimensions.add(dimension);
       }
 
       modelQueries[modelName] = {
-        dimensions: new Set<string>(
-          index === 0
-            ? new Set([...dimensions, ...primaryKeyDimensionNames])
-            : new Set(primaryKeyDimensionNames),
-        ),
+        dimensions: new Set<string>(dimensions),
         metrics: new Set<string>(),
       };
     }
@@ -238,13 +218,13 @@ export function processQueryAndExpandToSegments(
     metricModels.length === 0
       ? [
           mergeQuerySegmentWithFilters(
-            getQuerySegment(repository, queryAnalysis, null, 0, false),
+            getQuerySegment(queryAnalysis, null),
             query.filters,
           ),
         ]
-      : metricModels.map((model, idx) =>
+      : metricModels.map((model) =>
           mergeQuerySegmentWithFilters(
-            getQuerySegment(repository, queryAnalysis, model, idx, true),
+            getQuerySegment(queryAnalysis, model),
             query.filters,
           ),
         );