feat(api-gateway): Expose usedPreAggregations on data responses

Signed-off-by: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com>

Author: MazterQyou

docs-mintlify/api-reference/core-data.yaml

@@ -64,8 +64,9 @@ paths:
       description: >-
         Run a SQL query against the Cube [SQL API](/reference/core-data-apis/sql-api) and stream the
         results. The response is newline-delimited JSON: the first line carries the `schema` (column
-        names and types) and optionally `lastRefreshTime`; each subsequent line carries a `data`
-        chunk with one or more result rows.
+        names and types), optionally `lastRefreshTime`, and optionally `usedPreAggregations` naming
+        the pre-aggregations the result was served from; each subsequent line carries a `data` chunk
+        with one or more result rows.
       requestBody:
         required: true
         content:
@@ -127,6 +128,10 @@ paths:
                   lastRefreshTime:
                     type: string
                     format: date-time
+                  usedPreAggregations:
+                    type: object
+                    additionalProperties:
+                      type: object
                   data:
                     type: array
                     items:
@@ -140,6 +145,12 @@ paths:
                       - name: value
                         column_type: Int64
                     lastRefreshTime: '2025-01-13T12:00:00.000Z'
+                    usedPreAggregations:
+                      schema.orders_main20240101:
+                        preAggregationId: Orders.main
+                        targetTableName: schema.orders_main20240101_abc_def_1712
+                        lastUpdatedAt: 1712000000000
+                        type: rollup
                 dataLine:
                   summary: Data chunk
                   value:
@@ -672,6 +683,27 @@ components:
             type: object
         lastRefreshTime:
           type: string
+        external:
+          type: boolean
+          description: '`true` when the result was served from an external (CubeStore) pre-aggregation.'
+        usedPreAggregations:
+          type: object
+          description: >-
+            Pre-aggregations this result was served from, keyed by pre-aggregation table name. Each
+            entry carries `preAggregationId`, `targetTableName`, `lastUpdatedAt` and `type`. Absent
+            when the query hit no pre-aggregation. In dev mode and for the Playground the entries
+            also carry `refreshKeyValues`.
+          additionalProperties:
+            type: object
+            properties:
+              preAggregationId:
+                type: string
+              targetTableName:
+                type: string
+              lastUpdatedAt:
+                type: integer
+              type:
+                type: string
     V1LoadResultAnnotation:
       type: object
       required:

docs-mintlify/reference/core-data-apis/graphql-api/reference.mdx

@@ -31,6 +31,9 @@ object with query metadata:
   `annotation` field of the [REST API][ref-rest-api] response.
 - `lastRefreshTime` — an ISO 8601 timestamp indicating when the data was last
   refreshed.
+- `usedPreAggregations` — the pre-aggregations the result was served from, the
+  same as the `usedPreAggregations` field of the [REST API][ref-rest-api]
+  response. Absent when the query hit no pre-aggregation.
 
 ## `CubeQueryArgs`
 

docs-mintlify/reference/core-data-apis/rest-api/reference.mdx

@@ -41,6 +41,11 @@ Response
 - `lastRefreshTime` - An ISO 8601 timestamp indicating when the data was last refreshed.
 - `external` - A boolean indicating whether the query was served from a
   pre-aggregation in Cube Store. Present only when the query hit a pre-aggregation.
+- `usedPreAggregations` - The pre-aggregations this result was served from, keyed
+  by pre-aggregation table name. Each entry carries `preAggregationId` (the name in
+  the data model, e.g. `Orders.main`, stable across rebuilds), `targetTableName`
+  (the table of one specific build), `lastUpdatedAt` and `type`. Absent when the
+  query hit no pre-aggregation.
 - `total` - The total number of rows returned for the query. Useful for
   paginating results.
 
