diff --git a/crates/dbx-core/src/query_result_sql.rs b/crates/dbx-core/src/query_result_sql.rs index 57f2d4fad..f367ee4e4 100644 --- a/crates/dbx-core/src/query_result_sql.rs +++ b/crates/dbx-core/src/query_result_sql.rs @@ -129,6 +129,10 @@ pub fn build_query_pagination_execution_plan( return plan; } + if options.database_type == Some(DatabaseType::SqlServer) && starts_with_cte(&options.query_base_sql) { + return plan; + } + if options.use_agent_cursor && options.pagination.offset == 0 { plan.sql_to_execute = options.query_base_sql; plan.page_limit = Some(options.pagination.limit); @@ -287,7 +291,7 @@ fn single_selectable_statement(original_sql: &str) -> Result { if statement.len() != base_sql.trim_end_matches(';').trim().len() { return Err(()); } - let upper = statement.trim_start().to_ascii_uppercase(); + let upper = statement.trim_start_matches(';').trim_start().to_ascii_uppercase(); if !(upper.starts_with("SELECT") || upper.starts_with("WITH")) { return Err(()); } @@ -298,6 +302,10 @@ fn single_selectable_statement(original_sql: &str) -> Result { Ok(statement) } +fn starts_with_cte(sql: &str) -> bool { + sql.trim_start().trim_start_matches(';').trim_start().to_ascii_uppercase().starts_with("WITH") +} + fn single_statement_error_reason(original_sql: &str) -> &'static str { let base_sql = original_sql.trim(); if base_sql.is_empty() { @@ -641,6 +649,35 @@ mod tests { assert_eq!(result.sql.unwrap(), "SELECT TOP 1000 * FROM TicketInfo"); } + #[test] + fn sqlserver_cte_pagination_plan_executes_original_sql() { + let sql = ";WITH ranked AS (SELECT id FROM dbo.users) SELECT * FROM ranked".to_string(); + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: sql.clone(), + query_base_sql: sql.clone(), + database_type: Some(DatabaseType::SqlServer), + pagination: QueryPagination { limit: 100, offset: 0, session_id: None }, + use_agent_cursor: false, + }); + + assert_eq!(plan.sql_to_execute, sql); + assert!(plan.page_sql.is_none()); + assert!(plan.count_sql.is_none()); + assert_eq!(plan.page_limit, None); + assert_eq!(plan.page_offset, None); + } + + #[test] + fn sqlserver_cte_count_query_is_not_wrapped_as_derived_table() { + let result = build_count_query_sql(CountQuerySqlOptions { + original_sql: ";WITH cte AS (SELECT 1 AS id) SELECT * FROM cte".to_string(), + database_type: Some(DatabaseType::SqlServer), + }); + + assert_eq!(result.ok, false); + assert!(result.sql.is_none()); + } + #[test] fn wraps_sqlserver_select_with_unnamed_column() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions {