fix(sqlserver): rewrite sql variant queries without panic

crates/dbx-core/src/db/sqlserver.rs

@@ -1090,13 +1090,33 @@ fn build_sqlserver_unsafe_type_query(
         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
+    // ordering and pagination survive the derived-table rewrite. If a sort key
-    // keys cannot be safely migrated, refuse the rewrite rather than silently
+    // is not part of the projection, keep the original ordering inside the
-    // dropping order semantics (callers fall back to the plain statement).
+    // derived table instead of falling back to the unsafe sql_variant query.
-    let order_by = match &statement.order_by {
+    let mut inner = statement.inner;
+    let order_by = match statement.order_by {
         Some(order_by) => {
-            let outer_order_by = sqlserver_outer_order_by(order_by, columns)?;
+            if let Some(outer_order_by) = sqlserver_outer_order_by(&order_by, columns) {
-            format!(" {outer_order_by}")
+                format!(" {outer_order_by}")
+            } else {
+                if !has_top_level_top(&inner) {
+                    // TOP cannot combine with OFFSET/FETCH in the same query
+                    // scope, and an OFFSET/FETCH tail already makes ORDER BY
+                    // legal inside a derived table, so the TOP (100) PERCENT
+                    // crutch is only needed for plain ORDER BY clauses.
+                    if !top_level_sqlserver_tokens(&order_by).iter().any(|token| token.text == "OFFSET") {
+                        inner = add_sqlserver_top_percent(&inner);
+                    }
+                    // A trailing `--` comment would swallow the appended ORDER
+                    // BY, so break the line first (the same hazard the closing
+                    // separator above guards against).
+                    let order_by_separator =
+                        if inner.lines().last().is_some_and(|line| line.contains("--")) { '\n' } else { ' ' };
+                    inner.push(order_by_separator);
+                    inner.push_str(&order_by);
+                }
+                String::new()
+            }
         }
         None => String::new(),
     };
@@ -1104,13 +1124,27 @@ fn build_sqlserver_unsafe_type_query(
     Some(SqlServerUnsafeTypeQuery {
         sql: format!(
             "{}SELECT {select_list} FROM ({}{inner_closing_line_break}) AS {source_alias}({source_alias_list}){order_by}",
-            statement.prefix, statement.inner
+            statement.prefix, inner
         ),
         spatial_columns,
         restored_columns,
     })
 }
 
+fn add_sqlserver_top_percent(sql: &str) -> String {
+    let tokens = top_level_sqlserver_tokens(sql);
+    let Some(select_index) = tokens.iter().position(|token| token.text == "SELECT") else {
+        return sql.to_string();
+    };
+    // TOP must follow the optional ALL/DISTINCT quantifier; inserting it right
+    // after SELECT would produce the invalid `SELECT TOP ... DISTINCT ...`.
+    let insert_end = match tokens.get(select_index + 1) {
+        Some(token) if matches!(token.text.as_str(), "ALL" | "DISTINCT") => token.start + token.text.len(),
+        _ => tokens[select_index].start + "SELECT".len(),
+    };
+    format!("{} TOP (100) PERCENT{}", &sql[..insert_end], &sql[insert_end..])
+}
+
 fn is_sqlserver_unsafe_column(column: &SqlServerDescribedColumn) -> bool {
     is_sqlserver_spatial_column(column)
         || is_sqlserver_hierarchyid_column(column)
@@ -1209,8 +1243,7 @@ fn normalized_sqlserver_select_statement(sql: &str) -> Option<SqlServerNormalize
 /// ordering cannot be guaranteed on the outer query. `None` covers ORDER BY keys
 /// that reference columns absent from the projection, non-trivial expressions,
 /// and keys targeting transformed columns whose original ordering cannot be
-/// preserved; callers must then fall back to executing the plain statement
+/// preserved; callers keep the original ordering inside the derived table.
-/// rather than silently dropping order and pagination semantics.
 fn sqlserver_outer_order_by(order_by: &str, columns: &[SqlServerDescribedColumn]) -> Option<String> {
     // OFFSET/FETCH carry no column references, so keep that tail verbatim and
     // rebuild only the sort-key list against the outer projection.
@@ -4942,67 +4975,212 @@ mod tests {
     }
 
     #[test]
-    fn sqlserver_falls_back_when_order_by_cannot_be_migrated() {
+    fn sqlserver_keeps_unmapped_order_by_inside_variant_rewrite() {
-        // ORDER BY references a column absent from the projection: cannot be
+        // ORDER BY references a column absent from the projection. Keep it in
-        // re-applied on the outer query, so the rewrite must be refused rather
+        // the inner query so the sql_variant result is still cast safely.
-        // than silently dropping order semantics.
+        let rewritten = build_sqlserver_unsafe_type_query(
-        assert_eq!(
+            "SELECT id, value FROM sys.extended_properties ORDER BY major_id",
-            build_sqlserver_unsafe_type_query(
+            &[
-                "SELECT polygon FROM dbo.t ORDER BY landId",
+                SqlServerDescribedColumn {
-                &[SqlServerDescribedColumn {
+                    name: Some("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.contains(
+            "FROM (SELECT TOP (100) PERCENT id, value FROM sys.extended_properties ORDER BY major_id) AS [dbx_unsafe_source]"
+        ));
+        assert!(!rewritten.sql.ends_with("ORDER BY [major_id]"));
+
+        // Non-trivial ORDER BY expressions use the same safe inner fallback.
+        let rewritten = build_sqlserver_unsafe_type_query(
+            "SELECT id, value FROM sys.extended_properties ORDER BY UPPER(id)",
+            &[
+                SqlServerDescribedColumn {
+                    name: Some("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.contains("ORDER BY UPPER(id)"));
+        assert!(rewritten.sql.contains("SELECT TOP (100) PERCENT"));
+
+        // ORDER BY targeting a rewritten geometry column keeps the native
+        // ordering inside the derived table rather than ordering by WKT.
+        let rewritten = build_sqlserver_unsafe_type_query(
+            "SELECT id, polygon FROM dbo.t ORDER BY polygon",
+            &[
+                SqlServerDescribedColumn {
+                    name: Some("id".to_string()),
+                    system_type_name: Some("int".to_string()),
+                    user_type_schema: None,
+                    user_type_name: None,
+                },
+                SqlServerDescribedColumn {
                     name: Some("polygon".to_string()),
                     system_type_name: Some("geometry".to_string()),
                     user_type_schema: Some("sys".to_string()),
                     user_type_name: Some("geometry".to_string()),
-                }],
+                },
-            ),
+            ],
-            None
+        )
-        );
+        .unwrap();
+        assert!(rewritten.sql.contains("ORDER BY polygon"));
+        assert!(rewritten.sql.contains("SELECT TOP (100) PERCENT"));
+    }
 
-        // Non-trivial ORDER BY expression: cannot be re-applied safely.
+    #[test]
-        assert_eq!(
+    fn sqlserver_skips_top_percent_when_inner_order_by_carries_offset_fetch() {
-            build_sqlserver_unsafe_type_query(
+        // TOP cannot combine with OFFSET/FETCH in the same query scope, and the
-                "SELECT id, polygon FROM dbo.t ORDER BY UPPER(id)",
+        // OFFSET/FETCH tail already makes ORDER BY legal inside the derived
-                &[
+        // table, so the inner fallback must not inject TOP (100) PERCENT.
-                    SqlServerDescribedColumn {
+        let rewritten = build_sqlserver_unsafe_type_query(
-                        name: Some("id".to_string()),
+            "SELECT id, value FROM sys.extended_properties ORDER BY major_id OFFSET 10 ROWS FETCH NEXT 5 ROWS ONLY",
-                        system_type_name: Some("int".to_string()),
+            &[
-                        user_type_schema: None,
+                SqlServerDescribedColumn {
-                        user_type_name: None,
+                    name: Some("id".to_string()),
-                    },
+                    system_type_name: Some("int".to_string()),
-                    SqlServerDescribedColumn {
+                    user_type_schema: None,
-                        name: Some("polygon".to_string()),
+                    user_type_name: None,
-                        system_type_name: Some("geometry".to_string()),
+                },
-                        user_type_schema: Some("sys".to_string()),
+                SqlServerDescribedColumn {
-                        user_type_name: Some("geometry".to_string()),
+                    name: Some("value".to_string()),
-                    },
+                    system_type_name: Some("sql_variant".to_string()),
-                ],
+                    user_type_schema: None,
-            ),
+                    user_type_name: None,
-            None
+                },
-        );
+            ],
+        )
+        .unwrap();
 
-        // ORDER BY targeting a rewritten geometry column would order by the WKT
+        assert!(rewritten.sql.contains(
-        // string instead of the geometry value: refuse rather than change semantics.
+            "FROM (SELECT id, value FROM sys.extended_properties ORDER BY major_id OFFSET 10 ROWS FETCH NEXT 5 ROWS ONLY) AS [dbx_unsafe_source]"
-        assert_eq!(
+        ));
-            build_sqlserver_unsafe_type_query(
+        assert!(!rewritten.sql.contains("TOP (100) PERCENT"));
-                "SELECT id, polygon FROM dbo.t ORDER BY polygon",
+    }
-                &[
+
-                    SqlServerDescribedColumn {
+    #[test]
-                        name: Some("id".to_string()),
+    fn sqlserver_adds_top_percent_after_distinct_quantifier() {
-                        system_type_name: Some("int".to_string()),
+        // TOP must follow the optional ALL/DISTINCT quantifier; inserting it
-                        user_type_schema: None,
+        // right after SELECT would produce invalid T-SQL.
-                        user_type_name: None,
+        let rewritten = build_sqlserver_unsafe_type_query(
-                    },
+            "SELECT DISTINCT id, value FROM sys.extended_properties ORDER BY major_id",
-                    SqlServerDescribedColumn {
+            &[
-                        name: Some("polygon".to_string()),
+                SqlServerDescribedColumn {
-                        system_type_name: Some("geometry".to_string()),
+                    name: Some("id".to_string()),
-                        user_type_schema: Some("sys".to_string()),
+                    system_type_name: Some("int".to_string()),
-                        user_type_name: Some("geometry".to_string()),
+                    user_type_schema: None,
-                    },
+                    user_type_name: None,
-                ],
+                },
-            ),
+                SqlServerDescribedColumn {
-            None
+                    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
+            .contains("SELECT DISTINCT TOP (100) PERCENT id, value FROM sys.extended_properties ORDER BY major_id"));
+        assert!(!rewritten.sql.contains("TOP (100) PERCENT DISTINCT"));
+
+        let rewritten_all = build_sqlserver_unsafe_type_query(
+            "SELECT ALL id, value FROM sys.extended_properties ORDER BY major_id",
+            &[
+                SqlServerDescribedColumn {
+                    name: Some("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_all.sql.contains("SELECT ALL TOP (100) PERCENT id, value FROM sys.extended_properties"));
+    }
+
+    #[test]
+    fn sqlserver_appends_inner_order_by_on_a_new_line_after_line_comment() {
+        // The derived table's last line ends with a `--` comment, so the
+        // appended ORDER BY must start on a new line instead of being
+        // swallowed by the comment.
+        let rewritten = build_sqlserver_unsafe_type_query(
+            "SELECT id, value\nFROM sys.extended_properties -- keep source comment\nORDER BY major_id",
+            &[
+                SqlServerDescribedColumn {
+                    name: Some("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.contains(
+            "FROM (SELECT TOP (100) PERCENT id, value\nFROM sys.extended_properties -- keep source comment\nORDER BY major_id\n) AS [dbx_unsafe_source]"
+        ));
+        assert!(!rewritten.sql.contains("-- keep source comment ORDER BY"));
+    }
+
+    #[test]
+    fn sqlserver_does_not_duplicate_top_query_order_by_in_variant_rewrite() {
+        let rewritten = build_sqlserver_unsafe_type_query(
+            "SELECT TOP 10 id, value FROM dbo.t ORDER BY source_id",
+            &[
+                SqlServerDescribedColumn {
+                    name: Some("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_eq!(rewritten.sql.matches("ORDER BY source_id").count(), 1);
+        assert!(rewritten.sql.contains("FROM (SELECT TOP 10 id, value FROM dbo.t ORDER BY source_id)"));
     }
 
     #[test]