@@ -420,8 +425,10 @@ This endpoint is part of the [SQL API][ref-sql-api].
 | `x-request-id` | Custom request identifier. When provided, this ID is used to track the query through the system and can be used to [cancel the query](#base_path/v1/running-query/requestid). If not provided, a unique ID is generated automatically. | ❌ No |
 
 Response: a stream of newline-delimited JSON objects. The first object contains
-the `schema` property with column names and types, and optionally
-`lastRefreshTime` indicating when the data was last refreshed.
+the `schema` property with column names and types, optionally
+`lastRefreshTime` indicating when the data was last refreshed, and optionally
+`usedPreAggregations` naming the pre-aggregations the result was served from
+(same shape as in the [JSON query](#base_path/v1/load) response).
 The following objects contain chunks of the result set under the `data` property.
 Each chunk includes one or more rows of the result set; the maximum number of rows
 per chunk is set by

docs-mintlify/scripts/extract-core-data.js

@@ -54,7 +54,9 @@ const EXTRA_PATHS = {
       description:
         'Run a SQL query against the Cube [SQL API](/reference/core-data-apis/sql-api) ' +
         'and stream the results. The response is newline-delimited JSON: the first line ' +
-        'carries the `schema` (column names and types) and optionally `lastRefreshTime`; ' +
+        'carries the `schema` (column names and types), optionally `lastRefreshTime`, and ' +
+        'optionally `usedPreAggregations` naming the pre-aggregations the result was served ' +
+        'from; ' +
         'each subsequent line carries a `data` chunk with one or more result rows.',
       requestBody: {
         required: true,
@@ -112,6 +114,7 @@ const EXTRA_PATHS = {
                     },
                   },
                   lastRefreshTime: { type: 'string', format: 'date-time' },
+                  usedPreAggregations: { type: 'object', additionalProperties: { type: 'object' } },
                   data: { type: 'array', items: { type: 'array', items: {} } },
                 },
               },
@@ -121,6 +124,14 @@ const EXTRA_PATHS = {
                   value: {
                     schema: [{ name: 'value', column_type: 'Int64' }],
                     lastRefreshTime: '2025-01-13T12:00:00.000Z',
+                    usedPreAggregations: {
+                      'schema.orders_main20240101': {
+                        preAggregationId: 'Orders.main',
+                        targetTableName: 'schema.orders_main20240101_abc_def_1712',
+                        lastUpdatedAt: 1712000000000,
+                        type: 'rollup',
+                      },
+                    },
                   },
                 },
                 dataLine: { summary: 'Data chunk', value: { data: [['123']] } },

packages/cubejs-api-gateway/openspec.yml

@@ -462,6 +462,23 @@ components:
             type: "object"
         lastRefreshTime:
           type: "string"
+        external:
+          type: "boolean"
+          description: "`true` when the result was served from an external (CubeStore) pre-aggregation."
+        usedPreAggregations:
+          type: "object"
+          description: "Pre-aggregations this result was served from, keyed by pre-aggregation table name. Each entry carries `preAggregationId`, `targetTableName`, `lastUpdatedAt` and `type`. Absent when the query hit no pre-aggregation. In dev mode and for the Playground the entries also carry `refreshKeyValues`."
+          additionalProperties:
+            type: "object"
+            properties:
+              preAggregationId:
+                type: "string"
+              targetTableName:
+                type: "string"
+              lastUpdatedAt:
+                type: "integer"
+              type:
+                type: "string"
     V1Error:
       type: "object"
       required:

packages/cubejs-api-gateway/src/gateway.ts

@@ -132,6 +132,44 @@ function systemAsyncHandler(handler: (req: Request & { context: ExtendedRequestC
 
 const DEV_TOKEN_SCOPE = 'dev-token';
 
+/**
+ * Fields of `usedPreAggregations` that are safe to report to any client: the
+ * identity of the pre-aggregation a result was served from, so the client can
+ * match the result to a build it is watching.
+ *
+ * `refreshKeyValues` is deliberately left out. Those are raw rows of the
+ * refresh key queries - typically aggregates such as `MAX(updated_at)` or
+ * `COUNT(*)` - and a `refreshKey.sql` is often written without the security
+ * context filtering that the cube itself applies, so the values can describe
+ * data the caller cannot otherwise reach. The full object, including them, is
+ * still returned in dev mode and to the Playground.
+ */
+function publicUsedPreAggregations(
+  usedPreAggregations: Record<string, any> | undefined
+): Record<string, any> | undefined {
+  // A query that hit no pre-aggregation reports nothing rather than an empty
+  // object, so the key is simply absent from the response.
+  if (!usedPreAggregations || Object.keys(usedPreAggregations).length === 0) {
+    return undefined;
+  }
+
+  const publicFields = ['preAggregationId', 'targetTableName', 'lastUpdatedAt', 'type'];
+
+  return Object.fromEntries(
+    Object.entries(usedPreAggregations).map(([tableName, usage]) => [
+      tableName,
+      // Undefined fields are dropped rather than kept: the native result
+      // pipeline deserializes a JS `undefined` into a JSON `null`, so leaving
+      // them in would put `"preAggregationId": null` on the wire.
+      Object.fromEntries(
+        publicFields
+          .filter((field) => usage?.[field] !== undefined)
+          .map((field) => [field, usage[field]])
+      ),
+    ])
+  );
+}
+
 function hasDevTokenScope(securityContext: unknown): boolean {
   if (typeof securityContext !== 'object' || securityContext === null) {
     return false;
@@ -1945,6 +1983,10 @@ class ApiGateway {
     const resObj = {
       query: normalizedQuery,
       lastRefreshTime: response.lastRefreshTime?.toISOString(),
+      // Identity of the pre-aggregations behind this result, so a client can
+      // join it to the build it is waiting on. The dev-mode block below
+      // replaces it with the unredacted object.
+      usedPreAggregations: publicUsedPreAggregations(response.usedPreAggregations),
       ...(
         getEnv('devMode') ||
           context.signedWithPlaygroundAuthSecret
@@ -2242,6 +2284,11 @@ class ApiGateway {
             // otherwise the SQL API reports "unknown" for every query cubesql
             // hands over as pre-generated SQL.
             lastRefreshTime: response.lastRefreshTime?.toISOString(),
+            // Same reason as `lastRefreshTime` above: the pre-aggregation
+            // identity has to travel with the pushed-down result too, or a
+            // cubesql query that goes through pre-generated SQL can never tell
+            // the client which pre-aggregation it read.
+            usedPreAggregations: publicUsedPreAggregations(response.usedPreAggregations),
             // Always false: this branch builds its sqlQuery with
             // `disableExternalPreAggregations` set (above), which makes
             // `externalPreAggregationQuery()` return false, and the

packages/cubejs-api-gateway/src/graphql.ts

@@ -675,6 +675,7 @@ export function makeSchema(metaConfig: any): GraphQLSchema {
           res.extensions = {
             annotation: results.annotation,
             lastRefreshTime: results.lastRefreshTime,
+            usedPreAggregations: results.usedPreAggregations,
           };
 
           return results.data.map(entry => R.toPairs(entry)

packages/cubejs-api-gateway/test/index.test.ts

@@ -1508,7 +1508,7 @@ describe('API Gateway', () => {
   describe('external pre-aggregation indicator', () => {
     // Helper mock that lets a test pretend the query orchestrator served
     // the result from an external (CubeStore) pre-aggregation, optionally
-    // with the dev-only `usedPreAggregations` object as well.
+    // with the `usedPreAggregations` object as well.
     class AdapterApiMockWithFlags extends AdapterApiMock {
       public constructor(
         private readonly external: boolean | undefined,
@@ -1536,16 +1536,23 @@ describe('API Gateway', () => {
         .expect(200);
 
       expect(res.body.external).toBe(false);
-      // Full pre-agg object stays dev/playground-only.
+      // No pre-aggregation was used, so there is nothing to name.
       expect(res.body.usedPreAggregations).toBeUndefined();
     });
 
-    test('external=true when query was served from an external pre-aggregation (no leak of names)', async () => {
+    // Pre-aggregation identity is reported to every API consumer so a client
+    // can match a result to the build it is waiting on. Refresh key values are
+    // not: they are rows of the refresh key queries, which are often written
+    // without the security context filtering the cube itself applies.
+    test('external=true exposes pre-aggregation identity without refresh key values', async () => {
       const { app } = await createApiGateway(
         new AdapterApiMockWithFlags(true, {
           'Foo.fooMain': {
+            preAggregationId: 'Foo.fooMain',
             targetTableName: 'stb_pre_aggs.foo_foo_main',
+            lastUpdatedAt: 1712000000000,
             type: 'rollup',
+            refreshKeyValues: [[{ max_updated_at: '2024-01-01T00:00:00.000Z' }]],
           },
         }),
       );
@@ -1556,16 +1563,24 @@ describe('API Gateway', () => {
         .expect(200);
 
       expect(res.body.external).toBe(true);
-      // Pre-aggregation names / table names must NOT be exposed to ordinary
-      // API consumers — only the boolean flag is safe.
-      expect(res.body.usedPreAggregations).toBeUndefined();
+      expect(res.body.usedPreAggregations).toEqual({
+        'Foo.fooMain': {
+          preAggregationId: 'Foo.fooMain',
+          targetTableName: 'stb_pre_aggs.foo_foo_main',
+          lastUpdatedAt: 1712000000000,
+          type: 'rollup',
+        },
+      });
     });
 
-    test('usedPreAggregations is exposed under playground auth alongside external', async () => {
+    test('refresh key values are exposed under playground auth alongside external', async () => {
       const usedPreAggregations = {
         'Foo.fooMain': {
+          preAggregationId: 'Foo.fooMain',
           targetTableName: 'stb_pre_aggs.foo_foo_main',
+          lastUpdatedAt: 1712000000000,
           type: 'rollup',
+          refreshKeyValues: [[{ max_updated_at: '2024-01-01T00:00:00.000Z' }]],
         },
       };
       const { app } = await createApiGateway(

packages/cubejs-api-gateway/test/sql-api-load.test.ts

@@ -5,12 +5,27 @@ const logger = (type: any, message: any) => console.log({ type, ...message });
 
 const LAST_REFRESH_TIME = new Date('2024-01-01T00:00:00.000Z');
 
+// Shape the orchestrator reports: identity plus the refresh key values, which
+// must not reach a regular client.
+const USED_PRE_AGGREGATIONS = {
+  'schema.orders_main20240101': {
+    preAggregationId: 'Orders.main',
+    targetTableName: 'schema.orders_main20240101_abc_def_1712',
+    lastUpdatedAt: 1712000000000,
+    type: 'rollup',
+    refreshKeyValues: [[{ max_updated_at: '2024-01-01T00:00:00.000Z' }]],
+  },
+};
+
 class FreshnessAdapterApiMock extends AdapterApiMock {
   public lastRefreshTime: Date | undefined;
 
-  public constructor(lastRefreshTime?: Date) {
+  public usedPreAggregations: Record<string, any> | undefined;
+
+  public constructor(lastRefreshTime?: Date, usedPreAggregations?: Record<string, any>) {
     super();
     this.lastRefreshTime = lastRefreshTime;
+    this.usedPreAggregations = usedPreAggregations;
   }
 
   public async executeQuery(_query: any) {
@@ -20,6 +35,7 @@ class FreshnessAdapterApiMock extends AdapterApiMock {
       // Always falsy for a pushdown query — see the note in
       // `sqlApiLoad`. Mirrors what the orchestrator actually echoes back.
       external: false,
+      usedPreAggregations: this.usedPreAggregations,
     };
   }
 }
@@ -88,4 +104,61 @@ describe('sqlApiLoad freshness metadata', () => {
     expect(response.results).toHaveLength(1);
     expect(response.results[0].lastRefreshTime).toBeUndefined();
   });
+
+  // Pre-aggregation identity lets a client join a chart's result to the build
+  // it is waiting on. It has to travel with the pushed-down result too, since
+  // that is the branch every cubesql query takes.
+  test('pushed-down sqlQuery result carries usedPreAggregations', async () => {
+    const response = await sqlApiLoad(
+      new FreshnessAdapterApiMock(LAST_REFRESH_TIME, USED_PRE_AGGREGATIONS),
+      ['SELECT * FROM test', []]
+    );
+
+    expect(response.results[0].usedPreAggregations).toEqual({
+      'schema.orders_main20240101': {
+        preAggregationId: 'Orders.main',
+        targetTableName: 'schema.orders_main20240101_abc_def_1712',
+        lastUpdatedAt: 1712000000000,
+        type: 'rollup',
+      },
+    });
+  });
+
+  // Refresh key values are rows of the refresh key queries, and a
+  // `refreshKey.sql` is often not filtered by the security context the cube
+  // itself applies. They stay dev-mode only.
+  test('usedPreAggregations omits refreshKeyValues', async () => {
+    const response = await sqlApiLoad(
+      new FreshnessAdapterApiMock(LAST_REFRESH_TIME, USED_PRE_AGGREGATIONS),
+      ['SELECT * FROM test', []]
+    );
+
+    const usage = response.results[0].usedPreAggregations['schema.orders_main20240101'];
+    expect(usage.refreshKeyValues).toBeUndefined();
+  });
+
+  // A query that hit no pre-aggregation reports nothing rather than `{}`.
+  test('usedPreAggregations is absent when no pre-aggregation was used', async () => {
+    const response = await sqlApiLoad(
+      new FreshnessAdapterApiMock(LAST_REFRESH_TIME, {}),
+      ['SELECT * FROM test', []]
+    );
+
+    expect(response.results[0].usedPreAggregations).toBeUndefined();
+  });
+
+  // The non-pushdown branch goes through `prepareResultTransformData`, which is
+  // also what the REST `/load` response is built from.
+  test('regular query result carries usedPreAggregations', async () => {
+    const response = await sqlApiLoad(
+      new FreshnessAdapterApiMock(LAST_REFRESH_TIME, USED_PRE_AGGREGATIONS)
+    );
+
+    const result = response.getResults()[0].getRootResultObject()[0];
+
+    expect(result.usedPreAggregations['schema.orders_main20240101'].preAggregationId)
+      .toBe('Orders.main');
+    expect(result.usedPreAggregations['schema.orders_main20240101'].refreshKeyValues)
+      .toBeUndefined();
+  });
 });