diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index c59eefa50..d75ab0af8 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -769,6 +769,19 @@ mod tests { assert!(!supports_streaming_offset_pagination(&oracle_req, 100)); } + #[test] + fn clickhouse_scalar_with_query_supports_streaming_pagination() { + let sql = "WITH 1 AS min_id SELECT dept, COUNT(*) FROM employees WHERE id >= min_id GROUP BY dept"; + let req = QueryResultExportRequest { + sql: sql.to_string(), + query_base_sql: sql.to_string(), + database_type: DatabaseType::ClickHouse, + ..request("csv", Some(1000), None) + }; + + assert!(supports_streaming_offset_pagination(&req, 100)); + } + #[test] fn keyset_candidate_accepts_simple_single_table_wildcard_query() { let candidate = safe_keyset_candidate("SELECT * FROM public.users").expect("safe keyset candidate"); diff --git a/crates/dbx-core/src/query_result_sql.rs b/crates/dbx-core/src/query_result_sql.rs index c6a0c50dd..ea8351fdd 100644 --- a/crates/dbx-core/src/query_result_sql.rs +++ b/crates/dbx-core/src/query_result_sql.rs @@ -352,33 +352,17 @@ fn cte_main_statement_is_select(sql: &str) -> bool { index += 1; } - loop { - if tokens.get(index).is_some_and(|token| token.text != "AS") { - index += 1; - } - if tokens.get(index).is_none_or(|token| token.text != "AS") { - return false; + while let Some(token) = tokens.get(index) { + if is_with_main_statement_keyword(&token.text) { + return token.text == "SELECT"; } index += 1; - - if tokens.get(index).is_some_and(|token| token.text == "NOT") { - index += 1; - if tokens.get(index).is_none_or(|token| token.text != "MATERIALIZED") { - return false; - } - index += 1; - } else if tokens.get(index).is_some_and(|token| token.text == "MATERIALIZED") { - index += 1; - } - - let Some(token) = tokens.get(index) else { - return false; - }; - if tokens.get(index + 1).is_some_and(|next| next.text == "AS") { - continue; - } - return token.text == "SELECT"; } + false +} + +fn is_with_main_statement_keyword(token: &str) -> bool { + matches!(token, "SELECT" | "INSERT" | "UPDATE" | "DELETE" | "MERGE") } fn single_statement_error_reason(original_sql: &str) -> &'static str { @@ -1440,6 +1424,50 @@ WHERE u.id = picked.id; ); } + #[test] + fn clickhouse_scalar_with_select_is_paginated() { + let sql = "WITH 1 AS min_id SELECT dept, COUNT(*) FROM employees WHERE id >= min_id GROUP BY dept"; + let result = build_paginated_query_sql(PaginatedQuerySqlOptions { + original_sql: sql.to_string(), + database_type: Some(DatabaseType::ClickHouse), + limit: 50, + offset: 100, + }); + + assert!(result.ok); + assert_eq!( + result.sql.unwrap(), + "WITH 1 AS min_id SELECT dept, COUNT(*) FROM employees WHERE id >= min_id GROUP BY dept LIMIT 50 OFFSET 100;" + ); + } + + #[test] + fn clickhouse_scalar_with_select_can_be_counted() { + let result = build_count_query_sql(CountQuerySqlOptions { + original_sql: "WITH 1 AS min_id SELECT dept, COUNT(*) FROM employees WHERE id >= min_id GROUP BY dept" + .to_string(), + database_type: Some(DatabaseType::ClickHouse), + }); + + assert!(result.ok); + assert_eq!( + result.sql.unwrap(), + "SELECT COUNT(*) AS dbx_total_rows FROM (WITH 1 AS min_id SELECT dept, COUNT(*) FROM employees WHERE id >= min_id GROUP BY dept) `dbx_count`;" + ); + } + + #[test] + fn clickhouse_scalar_with_update_is_not_paginated() { + let result = build_paginated_query_sql(PaginatedQuerySqlOptions { + original_sql: "WITH 1 AS min_id UPDATE employees SET dept = 'sales' WHERE id = min_id".to_string(), + database_type: Some(DatabaseType::ClickHouse), + limit: 50, + offset: 0, + }); + + assert_eq!(result, err("not_select")); + } + #[test] fn wraps_sqlserver_select_with_unnamed_column() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions {