From 379fadcc6c6e84d36ee7003bc4846817f30bb199 Mon Sep 17 00:00:00 2001 From: Ryan McCormack Date: Tue, 30 Jun 2026 11:26:39 -0400 Subject: [PATCH] feat(cubesql): add disable_post_processing to /v1/cubesql endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Queries that require DataFusion post-processing fetch intermediate rows from the source DB capped at CUBEJS_DB_QUERY_LIMIT (default 50k). On large datasets this silently truncates the intermediate result, producing wrong aggregates. The sql4sql (/v1/sql?format=sql) path already classifies plans as pushdown/regular/post-processing and reports it to the caller, but never executes — so it cannot truncate. The /v1/cubesql path executes and streams results, so to guarantee accuracy it must block before execution. This adds the same control to /v1/cubesql with a hard guarantee: if the final plan still requires post-processing after the cost model is biased against it, the request fails (the error is delivered in the streamed response body) rather than returning truncated data. Changes: - env.ts: register CUBESQL_DISABLE_POST_PROCESSING env var (server-wide default, off by default) so all services are covered without per-request changes - gateway.ts: read disable_post_processing from request body, fall back to env var, pass through to execSql - sql-server.ts, index.ts: thread disablePostProcessing param to native - node_export.rs: set CUBESQL_PENALIZE_POST_PROCESSING_VAR before planning (biases the cost model toward pushdown) and, after planning, fail if the plan root is not a single CubeScan(Wrapped) node Tests: - smoke-cubesql.test.ts: error path (asserts the post-processing error is surfaced in the streamed body), push-down success, default-off baseline - api-gateway/test/index.test.ts: gateway unit test asserting the param is forwarded to execSql; existing /v1/cubesql call-signature assertions updated for the new trailing argument - backend-shared/test/env.test.ts: unit test for CUBESQL_DISABLE_POST_PROCESSING --- .../configuration/environment-variables.mdx | 13 +++++ .../core-data-apis/rest-api/reference.mdx | 10 ++++ packages/cubejs-api-gateway/src/gateway.ts | 3 +- packages/cubejs-api-gateway/src/sql-server.ts | 4 +- .../cubejs-api-gateway/test/index.test.ts | 41 ++++++++++++++ packages/cubejs-backend-native/js/index.ts | 4 +- .../cubejs-backend-native/src/node_export.rs | 54 +++++++++++++++++++ packages/cubejs-backend-shared/src/env.ts | 1 + .../cubejs-backend-shared/test/env.test.ts | 11 ++++ .../cubejs-testing/test/smoke-cubesql.test.ts | 46 ++++++++++++++++ 10 files changed, 182 insertions(+), 5 deletions(-) diff --git a/docs-mintlify/reference/configuration/environment-variables.mdx b/docs-mintlify/reference/configuration/environment-variables.mdx index 6bf869d010021..3fce8a4ed18f6 100644 --- a/docs-mintlify/reference/configuration/environment-variables.mdx +++ b/docs-mintlify/reference/configuration/environment-variables.mdx @@ -1529,6 +1529,19 @@ If `true`, enables query pushdown in the [SQL API][ref-sql-api]. | --------------- | ---------------------- | --------------------- | | `true`, `false` | `true` | `true` | +## `CUBESQL_DISABLE_POST_PROCESSING` + +If `true`, the [SQL API][ref-sql-api] rejects queries that can't be fully pushed +down to the data source instead of running them with post-processing. This prevents +results from being computed over a truncated intermediate result capped by +[`CUBEJS_DB_QUERY_LIMIT`](#cubejs_db_query_limit). It can be overridden per request +with the `disable_post_processing` option on the +[`/v1/cubesql`](/reference/core-data-apis/rest-api/reference#base_path/v1/cubesql) endpoint. + +| Possible Values | Default in Development | Default in Production | +| --------------- | ---------------------- | --------------------- | +| `true`, `false` | `false` | `false` | + ## `CUBESQL_STREAM_MODE` If `true`, enables the [streaming mode][ref-sql-api-streaming] in the [SQL API][ref-sql-api]. 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 e10beca0cd5a5..96c8aaeb5af8e 100644 --- a/docs-mintlify/reference/core-data-apis/rest-api/reference.mdx +++ b/docs-mintlify/reference/core-data-apis/rest-api/reference.mdx @@ -411,6 +411,7 @@ This endpoint is part of the [SQL API][ref-sql-api]. | `query` | The SQL query to run | ✅ Yes | | `timezone` | The [time zone][ref-time-zone] for this query in the [TZ Database Name][link-tzdb] format, e.g., `America/Los_Angeles` | ❌ No | | `cache` | See [cache control][ref-cache-control]. `stale-if-slow` by default | ❌ No | +| `disable_post_processing` | If `true`, reject the query instead of running it with [post-processing][ref-query-wpp]. `false` by default | ❌ No | ### Headers @@ -425,6 +426,13 @@ the `schema` property with column names and types, and optionally The following objects contain chunks of the result set under the `data` property. Each chunk includes one or more rows of the result set. +If `disable_post_processing` is `true`, the query runs only if it fully pushes down +to the data source. Otherwise the response contains an `error` property instead of +results computed with [post-processing][ref-query-wpp], which are capped by +[`CUBEJS_DB_QUERY_LIMIT`][ref-env-db-query-limit] and may be truncated. The server-wide +default is set by the +[`CUBESQL_DISABLE_POST_PROCESSING`][ref-env-disable-pp] environment variable. + ### Example Simple request: @@ -1015,6 +1023,8 @@ Keep-Alive: timeout=5 [ref-folders]: /reference/data-modeling/view#folders [ref-cache-control]: /reference/core-data-apis/rest-api#cache-control [ref-time-zone]: /reference/core-data-apis/queries#time-zone +[ref-env-db-query-limit]: /reference/configuration/environment-variables#cubejs_db_query_limit +[ref-env-disable-pp]: /reference/configuration/environment-variables#cubesql_disable_post_processing [link-tzdb]: https://en.wikipedia.org/wiki/Tz_database [ref-control-plane-api]: /reference/control-plane-api [self-metadata-api]: #metadata-api diff --git a/packages/cubejs-api-gateway/src/gateway.ts b/packages/cubejs-api-gateway/src/gateway.ts index cf1526fdf2754..32e695f1be9c7 100644 --- a/packages/cubejs-api-gateway/src/gateway.ts +++ b/packages/cubejs-api-gateway/src/gateway.ts @@ -480,7 +480,8 @@ class ApiGateway { try { await this.assertApiScope('data', req.context?.securityContext); - await this.sqlServer.execSql(req.body.query, res, req.context?.securityContext, req.body.cache, req.body.timezone, req.body.throwContinueWait, req.context?.requestId); + const disablePostProcessing = req.body.disable_post_processing ?? getEnv('disablePostProcessing'); + await this.sqlServer.execSql(req.body.query, res, req.context?.securityContext, req.body.cache, req.body.timezone, req.body.throwContinueWait, req.context?.requestId, disablePostProcessing); } catch (e: any) { // Quickfix for https://github.com/cube-js/cube/issues/10450, // Right now, it's too complicated to fix the issue correctly, because diff --git a/packages/cubejs-api-gateway/src/sql-server.ts b/packages/cubejs-api-gateway/src/sql-server.ts index 8b362ae57ce66..fdd9dd1b58ace 100644 --- a/packages/cubejs-api-gateway/src/sql-server.ts +++ b/packages/cubejs-api-gateway/src/sql-server.ts @@ -76,8 +76,8 @@ export class SQLServer { return this.sqlInterfaceInstance; } - public async execSql(sqlQuery: string, stream: any, securityContext?: any, cacheMode?: CacheMode, timezone?: string, throwContinueWait?: boolean, requestId?: string) { - await execSql(this.getSqlInterfaceInstance(), sqlQuery, stream, securityContext, cacheMode, timezone, throwContinueWait, requestId); + public async execSql(sqlQuery: string, stream: any, securityContext?: any, cacheMode?: CacheMode, timezone?: string, throwContinueWait?: boolean, requestId?: string, disablePostProcessing?: boolean) { + await execSql(this.getSqlInterfaceInstance(), sqlQuery, stream, securityContext, cacheMode, timezone, throwContinueWait, requestId, disablePostProcessing); } public async sql4sql(sqlQuery: string, disablePostProcessing: boolean, securityContext?: unknown): Promise { diff --git a/packages/cubejs-api-gateway/test/index.test.ts b/packages/cubejs-api-gateway/test/index.test.ts index 6c2f9c27311d7..2db815625a137 100644 --- a/packages/cubejs-api-gateway/test/index.test.ts +++ b/packages/cubejs-api-gateway/test/index.test.ts @@ -1304,6 +1304,7 @@ describe('API Gateway', () => { undefined, undefined, expect.any(String), + false, ); }); @@ -1345,6 +1346,7 @@ describe('API Gateway', () => { 'America/Los_Angeles', undefined, expect.any(String), + false, ); }); @@ -1382,6 +1384,44 @@ describe('API Gateway', () => { undefined, true, expect.any(String), + false, + ); + }); + + test('disable_post_processing can be passed', async () => { + const { app, apiGateway } = await createApiGateway(); + + // Mock the sqlServer.execSql method + const execSqlMock = jest.fn(async (query, stream, securityContext, cacheMode, timezone) => { + stream.write(`${JSON.stringify({ + schema: [{ name: 'id', column_type: 'Int' }] + })}\n`); + stream.end(); + }); + + apiGateway.getSQLServer().execSql = execSqlMock; + + await request(app) + .post('/cubejs-api/v1/cubesql') + .set('Content-type', 'application/json') + .set('Authorization', 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.e30.t-IDcSemACt8x4iTMCda8Yhe3iZaWbvV5XKSTbuAn0M') + .send({ + query: 'SELECT id FROM test LIMIT 3', + disable_post_processing: true, + }) + .responseType('text') + .expect(200); + + // Verify disable_post_processing is forwarded as the trailing argument + expect(execSqlMock).toHaveBeenCalledWith( + 'SELECT id FROM test LIMIT 3', + expect.anything(), + {}, + undefined, + undefined, + undefined, + expect.any(String), + true, ); }); @@ -1419,6 +1459,7 @@ describe('API Gateway', () => { undefined, undefined, 'test-request-id-12345', + false, ); }); diff --git a/packages/cubejs-backend-native/js/index.ts b/packages/cubejs-backend-native/js/index.ts index b5dbf3cdc4ad2..88ed6d0b8ca56 100644 --- a/packages/cubejs-backend-native/js/index.ts +++ b/packages/cubejs-backend-native/js/index.ts @@ -439,10 +439,10 @@ export const shutdownInterface = async (instance: SqlInterfaceInstance, shutdown await native.shutdownInterface(instance, shutdownMode); }; -export const execSql = async (instance: SqlInterfaceInstance, sqlQuery: string, stream: any, securityContext?: any, cacheMode: CacheMode = 'stale-if-slow', timezone?: string, throwContinueWait?: boolean, requestId?: string): Promise => { +export const execSql = async (instance: SqlInterfaceInstance, sqlQuery: string, stream: any, securityContext?: any, cacheMode: CacheMode = 'stale-if-slow', timezone?: string, throwContinueWait?: boolean, requestId?: string, disablePostProcessing?: boolean): Promise => { const native = loadNative(); - await native.execSql(instance, sqlQuery, stream, securityContext ? JSON.stringify(securityContext) : null, cacheMode, timezone, throwContinueWait, requestId); + await native.execSql(instance, sqlQuery, stream, securityContext ? JSON.stringify(securityContext) : null, cacheMode, timezone, throwContinueWait, requestId, disablePostProcessing ?? false); }; // TODO parse result from native code diff --git a/packages/cubejs-backend-native/src/node_export.rs b/packages/cubejs-backend-native/src/node_export.rs index c5bdc1cc687d4..c8f6135fa60ea 100644 --- a/packages/cubejs-backend-native/src/node_export.rs +++ b/packages/cubejs-backend-native/src/node_export.rs @@ -24,6 +24,13 @@ use crate::transport::NodeBridgeTransport; use crate::utils::{batch_to_rows, NonDebugInRelease}; use cubenativeutils::wrappers::neon::neon_guarded_funcion_call; use cubenativeutils::wrappers::NativeContextHolder; +use cubesql::compile::datafusion::logical_plan::LogicalPlan; +use cubesql::compile::datafusion::scalar::ScalarValue; +use cubesql::compile::datafusion::variable::VarType; +use cubesql::compile::engine::df::scan::CubeScanNode; +use cubesql::compile::engine::df::wrapper::CubeScanWrappedSqlNode; +use cubesql::compile::DatabaseVariable; +use cubesql::sql::CUBESQL_PENALIZE_POST_PROCESSING_VAR; use cubesqlplanner::cube_bridge::base_query_options::NativeBaseQueryOptions; use cubesqlplanner::planner::base_query::BaseQuery; use std::rc::Rc; @@ -243,6 +250,7 @@ async fn handle_sql_query( timezone: Option, throw_continue_wait: bool, request_id: Option, + disable_post_processing: bool, ) -> Result<(), CubeError> { let span_id = Some(Arc::new(SpanId::new( request_id.unwrap_or_else(|| Uuid::new_v4().to_string()), @@ -301,6 +309,16 @@ async fn handle_sql_query( *cm = throw_continue_wait; } + if disable_post_processing { + session.state.set_variables(vec![DatabaseVariable { + name: CUBESQL_PENALIZE_POST_PROCESSING_VAR.to_string(), + value: ScalarValue::Boolean(Some(true)), + var_type: VarType::UserDefined, + readonly: false, + additional_params: None, + }]); + } + let session_clone = Arc::clone(&session); let span_id_clone = span_id.clone(); @@ -324,6 +342,27 @@ async fn handle_sql_query( ) .await?; + if disable_post_processing { + // A query needs no post-processing only when its plan root is a single + // CubeScan(Wrapped) node — i.e. the whole query maps to one source query. + // Anything else means DataFusion processes intermediate rows in-memory, which + // are capped at CUBEJS_DB_QUERY_LIMIT (default 50k) and can be silently + // truncated, so we fail loudly instead of returning inaccurate results. + if let Ok(logical_plan) = query_plan.try_as_logical_plan() { + let is_pure_pushdown = matches!( + logical_plan, + LogicalPlan::Extension(ext) + if ext.node.as_any().is::() + || ext.node.as_any().is::() + ); + if !is_pure_pushdown { + return Err(CubeError::user( + "Provided query cannot be executed without post-processing (disabled to guarantee result accuracy).".to_string(), + )); + } + } + } + let mut stream = get_df_batches(&query_plan).await?; let semaphore = Arc::new(Semaphore::new(0)); @@ -562,6 +601,20 @@ fn exec_sql(mut cx: FunctionContext) -> JsResult { Err(_) => None, }; + let disable_post_processing: bool = match cx.argument::(8) { + Ok(val) => { + if val.is_a::(&mut cx) || val.is_a::(&mut cx) { + false + } else { + match val.downcast::(&mut cx) { + Ok(v) => v.value(&mut cx), + Err(_) => false, + } + } + } + Err(_) => false, + }; + let js_stream_on_fn = Arc::new( node_stream .get::(&mut cx, "on")? @@ -615,6 +668,7 @@ fn exec_sql(mut cx: FunctionContext) -> JsResult { timezone, throw_continue_wait, request_id, + disable_post_processing, ) .await; diff --git a/packages/cubejs-backend-shared/src/env.ts b/packages/cubejs-backend-shared/src/env.ts index 83230a3d7c2b3..969977ba324eb 100644 --- a/packages/cubejs-backend-shared/src/env.ts +++ b/packages/cubejs-backend-shared/src/env.ts @@ -317,6 +317,7 @@ const variables: Record any> = { .asInt(), nativeSqlPlanner: () => get('CUBEJS_TESSERACT_SQL_PLANNER').default('false').asBool(), nativeSqlPlannerPreAggregations: () => get('CUBEJS_TESSERACT_PRE_AGGREGATIONS').default('false').asBool(), + disablePostProcessing: () => get('CUBESQL_DISABLE_POST_PROCESSING').default('false').asBool(), transpilationWorkerThreads: () => { const enabled = get('CUBEJS_TRANSPILATION_WORKER_THREADS') .default('true') diff --git a/packages/cubejs-backend-shared/test/env.test.ts b/packages/cubejs-backend-shared/test/env.test.ts index 0bf0e80b7b4ff..098741b24d738 100644 --- a/packages/cubejs-backend-shared/test/env.test.ts +++ b/packages/cubejs-backend-shared/test/env.test.ts @@ -117,6 +117,17 @@ describe('getEnv', () => { expect(getEnv('livePreview')).toBe(false); }); + test('disablePostProcessing', () => { + delete process.env.CUBESQL_DISABLE_POST_PROCESSING; + expect(getEnv('disablePostProcessing')).toBe(false); // default off + + process.env.CUBESQL_DISABLE_POST_PROCESSING = 'true'; + expect(getEnv('disablePostProcessing')).toBe(true); + + process.env.CUBESQL_DISABLE_POST_PROCESSING = 'false'; + expect(getEnv('disablePostProcessing')).toBe(false); + }); + test('maxRequestSize', () => { delete process.env.CUBEJS_MAX_REQUEST_SIZE; expect(getEnv('maxRequestSize')).toBe(50 * 1024 * 1024); // default 50mb diff --git a/packages/cubejs-testing/test/smoke-cubesql.test.ts b/packages/cubejs-testing/test/smoke-cubesql.test.ts index 389454d969d8b..a6f9141894b73 100644 --- a/packages/cubejs-testing/test/smoke-cubesql.test.ts +++ b/packages/cubejs-testing/test/smoke-cubesql.test.ts @@ -241,6 +241,52 @@ describe('SQL API', () => { ]); }); + describe('disable_post_processing on /v1/cubesql', () => { + async function execCubeSql(query: string, disablePostProcessing: boolean = false) { + const response = await fetch(`${birdbox.configuration.apiUrl}/cubesql`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: token, + }, + body: JSON.stringify({ + query, + disable_post_processing: disablePostProcessing, + }), + }); + const { status } = response; + const text = await response.text(); + // Error responses are plain JSON; streaming success responses are newline-delimited JSON. + const firstLine = text.split('\n').find(l => l.trim()) ?? '{}'; + const parsed = JSON.parse(firstLine); + return { status, body: parsed }; + } + + it('errors when a post-processing plan is required and disable_post_processing is true', async () => { + // SELECT version() compiles to a DataFusion plan with no CubeScan root, i.e. it + // requires post-processing. /v1/cubesql is a streaming endpoint: errors are delivered + // in the (chunked) response body with HTTP 200, not via the status code. The guard + // must surface the post-processing error instead of streaming (possibly truncated) data. + const { status, body } = await execCubeSql('SELECT version();', true); + expect(status).toBe(200); + expect(body.error).toMatch(/post-processing/i); + expect(body.schema).toBeUndefined(); + }); + + it('succeeds for a push-down query with disable_post_processing true', async () => { + // Simple aggregate maps to a CubeScanNode/CubeScanWrappedSqlNode — no post-processing. + const { status, body } = await execCubeSql('SELECT SUM(totalAmount) AS total FROM Orders;', true); + expect(status).toBe(200); + expect(body.schema).toBeDefined(); + }); + + it('executes a post-processing query normally when disable_post_processing is false', async () => { + // Confirms the flag is off-by-default and does not change existing behaviour. + const { status } = await execCubeSql('SELECT version();', false); + expect(status).toBe(200); + }); + }); + describe('sql4sql', () => { async function generateSql(query: string, disablePostPprocessing: boolean = false) { const response = await fetch(`${birdbox.configuration.apiUrl}/sql`, {