diff --git a/docs-mintlify/api-reference/core-data.yaml b/docs-mintlify/api-reference/core-data.yaml index 4dedf92766295..5adefa63857cd 100644 --- a/docs-mintlify/api-reference/core-data.yaml +++ b/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: @@ -109,7 +110,8 @@ paths: '200': description: >- Newline-delimited JSON stream. The first line carries `schema` (and optionally - `lastRefreshTime`); subsequent lines carry `data` chunks. + `lastRefreshTime`, `external` and `usedPreAggregations`); subsequent lines carry `data` + chunks. content: application/json: schema: @@ -127,6 +129,19 @@ paths: lastRefreshTime: type: string format: date-time + external: + type: boolean + usedPreAggregations: + type: object + additionalProperties: + type: object + properties: + preAggregationId: + type: string + lastUpdatedAt: + type: integer + type: + type: string data: type: array items: @@ -140,6 +155,11 @@ paths: - name: value column_type: Int64 lastRefreshTime: '2025-01-13T12:00:00.000Z' + usedPreAggregations: + schema.orders_main20240101: + preAggregationId: Orders.main + lastUpdatedAt: 1712000000000 + type: rollup dataLine: summary: Data chunk value: @@ -672,6 +692,25 @@ 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`, `lastUpdatedAt` and `type`. Absent when the query hit + no pre-aggregation. In dev mode and for the Playground the entries also carry + `targetTableName` and `refreshKeyValues`. + additionalProperties: + type: object + properties: + preAggregationId: + type: string + lastUpdatedAt: + type: integer + type: + type: string V1LoadResultAnnotation: type: object required: diff --git a/docs-mintlify/reference/core-data-apis/graphql-api/reference.mdx b/docs-mintlify/reference/core-data-apis/graphql-api/reference.mdx index eba1746674fcb..fc9bc9743592a 100644 --- a/docs-mintlify/reference/core-data-apis/graphql-api/reference.mdx +++ b/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` diff --git a/docs-mintlify/reference/core-data-apis/rest-api/reference.mdx b/docs-mintlify/reference/core-data-apis/rest-api/reference.mdx index 04e6466bf06af..39a8619e9d413 100644 --- a/docs-mintlify/reference/core-data-apis/rest-api/reference.mdx +++ b/docs-mintlify/reference/core-data-apis/rest-api/reference.mdx @@ -40,7 +40,26 @@ Response - `type` - Data type - `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. + pre-aggregation in Cube Store. +- `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), `lastUpdatedAt` and + `type`. Absent when the query hit no pre-aggregation. + + A `rollupJoin` or `rollupLambda` pre-aggregation is reported as the rollups it + references, not as itself, because that is what the query actually reads. Look + for those rollups rather than for the joining or lambda pre-aggregation. + + + + These fields name your pre-aggregations and are returned to every API consumer - + including browsers holding a public embedding token. Fields that describe + storage layout or query results stay behind + [`CUBEJS_DEV_MODE`](/reference/configuration/environment-variables#cubejs_dev_mode) + and the Playground: `targetTableName` (the physical table of one specific build, + version hashes included) and `refreshKeyValues`. + + - `total` - The total number of rows returned for the query. Useful for paginating results. @@ -420,8 +439,12 @@ 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). Neither field is +reported for queries served in [streaming mode](/reference/core-data-apis/sql-api#streaming), +which records no result metadata. 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 diff --git a/docs-mintlify/scripts/extract-core-data.js b/docs-mintlify/scripts/extract-core-data.js index a1231438e7016..45928d08240a1 100644 --- a/docs-mintlify/scripts/extract-core-data.js +++ b/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, @@ -95,7 +97,8 @@ const EXTRA_PATHS = { '200': { description: 'Newline-delimited JSON stream. The first line carries `schema` (and optionally ' + - '`lastRefreshTime`); subsequent lines carry `data` chunks.', + '`lastRefreshTime`, `external` and `usedPreAggregations`); subsequent lines carry ' + + '`data` chunks.', content: { 'application/json': { schema: { @@ -112,6 +115,18 @@ const EXTRA_PATHS = { }, }, lastRefreshTime: { type: 'string', format: 'date-time' }, + external: { type: 'boolean' }, + usedPreAggregations: { + type: 'object', + additionalProperties: { + type: 'object', + properties: { + preAggregationId: { type: 'string' }, + lastUpdatedAt: { type: 'integer' }, + type: { type: 'string' }, + }, + }, + }, data: { type: 'array', items: { type: 'array', items: {} } }, }, }, @@ -121,6 +136,13 @@ const EXTRA_PATHS = { value: { schema: [{ name: 'value', column_type: 'Int64' }], lastRefreshTime: '2025-01-13T12:00:00.000Z', + usedPreAggregations: { + 'schema.orders_main20240101': { + preAggregationId: 'Orders.main', + lastUpdatedAt: 1712000000000, + type: 'rollup', + }, + }, }, }, dataLine: { summary: 'Data chunk', value: { data: [['123']] } }, diff --git a/packages/cubejs-api-gateway/openspec.yml b/packages/cubejs-api-gateway/openspec.yml index 0a5c71ddddc04..870005283f393 100644 --- a/packages/cubejs-api-gateway/openspec.yml +++ b/packages/cubejs-api-gateway/openspec.yml @@ -462,6 +462,21 @@ 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`, `lastUpdatedAt` and `type`. Absent when the query hit no pre-aggregation. In dev mode and for the Playground the entries also carry `targetTableName` and `refreshKeyValues`." + additionalProperties: + type: "object" + properties: + preAggregationId: + type: "string" + lastUpdatedAt: + type: "integer" + type: + type: "string" V1Error: type: "object" required: diff --git a/packages/cubejs-api-gateway/src/gateway.ts b/packages/cubejs-api-gateway/src/gateway.ts index 1b5389d9e5c86..e67c7755b4b0f 100644 --- a/packages/cubejs-api-gateway/src/gateway.ts +++ b/packages/cubejs-api-gateway/src/gateway.ts @@ -133,6 +133,65 @@ function systemAsyncHandler(handler: (req: Request & { context: ExtendedRequestC const DEV_TOKEN_SCOPE = 'dev-token'; +/** + * A query that hit no pre-aggregation reports nothing rather than an empty + * object, so the key is simply absent from the response. Applied to the dev + * mode object as well, otherwise `'usedPreAggregations' in response` would + * answer differently in dev mode and in production. The Rust side normalizes + * the same way in `is_reportable_used_pre_aggregations`. + */ +function nonEmptyUsedPreAggregations( + usedPreAggregations: Record | undefined +): Record | undefined { + return usedPreAggregations && Object.keys(usedPreAggregations).length > 0 + ? usedPreAggregations + : undefined; +} + +/** + * 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. + * + * `targetTableName` is left out too. It names the physical table of one + * specific build, down to the content and structure version hashes, which a + * data API consumer cannot query anyway; `preAggregationId` plus the entry key + * identify the pre-aggregation and `lastUpdatedAt` dates the build. + * + * The full object, including both, is still returned in dev mode and to the + * Playground. + */ +function publicUsedPreAggregations( + usedPreAggregations: Record | undefined +): Record | undefined { + const used = nonEmptyUsedPreAggregations(usedPreAggregations); + if (!used) { + return undefined; + } + + const publicFields = ['preAggregationId', 'lastUpdatedAt', 'type']; + + return Object.fromEntries( + Object.entries(used).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; @@ -1951,12 +2010,16 @@ 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 ? { refreshKeyValues: response.refreshKeyValues, - usedPreAggregations: response.usedPreAggregations, + usedPreAggregations: nonEmptyUsedPreAggregations(response.usedPreAggregations), transformedQuery: sqlQuery.canUseTransformedQuery, requestId: context.requestId, } @@ -2248,6 +2311,16 @@ 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. + // + // Always the redacted projection, with no dev mode override unlike + // `prepareResultTransformData`: this branch serves the SQL API, + // which is not a Playground path, and its result object carries + // none of the other dev-only fields either. + usedPreAggregations: publicUsedPreAggregations(response.usedPreAggregations), // Always false: this branch builds its sqlQuery with // `disableExternalPreAggregations` set (above), which makes // `externalPreAggregationQuery()` return false, and the diff --git a/packages/cubejs-api-gateway/src/graphql.ts b/packages/cubejs-api-gateway/src/graphql.ts index 4a328d1cd6a02..da94b4f4c4ce1 100644 --- a/packages/cubejs-api-gateway/src/graphql.ts +++ b/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) diff --git a/packages/cubejs-api-gateway/test/graphql.test.ts b/packages/cubejs-api-gateway/test/graphql.test.ts index 830c8d20f5797..d7a519cfb75fe 100644 --- a/packages/cubejs-api-gateway/test/graphql.test.ts +++ b/packages/cubejs-api-gateway/test/graphql.test.ts @@ -343,6 +343,61 @@ describe('GraphQL Schema', () => { expect(res.body).toMatchSnapshot(); }); + // `graphql.ts` copies the same pre-aggregation identity into `extensions` + // as the REST response carries. Its own app so the snapshot above keeps + // describing a query that hit no pre-aggregation. + test('should return usedPreAggregations in extensions', async () => { + const usedPreAggregations = { + 'schema.orders_main20240101': { + preAggregationId: 'Orders.main', + lastUpdatedAt: 1712000000000, + type: 'rollup', + }, + }; + + const preAggApp = express(); + + preAggApp.use('/graphql', jsonParser, (req: any, res: any) => { + const schema = makeSchema(metaConfig); + + return graphqlHTTP({ + schema, + context: { + req, + res, + apiGateway: { + async load({ query, res: response }) { + response({ + query, + annotation: mockAnnotation, + lastRefreshTime: mockLastRefreshTime, + usedPreAggregations, + data: [ + { 'Orders.count': 10, 'Orders.totalAmount': 500, 'Orders.status': 'completed' }, + ], + }); + }, + }, + }, + extensions: () => res.extensions || {}, + })(req, res); + }); + + const query = `query CubeQuery { + cube { + orders { count totalAmount status } + } + }`; + + const res = await request(preAggApp) + .post('/graphql') + .set('Content-Type', 'application/json') + .send(gqlQuery(query)); + + expect(res.body.errors).toBeUndefined(); + expect(res.body.extensions.usedPreAggregations).toEqual(usedPreAggregations); + }); + it('should accumulate all measures and dimensions', async () => { const query = `query CubeQuery { cube { diff --git a/packages/cubejs-api-gateway/test/index.test.ts b/packages/cubejs-api-gateway/test/index.test.ts index 87d0825e0bdf7..76177aaafcef0 100644 --- a/packages/cubejs-api-gateway/test/index.test.ts +++ b/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,25 @@ 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 and + // the physical table name of the build are not: the former are rows of the + // refresh key queries, which are often written without the security context + // filtering the cube itself applies, and the latter describes storage + // layout a data API consumer cannot query anyway. + test('external=true exposes pre-aggregation identity only', 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 +1565,23 @@ 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', + lastUpdatedAt: 1712000000000, + type: 'rollup', + }, + }); }); - test('usedPreAggregations is exposed under playground auth alongside external', async () => { + test('the full pre-aggregation object is exposed under playground auth', 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( diff --git a/packages/cubejs-api-gateway/test/sql-api-load.test.ts b/packages/cubejs-api-gateway/test/sql-api-load.test.ts index de5907f503394..49aa13af49d27 100644 --- a/packages/cubejs-api-gateway/test/sql-api-load.test.ts +++ b/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 | undefined; + + public constructor(lastRefreshTime?: Date, usedPreAggregations?: Record) { 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, }; } } @@ -33,7 +49,11 @@ function createGateway(adapterApi: any) { }); } -async function sqlApiLoad(adapterApi: any, sqlQuery?: [string, any[]]) { +async function sqlApiLoad( + adapterApi: any, + sqlQuery?: [string, any[]], + signedWithPlaygroundAuthSecret = false, +) { const apiGateway = createGateway(adapterApi); let response: any; @@ -46,7 +66,7 @@ async function sqlApiLoad(adapterApi: any, sqlQuery?: [string, any[]]) { context: { requestId: 'sql-api-load-test', securityContext: {}, - signedWithPlaygroundAuthSecret: false, + signedWithPlaygroundAuthSecret, } as any, res: (r: any) => { response = r; @@ -88,4 +108,107 @@ 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', + 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. `targetTableName` names the physical table of one build, + // hashes included. Both stay dev-mode only. + test('usedPreAggregations omits refreshKeyValues and targetTableName', 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(); + expect(usage.targetTableName).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(); + }); + + // Playground and dev mode get the object unredacted, but the same + // nothing-to-report normalization: a client checking for the key must not see + // it appear only because the deployment runs in dev mode. + test('usedPreAggregations is absent under playground auth when empty', async () => { + const response = await sqlApiLoad( + new FreshnessAdapterApiMock(LAST_REFRESH_TIME, {}), + undefined, + true, + ); + + const result = response.getResults()[0].getRootResultObject()[0]; + + expect(result.usedPreAggregations).toBeUndefined(); + }); + + test('usedPreAggregations keeps the full object under playground auth', async () => { + const response = await sqlApiLoad( + new FreshnessAdapterApiMock(LAST_REFRESH_TIME, USED_PRE_AGGREGATIONS), + undefined, + true, + ); + + const result = response.getResults()[0].getRootResultObject()[0]; + + expect(result.usedPreAggregations['schema.orders_main20240101']) + .toEqual(USED_PRE_AGGREGATIONS['schema.orders_main20240101']); + }); + + // Deliberate asymmetry with the branch above: the pushed-down result keeps + // the redacted projection even under Playground auth, because that branch + // serves the SQL API and carries no dev-only fields at all. + test('pushed-down sqlQuery result stays redacted under playground auth', async () => { + const response = await sqlApiLoad( + new FreshnessAdapterApiMock(LAST_REFRESH_TIME, USED_PRE_AGGREGATIONS), + ['SELECT * FROM test', []], + true, + ); + + const usage = response.results[0].usedPreAggregations['schema.orders_main20240101']; + expect(usage.targetTableName).toBeUndefined(); + expect(usage.refreshKeyValues).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(); + expect(result.usedPreAggregations['schema.orders_main20240101'].targetTableName) + .toBeUndefined(); + }); }); diff --git a/packages/cubejs-backend-native/src/node_export.rs b/packages/cubejs-backend-native/src/node_export.rs index f5f6d4af34e03..5922444bcb06c 100644 --- a/packages/cubejs-backend-native/src/node_export.rs +++ b/packages/cubejs-backend-native/src/node_export.rs @@ -1,3 +1,4 @@ +use cubesql::compile::engine::df::scan::parse_used_pre_aggregations; use cubesql::compile::parser::parse_sql_to_statement; use cubesql::compile::{convert_statement_to_cube_query, get_df_batches}; use cubesql::config::processing_loop::ShutdownMode; @@ -391,16 +392,21 @@ async fn handle_sql_query( // branch and never calls `load_data`, so neither the span nor the // stream schema carries the metadata and the header omits it. // - // Both values take the same precedence: whatever the span reported + // Every value takes the same precedence: whatever the span reported // wins outright, and the schema is consulted only when the span was // silent. `external` must not be OR-ed with the schema — a span that // folded to `false` because only some of its loads were external // would then be overridden back to `true`, undoing the conservative // fold in `SpanId::set_external`. - let (span_last_refresh_time, span_external) = match span_id_for_schema.as_ref() { - Some(span_id) => (span_id.last_refresh_time().await, span_id.external().await), - None => (None, None), - }; + let (span_last_refresh_time, span_external, span_used_pre_aggregations) = + match span_id_for_schema.as_ref() { + Some(span_id) => ( + span_id.last_refresh_time().await, + span_id.external().await, + span_id.used_pre_aggregations().await, + ), + None => (None, None, None), + }; let last_refresh_time = span_last_refresh_time.or_else(|| { stream @@ -428,6 +434,22 @@ async fn handle_sql_query( schema_response.insert("external".into(), serde_json::Value::Bool(true)); } + // Names the pre-aggregations behind the result so a client can join + // it to the build it is watching. Same precedence as above; the + // span already holds the union across every scan of the plan, while + // the stream schema only ever describes the last one. + let used_pre_aggregations = span_used_pre_aggregations.or_else(|| { + stream + .schema() + .metadata() + .get("usedPreAggregations") + .map(String::as_str) + .and_then(parse_used_pre_aggregations) + }); + if let Some(used_pre_aggregations) = used_pre_aggregations { + schema_response.insert("usedPreAggregations".into(), used_pre_aggregations); + } + write_jsonl_message( channel.clone(), stream_methods.write.clone(), diff --git a/packages/cubejs-backend-native/src/orchestrator.rs b/packages/cubejs-backend-native/src/orchestrator.rs index 4524a7e502b62..d5b45e308f904 100644 --- a/packages/cubejs-backend-native/src/orchestrator.rs +++ b/packages/cubejs-backend-native/src/orchestrator.rs @@ -38,6 +38,7 @@ pub struct ResultWrapper { transformed_data: Option, pub last_refresh_time: Option, pub external: bool, + pub used_pre_aggregations: Option, } impl ResultWrapper { @@ -115,6 +116,7 @@ impl ResultWrapper { transformed_data: None, last_refresh_time: None, external: false, + used_pre_aggregations: None, }) } diff --git a/packages/cubejs-backend-native/src/transport.rs b/packages/cubejs-backend-native/src/transport.rs index f6bca69f71deb..e497d6a62c738 100644 --- a/packages/cubejs-backend-native/src/transport.rs +++ b/packages/cubejs-backend-native/src/transport.rs @@ -16,7 +16,7 @@ use async_trait::async_trait; use cubeorchestrator::query_result_transform::RequestResultData; use cubesql::compile::engine::df::scan::{ build_response_schema, convert_transport_response, transform_response, CacheMode, MemberField, - RecordBatch, SchemaRef, + RecordBatch, ResultMetadata, SchemaRef, }; use cubesql::compile::engine::df::wrapper::SqlQuery; use cubesql::transport::{ @@ -444,6 +444,7 @@ impl TransportService for NodeBridgeTransport { wrapper.last_refresh_time = result_data.last_refresh_time; wrapper.external = result_data.external.unwrap_or(false); + wrapper.used_pre_aggregations = result_data.used_pre_aggregations; native_wrapped_results.push(wrapper); } @@ -532,8 +533,11 @@ impl TransportService for NodeBridgeTransport { .map(|mut wrapper| { let updated_schema = build_response_schema( &schema, - wrapper.last_refresh_time.clone(), - wrapper.external, + &ResultMetadata { + last_refresh_time: wrapper.last_refresh_time.clone(), + external: wrapper.external, + used_pre_aggregations: wrapper.used_pre_aggregations.clone(), + }, ); transform_response(&mut wrapper, updated_schema, &member_fields) diff --git a/packages/cubejs-backend-native/test/sql.test.ts b/packages/cubejs-backend-native/test/sql.test.ts index e79f2966033c6..089edd527c425 100644 --- a/packages/cubejs-backend-native/test/sql.test.ts +++ b/packages/cubejs-backend-native/test/sql.test.ts @@ -576,11 +576,12 @@ describe('SQLInterface', () => { // header used to come back without them while the same base measures queried // plainly did carry them. // - // Both queries carry an explicit LIMIT to pin them to the buffered path. + // Both queries carry a small explicit LIMIT to pin them to the buffered path. // `CubeScanExecutionPlan::execute` switches to `load_stream` when stream mode - // is on and the request has no limit, and that branch never runs `load_data`, - // so no freshness metadata is recorded at all — a known gap, and this suite - // runs under CUBESQL_STREAM_MODE=true in CI. + // is on and the request either has no limit or one above + // `CUBESQL_NON_STREAMING_QUERY_MAX_ROW_LIMIT`, and that branch never runs + // `load_data`, so no result metadata is recorded at all — a known gap, and + // this suite runs under CUBESQL_STREAM_MODE=true in CI. test.each([ [ 'plain measure projection', @@ -591,7 +592,7 @@ describe('SQLInterface', () => { 'SELECT customer_gender, ROUND(MEASURE(maxPrice) / MEASURE(count), 2) AS avg_value, MEASURE(count) AS cnt FROM KibanaSampleDataEcommerce GROUP BY 1 ORDER BY 3 DESC LIMIT 10;', ], ])( - 'lastRefreshTime and external survive in /cubesql JSONL header for a %s', + 'lastRefreshTime, external and usedPreAggregations survive in /cubesql JSONL header for a %s', async (_name, sql) => { const methods = { ...interfaceMethods(), @@ -618,6 +619,13 @@ describe('SQLInterface', () => { }, lastRefreshTime: '2024-01-01T00:00:00.000Z', external: true, + usedPreAggregations: { + 'schema.kibana_main': { + preAggregationId: 'KibanaSampleDataEcommerce.main', + lastUpdatedAt: 1712000000000, + type: 'rollup', + }, + }, }, ], }; @@ -650,6 +658,99 @@ describe('SQLInterface', () => { expect(schemaLine).toBeDefined(); expect(schemaLine.lastRefreshTime).toBe('2024-01-01T00:00:00.000Z'); expect(schemaLine.external).toBe(true); + expect(schemaLine.usedPreAggregations['schema.kibana_main'].preAggregationId) + .toBe('KibanaSampleDataEcommerce.main'); + } finally { + await native.shutdownInterface(instance, 'fast'); + } + } + ); + + // Identity of the pre-aggregations behind the result, so a client can join a + // chart to the build it is watching. Arrow schema metadata is a string map, + // so the object round-trips through JSON on the way here. + test.each([ + [ + 'surfaced when reported', + { + 'schema.kibana_main': { + preAggregationId: 'KibanaSampleDataEcommerce.main', + lastUpdatedAt: 1712000000000, + type: 'rollup', + }, + }, + { + 'schema.kibana_main': { + preAggregationId: 'KibanaSampleDataEcommerce.main', + lastUpdatedAt: 1712000000000, + type: 'rollup', + }, + }, + ], + // A query that hit no pre-aggregation must not leave an empty key behind. + ['absent when the object is empty', {}, undefined], + ['absent when not reported', undefined, undefined], + ])( + 'usedPreAggregations is %s in /cubesql JSONL schema header', + async (_name, usedPreAggregations, expected) => { + const methods = { + ...interfaceMethods(), + sqlApiLoad: jest.fn(async ({ streaming, query }: any) => { + if (streaming) { + return { stream: new FakeRowStream(query) }; + } + return { + results: [ + { + annotation: { + measures: {}, + dimensions: {}, + segments: {}, + timeDimensions: {}, + }, + data: { + members: ['KibanaSampleDataEcommerce.order_date'], + columns: [['2024-01-01T00:00:00.000']], + }, + lastRefreshTime: '2024-01-01T00:00:00.000Z', + external: true, + usedPreAggregations, + }, + ], + }; + }), + }; + + const instance = await native.registerInterface({ + ...methods, + canSwitchUserForSession: (_payload: any) => true, + }); + + let buf = ''; + const lines: any[] = []; + const write = jest.fn((chunk, _enc, callback) => { + const raw = (buf + chunk.toString('utf-8')).split('\n'); + buf = raw.pop() || ''; + for (const l of raw) { + if (l.trim().length) { + lines.push(JSON.parse(l)); + } + } + callback(); + }); + const cubeSqlStream = new Writable({ write }); + + try { + await native.execSql( + instance, + 'SELECT order_date FROM KibanaSampleDataEcommerce LIMIT 1;', + cubeSqlStream + ); + + const schemaLine = lines.find((o) => o.schema); + expect(schemaLine).toBeDefined(); + expect(schemaLine.usedPreAggregations).toEqual(expected); + expect(schemaLine.lastRefreshTime).toBe('2024-01-01T00:00:00.000Z'); } finally { await native.shutdownInterface(instance, 'fast'); } diff --git a/packages/cubejs-client-core/src/types.ts b/packages/cubejs-client-core/src/types.ts index 4a3ee92658c17..668d7f27cba03 100644 --- a/packages/cubejs-client-core/src/types.ts +++ b/packages/cubejs-client-core/src/types.ts @@ -197,7 +197,17 @@ export type TransformedQuery = { export type PreAggregationType = 'rollup' | 'rollupJoin' | 'rollupLambda' | 'originalSql'; export type UsedPreAggregation = { - targetTableName: string; + /** + * Identity of the pre-aggregation in the data model, e.g. `Orders.main`. + * Stable across rebuilds, unlike `targetTableName`. + */ + preAggregationId?: string; + /** + * Physical table of one specific build, content and structure versions + * included. Returned in dev mode and to the Playground only. + */ + targetTableName?: string; + lastUpdatedAt?: number; type: PreAggregationType; }; @@ -210,6 +220,11 @@ export type LoadResponseResult = { dbType: string; extDbType: string; requestId?: string; + /** + * Pre-aggregations this result was served from, keyed by pre-aggregation + * table name. Absent when the query hit none. Only carries identity fields; + * `refreshKeyValues` is added in dev mode and for the Playground. + */ usedPreAggregations?: Record; transformedQuery?: TransformedQuery; total?: number; diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts b/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts index df5015a4a957f..857aaaa2da225 100644 --- a/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts +++ b/packages/cubejs-query-orchestrator/src/orchestrator/PreAggregations.ts @@ -139,11 +139,28 @@ type PreAggJob = { dataSource: string, }; +/** + * Types a pre-aggregation description can have by the time it reaches the + * orchestrator. Narrower than the same-named type in `@cubejs-backend/client-core`, + * which also lists `rollupJoin` and `rollupLambda`: the schema compiler expands + * those into the rollups they reference (`preAggregationDescriptionsFor`), so no + * description with either type is ever built or reported here. + */ +export type PreAggregationType = 'rollup' | 'originalSql'; + export type LoadPreAggregationResult = { targetTableName: string; refreshKeyValues: any[]; lastUpdatedAt: number; buildRangeEnd: string; + /** + * Identity of the pre-aggregation this table belongs to, stamped by + * `loadAllPreAggregationsIfNeeded` from the query's pre-aggregation + * description rather than by the loaders. Reported to clients as part of + * `usedPreAggregations` so they can match a result to a build. + */ + preAggregationId?: string; + type?: PreAggregationType; lambdaTable?: InlineTable; queryKey?: any[]; rollupLambdaId?: string; @@ -165,7 +182,7 @@ export type LambdaQuery = { export type PreAggregationDescription = { preAggregationsSchema: string; - type: 'rollup' | 'originalSql'; + type: PreAggregationType; preAggregationId: string; priority: number; dataSource: string; @@ -572,6 +589,7 @@ export class PreAggregations { const loadResult = await loader.loadPreAggregations(); const usedPreAggregation = { ...loadResult, + preAggregationId: p.preAggregationId, type: p.type, }; if (!usedPreAggregation.isMultiTableUnion) { diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts b/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts index c7cfaf8827542..01c1e202ae3f6 100644 --- a/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts +++ b/packages/cubejs-query-orchestrator/src/orchestrator/QueryOrchestrator.ts @@ -232,6 +232,8 @@ export class QueryOrchestrator { targetTableName: pa.targetTableName, refreshKeyValues: pa.refreshKeyValues, lastUpdatedAt: pa.lastUpdatedAt, + preAggregationId: pa.preAggregationId, + type: pa.type, })), )(preAggregationsTablesToTempTables); diff --git a/rust/cubesql/cubeclient/src/models/v1_load_result.rs b/rust/cubesql/cubeclient/src/models/v1_load_result.rs index 96f2e55bc9a99..55721b62b2a40 100644 --- a/rust/cubesql/cubeclient/src/models/v1_load_result.rs +++ b/rust/cubesql/cubeclient/src/models/v1_load_result.rs @@ -36,6 +36,17 @@ pub struct V1LoadResult { /// pre-aggregations hit the source DB and rely on its own caching. #[serde(rename = "external", skip_serializing_if = "Option::is_none")] pub external: Option, + /// Pre-aggregations this result was served from, keyed by pre-aggregation + /// table name, as sent by the API gateway. Carries identity only + /// (`preAggregationId`, `lastUpdatedAt`, `type`, plus `targetTableName` and + /// `refreshKeyValues` in dev mode and for the Playground) so a client can + /// match the result to a build it is watching. Kept as raw JSON: cubesql + /// passes it through without reading into it. + #[serde( + rename = "usedPreAggregations", + skip_serializing_if = "Option::is_none" + )] + pub used_pre_aggregations: Option, } impl Default for V1LoadResult { @@ -47,6 +58,7 @@ impl Default for V1LoadResult { refresh_key_values: None, last_refresh_time: None, external: None, + used_pre_aggregations: None, } } } @@ -60,6 +72,7 @@ impl V1LoadResult { refresh_key_values: None, last_refresh_time: None, external: None, + used_pre_aggregations: None, } } } diff --git a/rust/cubesql/cubesql/src/compile/engine/df/scan.rs b/rust/cubesql/cubesql/src/compile/engine/df/scan.rs index d6ab813b31a00..3f2dfd434d5d6 100644 --- a/rust/cubesql/cubesql/src/compile/engine/df/scan.rs +++ b/rust/cubesql/cubesql/src/compile/engine/df/scan.rs @@ -931,6 +931,16 @@ async fn load_data( .unwrap_or(false), ) .await; + + if let Some(used_pre_aggregations) = metadata + .get("usedPreAggregations") + .map(String::as_str) + .and_then(parse_used_pre_aggregations) + { + span_id + .merge_used_pre_aggregations(used_pre_aggregations) + .await; + } } match (options.max_records, data.num_rows()) { @@ -1357,30 +1367,55 @@ pub fn transform_response( transform_response_body!(response, schema, member_fields) } -/// Builds a schema with `lastRefreshTime` / `external` metadata. +/// Result metadata of a single load response, as reported by the API gateway. +/// Grouped in one struct so every value is named at the call site instead of +/// riding along as a positional argument. +#[derive(Debug, Clone, Default)] +pub struct ResultMetadata { + pub last_refresh_time: Option, + /// `true` when the result was served from an external (CubeStore) + /// pre-aggregation. + pub external: bool, + /// `usedPreAggregations` object of the load response, passed through + /// verbatim. + pub used_pre_aggregations: Option, +} + +/// Builds a schema with `lastRefreshTime` / `external` / `usedPreAggregations` +/// metadata. /// /// `lastRefreshTime` is passed through unchanged. The `external` marker is /// added when the flag is set so downstream code can tell that the result /// was served from an external (CubeStore) pre-aggregation — the case /// where cubesql's own cache-freshness decisions actually need to look at /// the pre-agg refresh, as internal pre-aggregations hit the source DB -/// and rely on its own caching. -pub fn build_response_schema( - schema: &SchemaRef, - last_refresh_time: Option, - external: bool, -) -> SchemaRef { - if last_refresh_time.is_none() && !external { +/// and rely on its own caching. `usedPreAggregations` rides along the same +/// way, JSON-encoded because Arrow schema metadata is a string map; it names +/// the pre-aggregations behind the result so a client can match it to a build. +pub fn build_response_schema(schema: &SchemaRef, result_metadata: &ResultMetadata) -> SchemaRef { + let used_pre_aggregations = result_metadata + .used_pre_aggregations + .as_ref() + .filter(|v| is_reportable_used_pre_aggregations(v)) + .and_then(|v| serde_json::to_string(v).ok()); + + if result_metadata.last_refresh_time.is_none() + && !result_metadata.external + && used_pre_aggregations.is_none() + { return schema.clone(); } let mut metadata = schema.metadata().clone(); - if let Some(t) = last_refresh_time { - metadata.insert("lastRefreshTime".to_string(), t); + if let Some(t) = &result_metadata.last_refresh_time { + metadata.insert("lastRefreshTime".to_string(), t.clone()); } - if external { + if result_metadata.external { metadata.insert("external".to_string(), "true".to_string()); } + if let Some(used_pre_aggregations) = used_pre_aggregations { + metadata.insert("usedPreAggregations".to_string(), used_pre_aggregations); + } Arc::new(Schema::new_with_metadata( schema.fields().to_vec(), @@ -1388,6 +1423,35 @@ pub fn build_response_schema( )) } +/// Whether a `usedPreAggregations` value is worth passing on: a query that hit +/// no pre-aggregation reports an empty object, and anything that is not an +/// object at all is not something a reader can merge into a span or hand to a +/// client. The API gateway is the only writer and always sends an object, so +/// the type check is defensive. +fn is_reportable_used_pre_aggregations(value: &serde_json::Value) -> bool { + matches!(value, serde_json::Value::Object(map) if !map.is_empty()) +} + +/// Reads the `usedPreAggregations` schema metadata written by +/// `build_response_schema` back into a value, applying the same +/// nothing-to-report normalization so that every reader of the metadata agrees +/// on it - the writer and the readers live in different crates. A blob that +/// does not parse is reported and dropped rather than failing the query: the +/// metadata is reporting only, and no result depends on it. +pub fn parse_used_pre_aggregations(encoded: &str) -> Option { + match serde_json::from_str::(encoded) { + Ok(value) if is_reportable_used_pre_aggregations(&value) => Some(value), + Ok(_) => None, + Err(e) => { + warn!( + "Unable to parse usedPreAggregations of a load response: {}", + e + ); + None + } + } +} + pub fn convert_transport_response( response: V1LoadResponse, schema: SchemaRef, @@ -1401,13 +1465,20 @@ pub fn convert_transport_response( data, last_refresh_time, external, + used_pre_aggregations, .. } = result; let V1LoadResultDataColumnar { members, columns } = data; let mut response = JsonColumnarValueObject::try_new(members, columns)?; - let updated_schema = - build_response_schema(&schema, last_refresh_time, external.unwrap_or(false)); + let updated_schema = build_response_schema( + &schema, + &ResultMetadata { + last_refresh_time, + external: external.unwrap_or(false), + used_pre_aggregations, + }, + ); transform_response(&mut response, updated_schema, &member_fields) }) @@ -1448,20 +1519,26 @@ mod tests { #[test] fn build_response_schema_no_metadata_when_nothing_to_add() { let schema = build_schema(); - let updated = build_response_schema(&schema, None, false); + let updated = build_response_schema(&schema, &ResultMetadata::default()); assert!(updated.metadata().is_empty()); } #[test] fn build_response_schema_passes_through_last_refresh_time() { let schema = build_schema(); - let updated = - build_response_schema(&schema, Some("2024-01-01T00:00:00.000Z".to_string()), false); + let updated = build_response_schema( + &schema, + &ResultMetadata { + last_refresh_time: Some("2024-01-01T00:00:00.000Z".to_string()), + ..Default::default() + }, + ); assert_eq!( updated.metadata().get("lastRefreshTime"), Some(&"2024-01-01T00:00:00.000Z".to_string()) ); assert!(updated.metadata().get("external").is_none()); + assert!(updated.metadata().get("usedPreAggregations").is_none()); } #[test] @@ -1470,7 +1547,14 @@ mod tests { // passed through unchanged. The marker reports the external hit. let schema = build_schema(); let stale = "2000-01-01T00:00:00.000Z".to_string(); - let updated = build_response_schema(&schema, Some(stale.clone()), true); + let updated = build_response_schema( + &schema, + &ResultMetadata { + last_refresh_time: Some(stale.clone()), + external: true, + ..Default::default() + }, + ); assert_eq!(updated.metadata().get("lastRefreshTime"), Some(&stale)); assert_eq!( @@ -1484,7 +1568,13 @@ mod tests { let schema = build_schema(); // No incoming last_refresh_time, but external flag is set — emit // only the marker; do NOT synthesize a `lastRefreshTime`. - let updated = build_response_schema(&schema, None, true); + let updated = build_response_schema( + &schema, + &ResultMetadata { + external: true, + ..Default::default() + }, + ); assert!(updated.metadata().get("lastRefreshTime").is_none()); assert_eq!( updated.metadata().get("external"), @@ -1492,6 +1582,78 @@ mod tests { ); } + #[test] + fn build_response_schema_json_encodes_used_pre_aggregations() { + // Arrow schema metadata is a string map, so the object travels as JSON + // and has to come back out of it unchanged. + let schema = build_schema(); + let used_pre_aggregations = serde_json::json!({ + "schema.orders_main20240101": { + "preAggregationId": "Orders.main", + "targetTableName": "schema.orders_main20240101_abc_def_1712", + "lastUpdatedAt": 1712000000000u64, + "type": "rollup", + } + }); + let updated = build_response_schema( + &schema, + &ResultMetadata { + used_pre_aggregations: Some(used_pre_aggregations.clone()), + ..Default::default() + }, + ); + + let encoded = updated.metadata().get("usedPreAggregations").unwrap(); + assert_eq!( + serde_json::from_str::(encoded).unwrap(), + used_pre_aggregations + ); + } + + #[test] + fn parse_used_pre_aggregations_normalizes_nothing_to_report() { + // Same normalization as the writer applies, so a reader of the metadata + // never has to special-case an empty object or a null + assert_eq!(parse_used_pre_aggregations("null"), None); + assert_eq!(parse_used_pre_aggregations("{}"), None); + // Nothing a reader could merge or report, however well-formed + assert_eq!(parse_used_pre_aggregations("\"nonsense\""), None); + assert_eq!(parse_used_pre_aggregations("[]"), None); + // Not JSON at all - reported and dropped, never fatal + assert_eq!(parse_used_pre_aggregations("{oops"), None); + + let used_pre_aggregations = serde_json::json!({ + "schema.orders_main": { "preAggregationId": "Orders.main" } + }); + assert_eq!( + parse_used_pre_aggregations(&used_pre_aggregations.to_string()), + Some(used_pre_aggregations) + ); + } + + #[test] + fn build_response_schema_skips_unreportable_used_pre_aggregations() { + // A query that hit no pre-aggregation reports an empty object; passing + // that on would make every plain SQL result carry a useless key. The + // same goes for a value no reader could merge. + let schema = build_schema(); + for used_pre_aggregations in [ + serde_json::json!({}), + serde_json::Value::Null, + serde_json::json!("nonsense"), + ] { + let updated = build_response_schema( + &schema, + &ResultMetadata { + used_pre_aggregations: Some(used_pre_aggregations), + ..Default::default() + }, + ); + + assert!(updated.metadata().is_empty()); + } + } + /// Collects everything a splitter produces for a single input batch: the chunk /// returned by `split` plus every buffered leftover, as row values. fn split_to_rows(splitter: &mut RecordBatchSplitter, batch: RecordBatch) -> Vec> { diff --git a/rust/cubesql/cubesql/src/transport/service.rs b/rust/cubesql/cubesql/src/transport/service.rs index 764748ff7f4bc..3a3e5bd8d0cc8 100644 --- a/rust/cubesql/cubesql/src/transport/service.rs +++ b/rust/cubesql/cubesql/src/transport/service.rs @@ -87,6 +87,7 @@ pub struct SpanId { is_data_query: RWLockAsync, last_refresh_time: RWLockAsync>>, external: RWLockAsync>, + used_pre_aggregations: RWLockAsync>, } impl SpanId { @@ -98,6 +99,7 @@ impl SpanId { is_data_query: tokio::sync::RwLock::new(false), last_refresh_time: tokio::sync::RwLock::new(None), external: tokio::sync::RwLock::new(None), + used_pre_aggregations: tokio::sync::RwLock::new(serde_json::Map::new()), } } @@ -149,6 +151,35 @@ impl SpanId { *self.external.read().await } + /// Records the pre-aggregations a load result contributing to this span was + /// served from. A span may cover several load requests (e.g. a join of cube + /// scans), and the entries are keyed by pre-aggregation table name, so they + /// merge: the span reports the union of every pre-aggregation read. Two + /// loads reporting the same table report the same identity for it, so a + /// later entry overwriting an earlier one is a no-op. + /// + /// Anything but a JSON object is ignored - the value is passed through from + /// the API gateway, and only an object can be merged. + pub async fn merge_used_pre_aggregations(&self, used_pre_aggregations: serde_json::Value) { + if let serde_json::Value::Object(entries) = used_pre_aggregations { + let mut write = self.used_pre_aggregations.write().await; + for (table_name, usage) in entries { + write.insert(table_name, usage); + } + } + } + + /// `None` when no load has reported a pre-aggregation, so callers can fall + /// back to another source instead of reporting an empty object. + pub async fn used_pre_aggregations(&self) -> Option { + let read = self.used_pre_aggregations.read().await; + if read.is_empty() { + None + } else { + Some(serde_json::Value::Object(read.clone())) + } + } + pub fn duration(&self) -> u64 { self.span_start .elapsed() @@ -1151,4 +1182,45 @@ mod tests { span_id.set_external(false).await; assert_eq!(span_id.external().await, Some(false)); } + + #[tokio::test] + async fn span_id_used_pre_aggregations_merge_across_loads() { + // Silent span — consumers fall back to the stream schema, so an empty + // union must not be reported as an empty object + let span_id = SpanId::new("test".to_string(), serde_json::json!({})); + assert_eq!(span_id.used_pre_aggregations().await, None); + + span_id + .merge_used_pre_aggregations(serde_json::json!({ + "schema.orders_main": { "preAggregationId": "Orders.main" } + })) + .await; + span_id + .merge_used_pre_aggregations(serde_json::json!({ + "schema.users_main": { "preAggregationId": "Users.main" } + })) + .await; + + // A join of two cube scans reports both pre-aggregations, not the last + // one to arrive + assert_eq!( + span_id.used_pre_aggregations().await, + Some(serde_json::json!({ + "schema.orders_main": { "preAggregationId": "Orders.main" }, + "schema.users_main": { "preAggregationId": "Users.main" } + })) + ); + + // Anything but an object has nothing mergeable in it and is dropped + span_id + .merge_used_pre_aggregations(serde_json::json!("nonsense")) + .await; + assert_eq!( + span_id + .used_pre_aggregations() + .await + .and_then(|v| v.as_object().map(|o| o.len())), + Some(2) + ); + } }