feat(cubejs-client-core): add custom granularities support to drilldown()

packages/cubejs-client-core/src/ResultSet.ts

@@ -215,7 +215,7 @@ export default class ResultSet<T extends Record<string, any> = any> {
     normalizedPivotConfig?.y.forEach((member, currentIndex) => values.push([member, yValues[currentIndex]]));
 
     const { filters: parentFilters = [], segments = [] } = this.query();
-    const { measures } = this.loadResponses[0].annotation;
+    const { measures, timeDimensions: timeDimensionsAnnotation } = this.loadResponses[0].annotation;
     let [, measureName] = values.find(([member]) => member === 'measures') || [];
 
     if (measureName === undefined) {
@@ -240,7 +240,9 @@ export default class ResultSet<T extends Record<string, any> = any> {
         const [cubeName, dimension, granularity] = member.split('.');
 
         if (granularity !== undefined) {
-          const range = dayRange(value, value).snapTo(granularity);
+          // dayRange.snapTo now handles both predefined and custom granularities
+          const range = dayRange(value, value, timeDimensionsAnnotation).snapTo(granularity);
+
           const originalTimeDimension = query.timeDimensions?.find((td) => td.dimension);
 
           let dateRange = [
@@ -469,7 +471,7 @@ export default class ResultSet<T extends Record<string, any> = any> {
       !['hour', 'minute', 'second'].includes(timeDimension.granularity);
 
     const [start, end] = dateRange;
-    const range = dayRange(start, end);
+    const range = dayRange(start, end, annotations);
 
     if (isPredefinedGranularity(timeDimension.granularity)) {
       return TIME_SERIES[timeDimension.granularity](

packages/cubejs-client-core/src/time.ts

@@ -79,25 +79,6 @@ export const isPredefinedGranularity = (granularity: TimeDimensionGranularity):
 export const DateRegex = /^\d\d\d\d-\d\d-\d\d$/;
 export const LocalDateRegex = /^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}.\d{3}Z?$/;
 
-export const dayRange = (from: any, to: any): DayRange => ({
-  by: (value: any) => {
-    const results = [];
-
-    let start = internalDayjs(from);
-    const end = internalDayjs(to);
-
-    while (start.startOf(value).isBefore(end) || start.isSame(end)) {
-      results.push(start);
-      start = start.add(1, value);
-    }
-
-    return results;
-  },
-  snapTo: (value: any): DayRange => dayRange(internalDayjs(from).startOf(value), internalDayjs(to).endOf(value)),
-  start: internalDayjs(from),
-  end: internalDayjs(to),
-});
-
 /**
  * Parse PostgreSQL-like interval string into object
  * E.g. '2 years 15 months 100 weeks 99 hours 15 seconds'
@@ -202,6 +183,70 @@ function alignToOrigin(startDate: dayjs.Dayjs, interval: ParsedInterval, origin:
   return alignedDate;
 }
 
+export const dayRange = (from: any, to: any, annotations?: Record<string, { granularity?: Granularity }>): DayRange => ({
+  by: (value: any) => {
+    const results = [];
+
+    let start = internalDayjs(from);
+    const end = internalDayjs(to);
+
+    while (start.startOf(value).isBefore(end) || start.isSame(end)) {
+      results.push(start);
+      start = start.add(1, value);
+    }
+
+    return results;
+  },
+  snapTo: (value: any): DayRange => {
+    // Check if this is a custom granularity
+    if (!isPredefinedGranularity(value) && annotations) {
+      // Try to find the custom granularity metadata
+      // The annotation key might be in format "Cube.dimension.granularity"
+      // So we need to search through all annotations
+      let customGranularity: Granularity | undefined;
+
+      for (const key of Object.keys(annotations)) {
+        if (key.endsWith(`.${value}`) && annotations[key].granularity) {
+          customGranularity = annotations[key].granularity;
+          break;
+        }
+      }
+
+      if (customGranularity?.interval) {
+        // For custom granularities, calculate the range for the bucket
+        const intervalParsed = parseSqlInterval(customGranularity.interval);
+        let intervalStart = internalDayjs(from);
+
+        // origin and offset are mutually exclusive
+        // If either is specified, align to it
+        if (customGranularity.origin || customGranularity.offset) {
+          let origin;
+          if (customGranularity.origin) {
+            // Absolute origin time
+            origin = internalDayjs(customGranularity.origin);
+          } else {
+            // offset is relative to start of year
+            origin = addInterval(internalDayjs().startOf('year'), parseSqlInterval(customGranularity.offset!));
+          }
+
+          // Align the value to the origin to find the actual bucket start
+          intervalStart = alignToOrigin(intervalStart, intervalParsed, origin);
+        }
+
+        // End is start + interval - 1 millisecond (to stay within the bucket)
+        const intervalEnd = addInterval(intervalStart, intervalParsed).subtract(1, 'millisecond');
+
+        return dayRange(intervalStart, intervalEnd, annotations);
+      }
+    }
+
+    // Default behavior for predefined granularities
+    return dayRange(internalDayjs(from).startOf(value), internalDayjs(to).endOf(value), annotations);
+  },
+  start: internalDayjs(from),
+  end: internalDayjs(to),
+});
+
 /**
  * Returns the time series points for the custom interval
  * TODO: It's almost a copy/paste of timeSeriesFromCustomInterval from

packages/cubejs-client-core/test/drill-down.test.ts

@@ -401,4 +401,242 @@ describe('drill down query', () => {
       timezone: 'UTC',
     });
   });
+
+  it('handles custom granularity with interval and origin', () => {
+    const EVALUATION_PERIOD = 5;
+    const LAST_EVALUATED_AT = '2020-08-01T00:00:00.000';
+
+    const customGranularityResponse = {
+      queryType: 'regularQuery',
+      results: [
+        {
+          query: {
+            measures: ['Transactions.count'],
+            timeDimensions: [
+              {
+                dimension: 'Transactions.createdAt',
+                granularity: 'alerting_monitor',
+                dateRange: ['2020-08-01T00:00:00.000', '2020-08-01T01:00:00.000'],
+              },
+            ],
+            filters: [],
+            timezone: 'UTC',
+            order: [],
+            dimensions: [],
+          },
+          data: [
+            {
+              'Transactions.createdAt.alerting_monitor': '2020-08-01T00:00:00.000',
+              'Transactions.createdAt': '2020-08-01T00:00:00.000',
+              'Transactions.count': 10,
+            },
+            {
+              'Transactions.createdAt.alerting_monitor': '2020-08-01T00:05:00.000',
+              'Transactions.createdAt': '2020-08-01T00:05:00.000',
+              'Transactions.count': 15,
+            },
+            {
+              'Transactions.createdAt.alerting_monitor': '2020-08-01T00:10:00.000',
+              'Transactions.createdAt': '2020-08-01T00:10:00.000',
+              'Transactions.count': 8,
+            },
+          ],
+          annotation: {
+            measures: {
+              'Transactions.count': {
+                title: 'Transactions Count',
+                shortTitle: 'Count',
+                type: 'number',
+                drillMembers: ['Transactions.id', 'Transactions.createdAt'],
+                drillMembersGrouped: {
+                  measures: [],
+                  dimensions: ['Transactions.id', 'Transactions.createdAt'],
+                },
+              },
+            },
+            dimensions: {},
+            segments: {},
+            timeDimensions: {
+              'Transactions.createdAt.alerting_monitor': {
+                title: 'Transaction created at',
+                shortTitle: 'Created at',
+                type: 'time',
+                granularity: {
+                  name: 'alerting_monitor',
+                  title: 'Alerting Monitor',
+                  interval: `${EVALUATION_PERIOD} minutes`,
+                  origin: LAST_EVALUATED_AT,
+                },
+              },
+              'Transactions.createdAt': {
+                title: 'Transaction created at',
+                shortTitle: 'Created at',
+                type: 'time',
+              },
+            },
+          },
+        },
+      ],
+      pivotQuery: {
+        measures: ['Transactions.count'],
+        timeDimensions: [
+          {
+            dimension: 'Transactions.createdAt',
+            granularity: 'alerting_monitor',
+            dateRange: ['2020-08-01T00:00:00.000', '2020-08-01T01:00:00.000'],
+          },
+        ],
+        filters: [],
+        timezone: 'UTC',
+        order: [],
+        dimensions: [],
+      },
+    };
+
+    const resultSet = new ResultSet(customGranularityResponse as any);
+
+    // Test drilling down on the second data point (00:05:00)
+    expect(
+      resultSet.drillDown({ xValues: ['2020-08-01T00:05:00.000'] })
+    ).toEqual({
+      measures: [],
+      segments: [],
+      dimensions: ['Transactions.id', 'Transactions.createdAt'],
+      filters: [
+        {
+          member: 'Transactions.count',
+          operator: 'measureFilter',
+        },
+      ],
+      timeDimensions: [
+        {
+          dimension: 'Transactions.createdAt',
+          // Should create a date range for the 5-minute interval starting at 00:05:00
+          dateRange: ['2020-08-01T00:05:00.000', '2020-08-01T00:09:59.999'],
+        },
+      ],
+      timezone: 'UTC',
+    });
+  });
+
+  it('handles custom granularity with non-aligned origin', () => {
+    const EVALUATION_PERIOD = 5;
+    const NON_ALIGNED_ORIGIN = '2020-08-01T00:02:00.000'; // Origin at 00:02 instead of 00:00
+
+    const customGranularityResponse = {
+      queryType: 'regularQuery',
+      results: [
+        {
+          query: {
+            measures: ['Transactions.count'],
+            timeDimensions: [
+              {
+                dimension: 'Transactions.createdAt',
+                granularity: 'alerting_monitor',
+                dateRange: ['2020-08-01T00:00:00.000', '2020-08-01T01:00:00.000'],
+              },
+            ],
+            filters: [],
+            timezone: 'UTC',
+            order: [],
+            dimensions: [],
+          },
+          data: [
+            {
+              // First bucket starts at 00:02:00 (origin)
+              'Transactions.createdAt.alerting_monitor': '2020-08-01T00:02:00.000',
+              'Transactions.createdAt': '2020-08-01T00:02:00.000',
+              'Transactions.count': 10,
+            },
+            {
+              // Second bucket starts at 00:07:00 (origin + 5 minutes)
+              'Transactions.createdAt.alerting_monitor': '2020-08-01T00:07:00.000',
+              'Transactions.createdAt': '2020-08-01T00:07:00.000',
+              'Transactions.count': 15,
+            },
+            {
+              // Third bucket starts at 00:12:00 (origin + 10 minutes)
+              'Transactions.createdAt.alerting_monitor': '2020-08-01T00:12:00.000',
+              'Transactions.createdAt': '2020-08-01T00:12:00.000',
+              'Transactions.count': 8,
+            },
+          ],
+          annotation: {
+            measures: {
+              'Transactions.count': {
+                title: 'Transactions Count',
+                shortTitle: 'Count',
+                type: 'number',
+                drillMembers: ['Transactions.id', 'Transactions.createdAt'],
+                drillMembersGrouped: {
+                  measures: [],
+                  dimensions: ['Transactions.id', 'Transactions.createdAt'],
+                },
+              },
+            },
+            dimensions: {},
+            segments: {},
+            timeDimensions: {
+              'Transactions.createdAt.alerting_monitor': {
+                title: 'Transaction created at',
+                shortTitle: 'Created at',
+                type: 'time',
+                granularity: {
+                  name: 'alerting_monitor',
+                  title: 'Alerting Monitor',
+                  interval: `${EVALUATION_PERIOD} minutes`,
+                  origin: NON_ALIGNED_ORIGIN,
+                },
+              },
+              'Transactions.createdAt': {
+                title: 'Transaction created at',
+                shortTitle: 'Created at',
+                type: 'time',
+              },
+            },
+          },
+        },
+      ],
+      pivotQuery: {
+        measures: ['Transactions.count'],
+        timeDimensions: [
+          {
+            dimension: 'Transactions.createdAt',
+            granularity: 'alerting_monitor',
+            dateRange: ['2020-08-01T00:00:00.000', '2020-08-01T01:00:00.000'],
+          },
+        ],
+        filters: [],
+        timezone: 'UTC',
+        order: [],
+        dimensions: [],
+      },
+    };
+
+    const resultSet = new ResultSet(customGranularityResponse as any);
+
+    // Test drilling down on the second data point (00:07:00)
+    // Since origin is 00:02:00, the bucket is 00:07:00 - 00:11:59
+    expect(
+      resultSet.drillDown({ xValues: ['2020-08-01T00:07:00.000'] })
+    ).toEqual({
+      measures: [],
+      segments: [],
+      dimensions: ['Transactions.id', 'Transactions.createdAt'],
+      filters: [
+        {
+          member: 'Transactions.count',
+          operator: 'measureFilter',
+        },
+      ],
+      timeDimensions: [
+        {
+          dimension: 'Transactions.createdAt',
+          // Should align to origin: bucket starts at 00:07:00 (origin + 5min)
+          dateRange: ['2020-08-01T00:07:00.000', '2020-08-01T00:11:59.999'],
+        },
+      ],
+      timezone: 'UTC',
+    });
+  });
 });