diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index f567c76241..aa46c014e1 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -1078,6 +1078,8 @@ fn build_sqlserver_unsafe_type_query( .filter(|(_, column)| is_sqlserver_hierarchyid_column(column)) .map(|(column_index, _)| SqlServerRestoredColumn { column_index, column_type: "hierarchyid".to_string() }) .collect(); + let inner_closing_line_break = + if statement.inner.lines().last().is_some_and(|line| line.contains("--")) { "\n" } else { "" }; // Re-apply the original ORDER BY / OFFSET / FETCH on the outer query so // ordering and pagination survive the derived-table rewrite. When the sort @@ -1093,8 +1095,8 @@ fn build_sqlserver_unsafe_type_query( Some(SqlServerUnsafeTypeQuery { sql: format!( - "SELECT {select_list} FROM ({}) AS {source_alias}({source_alias_list}){order_by}", - statement.inner + "{}SELECT {select_list} FROM ({}{inner_closing_line_break}) AS {source_alias}({source_alias_list}){order_by}", + statement.prefix, statement.inner ), spatial_columns, restored_columns, @@ -1139,6 +1141,9 @@ fn is_sqlserver_variant_column(column: &SqlServerDescribedColumn) -> bool { } struct SqlServerNormalizedStatement { + /// Leading comments and, for CTE queries, the `WITH ...` declaration that + /// must remain immediately before the rewritten outer SELECT. + prefix: String, /// Statement with the trailing ORDER BY (and any OFFSET/FETCH) removed so it /// can be used as a derived table subquery. inner: String, @@ -1150,10 +1155,21 @@ struct SqlServerNormalizedStatement { fn normalized_sqlserver_select_statement(sql: &str) -> Option { let statement = trim_sqlserver_statement(sql); let trimmed = statement.trim_start(); - if trimmed.is_empty() || !trimmed.get(..6).is_some_and(|prefix| prefix.eq_ignore_ascii_case("SELECT")) { + if trimmed.is_empty() { return None; } - if has_top_level_select_into(trimmed) { + let statement_tokens = top_level_sqlserver_tokens(trimmed); + let first_token = statement_tokens.first()?; + let select_start = if first_token.text == "SELECT" { + first_token.start + } else if first_token.text == "WITH" { + statement_tokens.iter().find(|token| token.text == "SELECT")?.start + } else { + return None; + }; + let prefix = trimmed[..select_start].to_string(); + let select = &trimmed[select_start..]; + if has_top_level_select_into(select) { return None; } @@ -1165,19 +1181,19 @@ fn normalized_sqlserver_select_statement(sql: &str) -> Option bool { if statements.len() != 1 { return false; } - let statement = statements[0].trim_start(); - statement.get(..6).is_some_and(|prefix| prefix.eq_ignore_ascii_case("SELECT")) + normalized_sqlserver_select_statement(&statements[0]).is_some() } fn sqlserver_source_column_name(index: usize) -> String { @@ -5220,6 +5235,38 @@ mod tests { assert_eq!(rewritten.sql.matches("CAST(").count(), 1); } + #[test] + fn sqlserver_wraps_cte_sql_variant_columns_as_nvarchar() { + let sql = "-- table metadata\nWITH props AS (SELECT major_id, value FROM sys.extended_properties)\n\ + SELECT major_id, value FROM props -- keep source comment\nORDER BY major_id;"; + let rewritten = build_sqlserver_unsafe_type_query( + sql, + &[ + SqlServerDescribedColumn { + name: Some("major_id".to_string()), + system_type_name: Some("int".to_string()), + user_type_schema: None, + user_type_name: None, + }, + SqlServerDescribedColumn { + name: Some("value".to_string()), + system_type_name: Some("sql_variant".to_string()), + user_type_schema: None, + user_type_name: None, + }, + ], + ) + .unwrap(); + + assert!(rewritten.sql.starts_with("-- table metadata\nWITH props AS")); + assert!(rewritten.sql.contains("SELECT [major_id] = [dbx_unsafe_source].[dbx_col_1]")); + assert!(rewritten.sql.contains("[value] = CAST([dbx_unsafe_source].[dbx_col_2] AS NVARCHAR(MAX))")); + assert!(rewritten + .sql + .contains("FROM (SELECT major_id, value FROM props -- keep source comment\n) AS [dbx_unsafe_source]")); + assert!(rewritten.sql.ends_with("ORDER BY [major_id]")); + } + #[test] fn sqlserver_does_not_wrap_non_variant_columns() { assert_eq!( @@ -5320,6 +5367,11 @@ mod tests { let variant = super::execute_query(&mut client, sql).await.unwrap(); assert_eq!(variant.rows, vec![vec![serde_json::json!(1), serde_json::json!("legacy")]]); + let cte_sql = "WITH source AS (SELECT id, payload FROM #dbx_issue_4002) SELECT id, payload FROM source"; + assert!(super::is_single_sqlserver_select(cte_sql)); + let cte_variant = super::execute_query(&mut client, cte_sql).await.unwrap(); + assert_eq!(cte_variant.rows, vec![vec![serde_json::json!(1), serde_json::json!("legacy")]]); + let duplicate_sql = "SELECT id AS HJRQ, payload AS HJRQ FROM #dbx_issue_4002"; let duplicate_probe = super::sqlserver_legacy_probe(duplicate_sql).unwrap(); let duplicate_columns =