fix(cubesql): Remove excessive limit on inner wrapped queries (#7864)

Author: paveltiunov

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

@@ -1180,7 +1180,8 @@ class ApiGateway {
     exportAnnotatedSql,
     memberExpressions,
     expressionParams,
-    disableExternalPreAggregations
+    disableExternalPreAggregations,
+    disableLimitEnforcing,
   }: QueryRequest) {
     const requestStarted = new Date();
 
@@ -1193,7 +1194,7 @@ class ApiGateway {
         query = this.parseMemberExpressionsInQueries(query);
       }
 
-      const [queryType, normalizedQueries] = await this.getNormalizedQueries(query, context);
+      const [queryType, normalizedQueries] = await this.getNormalizedQueries(query, context, disableLimitEnforcing);
 
       const sqlQueries = await Promise.all<any>(
         normalizedQueries.map(async (normalizedQuery) => (await this.getCompilerApi(context)).getSql(

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

@@ -159,6 +159,7 @@ export class SQLServer {
               memberExpressions: true,
               disableExternalPreAggregations: true,
               queryType: 'multi',
+              disableLimitEnforcing: true,
               context,
               res: (message) => {
                 resolve(message);

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

@@ -127,6 +127,7 @@ type QueryRequest = BaseRequest & {
   exportAnnotatedSql?: boolean;
   memberExpressions?: boolean;
   disableExternalPreAggregations?: boolean;
+  disableLimitEnforcing?: boolean;
 };
 
 type SqlApiRequest = BaseRequest & {

packages/cubejs-schema-compiler/src/adapter/BaseQuery.js

@@ -2500,10 +2500,12 @@ export class BaseQuery {
       },
       statements: {
         select: 'SELECT {{ select_concat | map(attribute=\'aliased\') | join(\', \') }} \n' +
-          'FROM (\n  {{ from }}\n) AS {{ from_alias }} \n' +
-          '{% if filter %} WHERE {{ filter }}{% endif %}' +
-          '{% if group_by %} GROUP BY {{ group_by | map(attribute=\'index\') | join(\', \') }}{% endif %}' +
-          '{% if order_by %} ORDER BY {{ order_by | map(attribute=\'expr\') | join(\', \') }}{% endif %}' +
+          'FROM (\n' +
+          '{{ from | indent(2, true) }}\n' +
+          ') AS {{ from_alias }}' +
+          '{% if filter %}\nWHERE {{ filter }}{% endif %}' +
+          '{% if group_by %}\nGROUP BY {{ group_by | map(attribute=\'index\') | join(\', \') }}{% endif %}' +
+          '{% if order_by %}\nORDER BY {{ order_by | map(attribute=\'expr\') | join(\', \') }}{% endif %}' +
           '{% if limit %}\nLIMIT {{ limit }}{% endif %}' +
           '{% if offset %}\nOFFSET {{ offset }}{% endif %}',
       },

packages/cubejs-testing/birdbox-fixtures/postgresql/schema/BigOrders.js

@@ -0,0 +1,42 @@
+cube(`BigOrders`, {
+  sql: `
+  select * 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
+  ) data
+  CROSS JOIN GENERATE_SERIES(1, 20000) value
+  `,
+  measures: {
+    totalAmount: {
+      sql: `amount`,
+      type: `sum`,
+    },
+    toRemove: {
+      type: `count`,
+    },
+  },
+  dimensions: {
+    id: {
+      sql: `id`,
+      type: `number`,
+      primaryKey: true,
+    },
+
+    status: {
+      sql: `status`,
+      type: `string`,
+    },
+
+    createdAt: {
+      sql: `created_at`,
+      type: `time`
+    }
+  },
+});

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

@@ -17,6 +17,16 @@ Array [
 ]
 `;
 
+exports[`SQL API Postgres (Data) no limit for non matching count push down: no limit for non matching count push down 1`] = `
+Array [
+  Object {
+    "a0": 2024-01-05T00:00:00.000Z,
+    "a1": 2024-01-01T00:00:00.000Z,
+    "a2": "100000",
+  },
+]
+`;
+
 exports[`SQL API Postgres (Data) powerbi min max push down: powerbi_min_max_push_down 1`] = `
 Array [
   Object {

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

@@ -180,5 +180,17 @@ from
   `);
       expect(res.rows).toMatchSnapshot('powerbi_min_max_push_down');
     });
+
+    test('no limit for non matching count push down', async () => {
+      const res = await connection.query(`
+    select
+      max("rows"."createdAt") as "a0",
+      min("rows"."createdAt") as "a1",
+      count(*) as "a2"
+    from
+      "public"."BigOrders" "rows"
+  `);
+      expect(res.rows).toMatchSnapshot('no limit for non matching count push down');
+    });
   });
 });

rust/cubesql/cubesql/src/compile/engine/df/planner.rs

@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use crate::config::ConfigObj;
 use async_trait::async_trait;
 use datafusion::{
     error::Result,
@@ -15,11 +16,20 @@ use super::scan::CubeScanExtensionPlanner;
 pub struct CubeQueryPlanner {
     pub transport: Arc<dyn TransportService>,
     pub meta: LoadRequestMeta,
+    pub config_obj: Arc<dyn ConfigObj>,
 }
 
 impl CubeQueryPlanner {
-    pub fn new(transport: Arc<dyn TransportService>, meta: LoadRequestMeta) -> Self {
-        Self { transport, meta }
+    pub fn new(
+        transport: Arc<dyn TransportService>,
+        meta: LoadRequestMeta,
+        config_obj: Arc<dyn ConfigObj>,
+    ) -> Self {
+        Self {
+            transport,
+            meta,
+            config_obj,
+        }
     }
 }
 
@@ -36,6 +46,7 @@ impl QueryPlanner for CubeQueryPlanner {
             CubeScanExtensionPlanner {
                 transport: self.transport.clone(),
                 meta: self.meta.clone(),
+                config_obj: self.config_obj.clone(),
             },
         )]);
         // Delegate most work of physical planning to the default physical planner

rust/cubesql/cubesql/src/compile/engine/df/scan.rs

@@ -34,6 +34,7 @@ use crate::{
         find_cube_scans_deep_search,
         rewrite::WrappedSelectType,
     },
+    config::ConfigObj,
     sql::AuthContextRef,
     transport::{CubeStreamReceiver, LoadRequestMeta, SpanId, TransportService},
     CubeError,
@@ -341,6 +342,7 @@ impl UserDefinedLogicalNode for WrappedSelectNode {
 pub struct CubeScanExtensionPlanner {
     pub transport: Arc<dyn TransportService>,
     pub meta: LoadRequestMeta,
+    pub config_obj: Arc<dyn ConfigObj>,
 }
 
 impl ExtensionPlanner for CubeScanExtensionPlanner {
@@ -369,6 +371,7 @@ impl ExtensionPlanner for CubeScanExtensionPlanner {
                     options: scan_node.options.clone(),
                     meta: self.meta.clone(),
                     span_id: scan_node.span_id.clone(),
+                    config_obj: self.config_obj.clone(),
                 }))
             } else if let Some(wrapper_node) = node.as_any().downcast_ref::<CubeScanWrapperNode>() {
                 // TODO
@@ -404,6 +407,7 @@ impl ExtensionPlanner for CubeScanExtensionPlanner {
                     options: scan_node.options.clone(),
                     meta: self.meta.clone(),
                     span_id: scan_node.span_id.clone(),
+                    config_obj: self.config_obj.clone(),
                 }))
             } else {
                 None
@@ -426,6 +430,7 @@ struct CubeScanExecutionPlan {
     // injected by extension planner
     meta: LoadRequestMeta,
     span_id: Option<Arc<SpanId>>,
+    config_obj: Arc<dyn ConfigObj>,
 }
 
 #[derive(Debug)]
@@ -582,14 +587,8 @@ impl ExecutionPlan for CubeScanExecutionPlan {
         _context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
         // TODO: move envs to config
-        let stream_mode = std::env::var("CUBESQL_STREAM_MODE")
-            .ok()
-            .map(|v| v.parse::<bool>().unwrap())
-            .unwrap_or(false);
-        let query_limit = std::env::var("CUBEJS_DB_QUERY_LIMIT")
-            .ok()
-            .map(|v| v.parse::<i32>().unwrap())
-            .unwrap_or(50000);
+        let stream_mode = self.config_obj.stream_mode();
+        let query_limit = self.config_obj.non_streaming_query_max_row_limit();
 
         let stream_mode = match (stream_mode, self.request.limit) {
             (true, None) => true,
@@ -1406,6 +1405,7 @@ mod tests {
             transport: get_test_transport(),
             meta: get_test_load_meta(DatabaseProtocol::PostgreSQL),
             span_id: None,
+            config_obj: crate::config::Config::test().config_obj(),
         };
 
         let runtime = Arc::new(

rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs

@@ -3,6 +3,7 @@ use crate::{
         engine::df::scan::{CubeScanNode, DataType, MemberField, WrappedSelectNode},
         rewrite::WrappedSelectType,
     },
+    config::ConfigObj,
     sql::AuthContextRef,
     transport::{
         AliasedColumn, LoadRequestMeta, MetaContext, SpanId, SqlGenerator, SqlTemplates,
@@ -107,6 +108,7 @@ pub struct CubeScanWrapperNode {
     pub request: Option<V1LoadRequestQuery>,
     pub member_fields: Option<Vec<MemberField>>,
     pub span_id: Option<Arc<SpanId>>,
+    pub config_obj: Arc<dyn ConfigObj>,
 }
 
 impl CubeScanWrapperNode {
@@ -115,6 +117,7 @@ impl CubeScanWrapperNode {
         meta: Arc<MetaContext>,
         auth_context: AuthContextRef,
         span_id: Option<Arc<SpanId>>,
+        config_obj: Arc<dyn ConfigObj>,
     ) -> Self {
         Self {
             wrapped_plan,
@@ -124,6 +127,7 @@ impl CubeScanWrapperNode {
             request: None,
             member_fields: None,
             span_id,
+            config_obj,
         }
     }
 
@@ -141,6 +145,7 @@ impl CubeScanWrapperNode {
             request: Some(request),
             member_fields: Some(member_fields),
             span_id: self.span_id.clone(),
+            config_obj: self.config_obj.clone(),
         }
     }
 }
@@ -214,11 +219,12 @@ impl CubeScanWrapperNode {
         load_request_meta: Arc<LoadRequestMeta>,
     ) -> result::Result<Self, CubeError> {
         let schema = self.schema();
+        let wrapped_plan = self.wrapped_plan.clone();
         let (sql, request, member_fields) = Self::generate_sql_for_node(
             Arc::new(self.clone()),
             transport,
             load_request_meta,
-            self.wrapped_plan.clone(),
+            self.clone().set_max_limit_for_node(wrapped_plan),
             true,
         )
         .await
@@ -256,6 +262,45 @@ impl CubeScanWrapperNode {
         Ok(self.with_sql_and_request(sql, request, member_fields))
     }
 
+    pub fn set_max_limit_for_node(self, node: Arc<LogicalPlan>) -> Arc<LogicalPlan> {
+        let stream_mode = self.config_obj.stream_mode();
+        if stream_mode {
+            return node;
+        }
+
+        let query_limit = self.config_obj.non_streaming_query_max_row_limit();
+        match node.as_ref() {
+            LogicalPlan::Extension(Extension {
+                node: extension_node,
+            }) => {
+                let cube_scan_node = extension_node
+                    .as_any()
+                    .downcast_ref::<CubeScanNode>()
+                    .cloned();
+                let wrapped_select_node = extension_node
+                    .as_any()
+                    .downcast_ref::<WrappedSelectNode>()
+                    .cloned();
+                if let Some(node) = cube_scan_node {
+                    let mut new_node = node.clone();
+                    new_node.request.limit = Some(query_limit);
+                    Arc::new(LogicalPlan::Extension(Extension {
+                        node: Arc::new(new_node),
+                    }))
+                } else if let Some(node) = wrapped_select_node {
+                    let mut new_node = node.clone();
+                    new_node.limit = Some(query_limit as usize);
+                    Arc::new(LogicalPlan::Extension(Extension {
+                        node: Arc::new(new_node),
+                    }))
+                } else {
+                    node.clone()
+                }
+            }
+            _ => node.clone(),
+        }
+    }
+
     pub fn generate_sql_for_node(
         plan: Arc<Self>,
         transport: Arc<dyn TransportService>,
@@ -1787,6 +1832,7 @@ impl UserDefinedLogicalNode for CubeScanWrapperNode {
             request: self.request.clone(),
             member_fields: self.member_fields.clone(),
             span_id: self.span_id.clone(),
+            config_obj: self.config_obj.clone(),
         })
     }
 }