fix(sqlserver): rewrite sql_variant CTE queries to avoid crash

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<SqlServerNormalizedStatement> {
     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<SqlServerNormalize
     // stay inside the derived table too (TOP makes ORDER BY legal in subqueries
     // on every SQL Server version) while the outer rewrite re-applies it for the
     // final result order.
-    let tokens = top_level_sqlserver_tokens(trimmed);
+    let tokens = top_level_sqlserver_tokens(select);
     let order_index = (0..tokens.len().saturating_sub(1))
         .rev()
         .find(|index| tokens[*index].text == "ORDER" && tokens.get(index + 1).is_some_and(|token| token.text == "BY"));
     let Some(order_index) = order_index else {
-        return Some(SqlServerNormalizedStatement { inner: trimmed.to_string(), order_by: None });
+        return Some(SqlServerNormalizedStatement { prefix, inner: select.to_string(), order_by: None });
     };
-    let order_by = trimmed[tokens[order_index].start..].trim().to_string();
-    if has_top_level_top(trimmed) {
-        return Some(SqlServerNormalizedStatement { inner: trimmed.to_string(), order_by: Some(order_by) });
+    let order_by = select[tokens[order_index].start..].trim().to_string();
+    if has_top_level_top(select) {
+        return Some(SqlServerNormalizedStatement { prefix, inner: select.to_string(), order_by: Some(order_by) });
     }
-    let inner = trimmed[..tokens[order_index].start].trim_end().to_string();
-    Some(SqlServerNormalizedStatement { inner, order_by: Some(order_by) })
+    let inner = select[..tokens[order_index].start].trim_end().to_string();
+    Some(SqlServerNormalizedStatement { prefix, inner, order_by: Some(order_by) })
 }
 
 /// Translate the original trailing ORDER BY clause (with optional OFFSET/FETCH)
@@ -1279,8 +1295,7 @@ fn is_single_sqlserver_select(sql: &str) -> 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 =