feat: Implement disable_post_processing in /v1/sql (#9331)

Author: mcheshkov

docs/pages/product/apis-integrations/rest-api/reference.mdx

@@ -119,12 +119,6 @@ If `disable_post_processing` is set to `true`, Cube will try to generate the SQL
 as if the query is run without [post-processing][ref-query-wpp], i.e., if it's run as a
 query with [pushdown][ref-query-wpd].
 
-<WarningBox>
-
-Currently, the `disable_post_processing` parameter is not yet supported.
-
-</WarningBox>
-
 The response will contain a JSON object with the following properties under the `sql` key:
 
 | Property, type | Description |

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

@@ -330,6 +330,7 @@ class ApiGateway {
       if (req.query.format === 'sql') {
         await this.sql4sql({
           query: req.query.query,
+          disablePostProcessing: req.query.disable_post_processing === 'true',
           context: req.context,
           res: this.resToResultFn(res)
         });
@@ -349,6 +350,7 @@ class ApiGateway {
       if (req.body.format === 'sql') {
         await this.sql4sql({
           query: req.body.query,
+          disablePostProcessing: req.body.disable_post_processing,
           context: req.context,
           res: this.resToResultFn(res)
         });
@@ -1308,13 +1310,14 @@ class ApiGateway {
 
   protected async sql4sql({
     query,
+    disablePostProcessing,
     context,
     res,
-  }: {query: string} & BaseRequest) {
+  }: {query: string, disablePostProcessing: boolean} & BaseRequest) {
     try {
       await this.assertApiScope('data', context.securityContext);
 
-      const result = await this.sqlServer.sql4sql(query, context.securityContext);
+      const result = await this.sqlServer.sql4sql(query, disablePostProcessing, context.securityContext);
       res({ sql: result });
     } catch (e: any) {
       this.handleError({

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

@@ -64,8 +64,8 @@ export class SQLServer {
     await execSql(this.sqlInterfaceInstance!, sqlQuery, stream, securityContext);
   }
 
-  public async sql4sql(sqlQuery: string, securityContext?: any): Promise<Sql4SqlResponse> {
-    return sql4sql(this.sqlInterfaceInstance!, sqlQuery, securityContext);
+  public async sql4sql(sqlQuery: string, disablePostProcessing: boolean, securityContext?: unknown): Promise<Sql4SqlResponse> {
+    return sql4sql(this.sqlInterfaceInstance!, sqlQuery, disablePostProcessing, securityContext);
   }
 
   protected buildCheckSqlAuth(options: SQLServerOptions): CheckSQLAuthFn {

packages/cubejs-backend-native/js/index.ts

@@ -405,10 +405,10 @@ export const execSql = async (instance: SqlInterfaceInstance, sqlQuery: string,
 };
 
 // TODO parse result from native code
-export const sql4sql = async (instance: SqlInterfaceInstance, sqlQuery: string, securityContext?: any): Promise<Sql4SqlResponse> => {
+export const sql4sql = async (instance: SqlInterfaceInstance, sqlQuery: string, disablePostProcessing: boolean, securityContext?: unknown): Promise<Sql4SqlResponse> => {
   const native = loadNative();
 
-  return native.sql4sql(instance, sqlQuery, securityContext ? JSON.stringify(securityContext) : null);
+  return native.sql4sql(instance, sqlQuery, disablePostProcessing, securityContext ? JSON.stringify(securityContext) : null);
 };
 
 export const buildSqlAndParams = (cubeEvaluator: any): String => {

packages/cubejs-backend-native/src/sql4sql.rs

@@ -2,11 +2,13 @@ use std::sync::Arc;
 
 use neon::prelude::*;
 
-use cubesql::compile::convert_sql_to_cube_query;
 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, CubeScanWrapperNode};
-use cubesql::sql::Session;
+use cubesql::compile::{convert_sql_to_cube_query, DatabaseVariable};
+use cubesql::sql::{Session, CUBESQL_PENALIZE_POST_PROCESSING_VAR};
 use cubesql::transport::MetaContext;
 use cubesql::CubeError;
 
@@ -157,8 +159,20 @@ async fn handle_sql4sql_query(
     services: Arc<NodeCubeServices>,
     native_auth_ctx: Arc<NativeAuthContext>,
     sql_query: &str,
+    disable_post_processing: bool,
 ) -> Result<Sql4SqlResponse, CubeError> {
     with_session(&services, native_auth_ctx.clone(), |session| async move {
+        if disable_post_processing {
+            let v = DatabaseVariable {
+                name: CUBESQL_PENALIZE_POST_PROCESSING_VAR.to_string(),
+                value: ScalarValue::Boolean(Some(true)),
+                var_type: VarType::UserDefined,
+                readonly: false,
+                additional_params: None,
+            };
+            session.state.set_variables(vec![v]);
+        }
+
         let transport = session.server.transport.clone();
         // todo: can we use compiler_cache?
         let meta_context = transport
@@ -176,8 +190,9 @@ async fn handle_sql4sql_query(
 pub fn sql4sql(mut cx: FunctionContext) -> JsResult<JsValue> {
     let interface = cx.argument::<JsBox<crate::node_export::SQLInterface>>(0)?;
     let sql_query = cx.argument::<JsString>(1)?.value(&mut cx);
+    let disable_post_processing = cx.argument::<JsBoolean>(2)?.value(&mut cx);
 
-    let security_context: Option<serde_json::Value> = match cx.argument::<JsValue>(2) {
+    let security_context: Option<serde_json::Value> = match cx.argument::<JsValue>(3) {
         Ok(string) => match string.downcast::<JsString, _>(&mut cx) {
             Ok(v) => v.value(&mut cx).parse::<serde_json::Value>().ok(),
             Err(_) => None,
@@ -208,7 +223,13 @@ pub fn sql4sql(mut cx: FunctionContext) -> JsResult<JsValue> {
     //  can do it relatively rare, and in a single loop for all JoinHandles
     //  this is just a watchdog for a Very Bad case, so latency requirement can be quite relaxed
     runtime.spawn(async move {
-        let result = handle_sql4sql_query(services, native_auth_ctx, &sql_query).await;
+        let result = handle_sql4sql_query(
+            services,
+            native_auth_ctx,
+            &sql_query,
+            disable_post_processing,
+        )
+        .await;
 
         if let Err(err) = deferred.try_settle_with(&channel, move |mut cx| {
             // `neon::result::ResultExt` is implemented only for Result<Handle, Handle>, even though Ok variant is not touched

packages/cubejs-testing/test/__snapshots__/smoke-cubesql.test.ts.snap

@@ -36,6 +36,64 @@ Object {
 }
 `;
 
+exports[`SQL API Cube SQL over HTTP sql4sql double aggregation post-processing with disabled post-processing 1`] = `
+Object {
+  "body": Object {
+    "sql": Object {
+      "query_type": "pushdown",
+      "sql": Array [
+        "SELECT \\"t\\".\\"avg_t_total_\\" \\"avg_t_total_\\" 
+FROM (
+  SELECT AVG(\\"t\\".\\"total\\") \\"avg_t_total_\\" 
+  FROM (
+    SELECT
+          \\"orders\\".status \\"status\\", sum(\\"orders\\".amount) \\"total\\"
+        FROM
+          (
+      select 1 as id, 100 as amount, 'new' status, '2024-01-01'::timestamptz created_at
+      UNION ALL
+      select 2 as id, 200 as amount, 'new' status, '2024-01-02'::timestamptz created_at
+      UNION ALL
+      select 3 as id, 300 as amount, 'processed' status, '2024-01-03'::timestamptz created_at
+      UNION ALL
+      select 4 as id, 500 as amount, 'processed' status, '2024-01-04'::timestamptz created_at
+      UNION ALL
+      select 5 as id, 600 as amount, 'shipped' status, '2024-01-05'::timestamptz created_at
+      ) AS \\"orders\\"  GROUP BY 1
+  ) AS \\"t\\"
+) AS \\"t\\"",
+        Array [],
+      ],
+      "status": "ok",
+    },
+  },
+  "headers": Headers {
+    Symbol(map): Object {
+      "access-control-allow-origin": Array [
+        "*",
+      ],
+      "connection": Array [
+        "keep-alive",
+      ],
+      "content-length": Array [
+        "878",
+      ],
+      "content-type": Array [
+        "application/json; charset=utf-8",
+      ],
+      "keep-alive": Array [
+        "timeout=5",
+      ],
+      "x-powered-by": Array [
+        "Express",
+      ],
+    },
+  },
+  "status": 200,
+  "statusText": "OK",
+}
+`;
+
 exports[`SQL API Cube SQL over HTTP sql4sql regular query 1`] = `
 Object {
   "body": Object {
@@ -244,6 +302,42 @@ Object {
 }
 `;
 
+exports[`SQL API Cube SQL over HTTP sql4sql strictly post-processing with disabled post-processing 1`] = `
+Object {
+  "body": Object {
+    "sql": Object {
+      "error": "Provided query can not be executed without post-processing.",
+      "query_type": "post_processing",
+      "status": "error",
+    },
+  },
+  "headers": Headers {
+    Symbol(map): Object {
+      "access-control-allow-origin": Array [
+        "*",
+      ],
+      "connection": Array [
+        "keep-alive",
+      ],
+      "content-length": Array [
+        "127",
+      ],
+      "content-type": Array [
+        "application/json; charset=utf-8",
+      ],
+      "keep-alive": Array [
+        "timeout=5",
+      ],
+      "x-powered-by": Array [
+        "Express",
+      ],
+    },
+  },
+  "status": 200,
+  "statusText": "OK",
+}
+`;
+
 exports[`SQL API Cube SQL over HTTP sql4sql wrapper 1`] = `
 Object {
   "body": Object {

packages/cubejs-testing/test/smoke-cubesql.test.ts

@@ -149,7 +149,7 @@ describe('SQL API', () => {
     });
 
     describe('sql4sql', () => {
-      async function generateSql(query: string) {
+      async function generateSql(query: string, disablePostPprocessing: boolean = false) {
         const response = await fetch(`${birdbox.configuration.apiUrl}/sql`, {
           method: 'POST',
           headers: {
@@ -159,6 +159,7 @@ describe('SQL API', () => {
           body: JSON.stringify({
             query,
             format: 'sql',
+            disable_post_processing: disablePostPprocessing,
           }),
         });
         const { status, statusText, headers } = response;
@@ -193,6 +194,10 @@ describe('SQL API', () => {
         expect(await generateSql(`SELECT version();`)).toMatchSnapshot();
       });
 
+      it('strictly post-processing with disabled post-processing', async () => {
+        expect(await generateSql(`SELECT version();`, true)).toMatchSnapshot();
+      });
+
       it('double aggregation post-processing', async () => {
         expect(await generateSql(`
           SELECT AVG(total)
@@ -206,6 +211,19 @@ describe('SQL API', () => {
         `)).toMatchSnapshot();
       });
 
+      it('double aggregation post-processing with disabled post-processing', async () => {
+        expect(await generateSql(`
+          SELECT AVG(total)
+          FROM (
+            SELECT
+              status,
+              SUM(totalAmount) AS total
+            FROM Orders
+            GROUP BY 1
+          ) t
+        `, true)).toMatchSnapshot();
+      });
+
       it('wrapper', async () => {
         expect(await generateSql(`
           SELECT