feat(query-orchestrator): compute interval refresh keys locally instead of querying

docs-mintlify/reference/configuration/environment-variables.mdx

@@ -1417,6 +1417,32 @@ this variable.
 | ------------------------- | ---------------------- | --------------------- |
 | A valid number in seconds | `600`                  | `600`                 |
 
+## `CUBEJS_REFRESH_KEY_LOCAL_TIME`
+
+To check an `every`-based [`refresh_key`](/reference/data-modeling/cube#refresh_key), Cube
+asks the database what time it is. If `true`, it uses its own clock instead, saving a round
+trip on nearly every request.
+
+| Possible Values | Default in Development | Default in Production |
+| --------------- | ---------------------- | --------------------- |
+| `true`, `false` | `false`                | `false`               |
+
+Refresh keys written with `sql` or marked
+[`incremental`](/reference/data-modeling/pre-aggregations#incremental) keep querying the
+database, since their values depend on your data and not just on the time. So do all
+refresh keys if you have set
+[`queryCacheOptions.refreshKeyRenewalThreshold`][ref-config-query-cache-options], which
+works by caching the query result.
+
+<Warning>
+
+Cube now trusts each machine's clock, so they all have to agree. Set this variable to the
+same value on **every** API instance and refresh worker and keep them synchronized with
+NTP. Instances on opposite sides of an interval boundary compute different refresh keys,
+which can build the same pre-aggregation twice.
+
+</Warning>
+
 ## `CUBEJS_REFRESH_WORKER`
 
 If `true`, this instance of Cube will **only** refresh pre-aggregations.
@@ -2307,6 +2333,7 @@ The port for a Cube deployment to listen to API connections on.
   https://motherduck.com/docs/authenticating-to-motherduck/#authentication-using-a-service-token
 [ref-config-db]: /admin/connect-to-data/data-sources
 [ref-config-driver-factory]: /reference/configuration/config#driver_factory
+[ref-config-query-cache-options]: /reference/configuration/config#orchestrator_options
 [ref-config-multiple-ds-cloud]: /admin/connect-to-data/multiple-data-sources
 [ref-config-multiple-ds-decorating-env]: /admin/connect-to-data/multiple-data-sources#under-the-hood
 [ref-preagg-data-source]: /docs/pre-aggregations/refreshing-pre-aggregations#pre-aggregation-data-source

packages/cubejs-backend-shared/src/env.ts

@@ -368,6 +368,9 @@ const variables: Record<string, (...args: any) => any> = {
   preciseDecimalInCubestore: () => get('CUBEJS_DB_PRECISE_DECIMAL_IN_CUBESTORE')
     .default('false')
     .asBoolStrict(),
+  refreshKeyLocalTime: () => get('CUBEJS_REFRESH_KEY_LOCAL_TIME')
+    .default('false')
+    .asBoolStrict(),
 
   /** ****************************************************************
    * Common db options                                               *

packages/cubejs-backend-shared/test/env.test.ts

@@ -107,6 +107,25 @@ describe('getEnv', () => {
     ).toBe(10);
   });
 
+  test('refreshKeyLocalTime', () => {
+    delete process.env.CUBEJS_REFRESH_KEY_LOCAL_TIME;
+    expect(getEnv('refreshKeyLocalTime')).toBe(false);
+
+    process.env.CUBEJS_REFRESH_KEY_LOCAL_TIME = 'true';
+    expect(getEnv('refreshKeyLocalTime')).toBe(true);
+
+    process.env.CUBEJS_REFRESH_KEY_LOCAL_TIME = 'false';
+    expect(getEnv('refreshKeyLocalTime')).toBe(false);
+  });
+
+  test('refreshKeyLocalTime(exception)', () => {
+    process.env.CUBEJS_REFRESH_KEY_LOCAL_TIME = 'yes';
+
+    expect(() => getEnv('refreshKeyLocalTime')).toThrowError();
+
+    delete process.env.CUBEJS_REFRESH_KEY_LOCAL_TIME;
+  });
+
   test('livePreview', () => {
     expect(getEnv('livePreview')).toBe(true);
 

packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts

@@ -25,7 +25,12 @@ import { ContinueWaitError } from './ContinueWaitError';
 import { LocalCacheDriver } from './LocalCacheDriver';
 import { DriverFactory, DriverFactoryByDataSource } from './DriverFactory';
 import { LoadPreAggregationResult, PreAggregationDescription } from './PreAggregations';
-import { getCacheHash, extractRequestUUID } from './utils';
+import {
+  getCacheHash,
+  extractRequestUUID,
+  evaluateLocalRefreshKey,
+  isValidLocalRefreshKey,
+} from './utils';
 import { CacheAndQueryDriverType, MetadataOperationType } from './QueryOrchestrator';
 
 export type CacheQueryResultOptions = {
@@ -52,12 +57,28 @@ export type CacheQueryResultOptions = {
 export type RefreshKeyCacheOptions =
   Pick<CacheQueryResultOptions, 'priority' | 'requestId' | 'waitForRenew' | 'dataSource'>;
 
+/**
+ * Everything needed to evaluate an `every` based refreshKey without touching a
+ * database: `FLOOR((utcOffset + unixTimestamp - dayOffset) / interval)`.
+ *
+ * `utcOffset` is frozen at compile time, exactly as it already is when baked into
+ * the emitted SQL string — the orchestrator must never recompute it, or the local
+ * and SQL paths would stop agreeing across a DST transition.
+ */
+export type LocalRefreshKeyDescriptor = {
+  interval: number;
+  utcOffset: number;
+  dayOffset: number;
+  cron?: boolean;
+};
+
 type QueryOptions = {
   external?: boolean;
   renewalThreshold?: number;
   updateWindowSeconds?: number;
   renewalThresholdOutsideUpdateWindow?: number;
   incremental?: boolean;
+  localRefreshKey?: LocalRefreshKeyDescriptor;
 };
 
 export type QueryWithParams = [
@@ -182,6 +203,10 @@ export class QueryCache {
 
   protected static readonly IN_MEMORY_CACHE_DISABLE_PERIOD = 5 * 60 * 1000;
 
+  protected readonly localRefreshKeyEnabled: boolean;
+
+  protected localRefreshKeyLogged: boolean = false;
+
   public constructor(
     protected readonly cachePrefix: string,
     protected readonly driverFactory: DriverFactoryByDataSource,
@@ -208,6 +233,49 @@ export class QueryCache {
     this.memoryCache = new LRUCache<string, CacheEntry>({
       max: options.maxInMemoryCacheEntries || 10000
     });
+
+    // Read from the environment rather than an option: the emitting half in the schema
+    // compiler is wired straight from the same variable, and an option here could be set
+    // on its own, turning the feature into a no-op where no descriptor is ever emitted.
+    this.localRefreshKeyEnabled = getEnv('refreshKeyLocalTime');
+  }
+
+  public localRefreshKeyResult(queryOptions?: QueryOptions): [{ refresh_key: number }] | null {
+    if (!this.localRefreshKeyEnabled || !isValidLocalRefreshKey(queryOptions?.localRefreshKey)) {
+      return null;
+    }
+
+    // `refreshKeyRenewalThreshold` throttles how often the SQL result is re-read, and that is
+    // also what bounds how often the key advances: a value cached for a day advances daily,
+    // whatever `every` says. A locally evaluated key has no cache entry to age out, so the only
+    // way to keep honouring the override is to leave these keys on the SQL path.
+    if (this.options.refreshKeyRenewalThreshold) {
+      this.logLocalRefreshKeyOnce('Local refresh key evaluation skipped', {
+        warning: 'refreshKeyRenewalThreshold is set, which throttles refresh key advancement. ' +
+          'Interval based refresh keys keep running as queries. Unset it to evaluate them locally.',
+        refreshKeyRenewalThreshold: this.options.refreshKeyRenewalThreshold,
+      });
+
+      return null;
+    }
+
+    this.logLocalRefreshKeyOnce('Local refresh key evaluation enabled', {
+      warning: 'Interval based refresh keys are evaluated from this instance clock. ' +
+        'Clocks must be in sync across all API instances and refresh workers.',
+    });
+
+    return evaluateLocalRefreshKey(<LocalRefreshKeyDescriptor>queryOptions?.localRefreshKey);
+  }
+
+  // Logged on first use rather than from the constructor, where subclass field
+  // initialization has not run yet and `this.logger` may not be usable.
+  private logLocalRefreshKeyOnce(message: string, meta: Record<string, unknown>) {
+    if (this.localRefreshKeyLogged) {
+      return;
+    }
+
+    this.localRefreshKeyLogged = true;
+    this.logger(message, meta);
   }
 
   public getCacheDriver(): CacheDriverInterface {
@@ -462,6 +530,13 @@ export class QueryCache {
     options: RefreshKeyCacheOptions,
   ) {
     const [query, values, queryOptions] = sqlQuery;
+
+    // A locally evaluated key is free: nothing to cache, no queue to wait on.
+    const local = this.localRefreshKeyResult(queryOptions);
+    if (local) {
+      return local;
+    }
+
     const cacheKey = QueryCache.refreshKeyIdentity(sqlQuery, options.dataSource);
 
     return this.cacheQueryResult(query, values, cacheKey, expiration, {

packages/cubejs-query-orchestrator/src/orchestrator/utils.ts

@@ -2,7 +2,7 @@ import crypto from 'crypto';
 
 import { getProcessUid } from '@cubejs-backend/shared';
 import { QueryKey, QueryKeyHash } from '@cubejs-backend/base-driver';
-import { CacheKey } from './QueryCache';
+import { CacheKey, LocalRefreshKeyDescriptor } from './QueryCache';
 
 /**
  * Unique process ID regexp.
@@ -32,6 +32,28 @@ export function getCacheHash(queryKey: QueryKey | CacheKey, processUid?: string)
   }
 }
 
+/**
+ * `nowMs` is not floored to whole seconds first: for integer `x`, fractional
+ * `f` in [0, 1) and `interval >= 1`, floor((x + f) / interval) === floor(x / interval),
+ * which is also why the fractional seconds of `EXTRACT(EPOCH FROM NOW())` never
+ * mattered on the SQL path.
+ */
+export function evaluateLocalRefreshKey(
+  descriptor: LocalRefreshKeyDescriptor,
+  nowMs: number = Date.now(),
+): [{ refresh_key: number }] {
+  const { utcOffset, interval, dayOffset } = descriptor;
+
+  return [{ refresh_key: Math.floor((utcOffset + nowMs / 1000 - dayOffset) / interval) }];
+}
+
+export function isValidLocalRefreshKey(descriptor?: LocalRefreshKeyDescriptor): boolean {
+  return !!descriptor &&
+    Number.isFinite(descriptor.interval) && descriptor.interval > 0 &&
+    Number.isFinite(descriptor.utcOffset) &&
+    Number.isFinite(descriptor.dayOffset);
+}
+
 /**
  * Extracts the UUID prefix from a request ID by stripping the `-span-N` suffix.
  */

packages/cubejs-query-orchestrator/test/unit/PreAggregations.test.ts

@@ -467,6 +467,128 @@ describe('PreAggregations', () => {
     });
   });
 
+  describe('local refresh key', () => {
+    const REFRESH_KEY_SQL = 'SELECT FLOOR((UNIX_TIMESTAMP()) / 600) as refresh_key';
+    const descriptor = { interval: 600, utcOffset: 0, dayOffset: 0, cron: false };
+
+    // The flag is read from the environment in the constructor, so it has to be toggled
+    // around construction rather than passed in as an option.
+    const newQueryCache = (localRefreshKey?: boolean) => {
+      const previous = process.env.CUBEJS_REFRESH_KEY_LOCAL_TIME;
+      if (localRefreshKey === undefined) {
+        delete process.env.CUBEJS_REFRESH_KEY_LOCAL_TIME;
+      } else {
+        process.env.CUBEJS_REFRESH_KEY_LOCAL_TIME = String(localRefreshKey);
+      }
+
+      try {
+        return new QueryCache(
+          'TEST',
+          mockDriverFactory as any,
+          // eslint-disable-next-line @typescript-eslint/no-empty-function
+          () => {},
+          {
+            cacheAndQueueDriver: 'memory',
+            queueOptions: async () => ({ executionTimeout: 1, concurrency: 2 }),
+          },
+        );
+      } finally {
+        if (previous === undefined) {
+          delete process.env.CUBEJS_REFRESH_KEY_LOCAL_TIME;
+        } else {
+          process.env.CUBEJS_REFRESH_KEY_LOCAL_TIME = previous;
+        }
+      }
+    };
+
+    const newLoadCache = (localRefreshKey?: boolean) => {
+      const cache = newQueryCache(localRefreshKey);
+      (cache.getCacheDriver() as LocalCacheDriver).reset();
+
+      const preAggregations = new PreAggregations(
+        'TEST',
+        mockDriverFactory as any,
+        // eslint-disable-next-line @typescript-eslint/no-empty-function
+        () => {},
+        cache,
+        { queueOptions: async () => ({ executionTimeout: 1, concurrency: 2 }) },
+      );
+
+      return new PreAggregationLoadCache(
+        mockDriverFactory as any,
+        cache,
+        preAggregations,
+        { dataSource: 'default' },
+      );
+    };
+
+    test('keyQueryResult evaluates locally without querying the datasource', async () => {
+      const loadCache = newLoadCache(true);
+
+      const result = await loadCache.keyQueryResult(
+        [REFRESH_KEY_SQL, [], { external: true, renewalThreshold: 60, localRefreshKey: descriptor }],
+        false,
+        10,
+      );
+
+      expect(result).toEqual([{ refresh_key: Math.floor(Date.now() / 1000 / descriptor.interval) }]);
+      expect(mockDriver!.executedQueries).toEqual([]);
+    });
+
+    test('keyQueryResult still queries when the flag is off', async () => {
+      const loadCache = newLoadCache(false);
+
+      await loadCache.keyQueryResult(
+        [REFRESH_KEY_SQL, [], { external: false, renewalThreshold: 60, localRefreshKey: descriptor }],
+        false,
+        10,
+      );
+
+      expect(mockDriver!.executedQueries).toEqual([REFRESH_KEY_SQL]);
+    });
+
+    test('keyQueryResult still queries an incremental key that carries no descriptor', async () => {
+      const loadCache = newLoadCache(true);
+      const incrementalSql = 'SELECT CASE WHEN NOW() < $1 THEN FLOOR((UNIX_TIMESTAMP()) / 3600) END as refresh_key';
+
+      await loadCache.keyQueryResult(
+        [incrementalSql, [], {
+          external: false,
+          renewalThreshold: 300,
+          incremental: true,
+          renewalThresholdOutsideUpdateWindow: 86400,
+        }],
+        false,
+        10,
+      );
+
+      expect(mockDriver!.executedQueries).toEqual([incrementalSql]);
+    });
+
+    // A single load reads the invalidation keys several times (contentVersion, the returned
+    // refreshKeyValues, the refresh queue key). If the clock were re-read, a load crossing an
+    // interval boundary would look a table up under one content version and enqueue it under
+    // another.
+    test('keyQueryResult is stable across an interval boundary within one load cache', async () => {
+      const loadCache = newLoadCache(true);
+      const key: [string, any[], Record<string, any>] =
+        [REFRESH_KEY_SQL, [], { external: true, renewalThreshold: 60, localRefreshKey: descriptor }];
+
+      const nowSpy = jest.spyOn(Date, 'now').mockReturnValue(600_000);
+      try {
+        const first = await loadCache.keyQueryResult(key, false, 10);
+        expect(first).toEqual([{ refresh_key: 1 }]);
+
+        nowSpy.mockReturnValue(1_200_000);
+        const second = await loadCache.keyQueryResult(key, false, 10);
+
+        expect(second).toEqual(first);
+      } finally {
+        nowSpy.mockRestore();
+      }
+    });
+  });
+
   describe('loadAllPreAggregationsIfNeeded', () => {
     let preAggregations: PreAggregations | null = null;