diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index 66aac91dfe..f8b7645b8e 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -354,6 +354,19 @@ fn restore_sqlserver_spatial_column_types(column_types: &mut [String], spatial_c } } +fn restore_sqlserver_column_types(column_types: &mut [String], restored_columns: &[SqlServerRestoredColumn]) { + for restored_column in restored_columns { + if let Some(column_type) = column_types.get_mut(restored_column.column_index) { + column_type.clone_from(&restored_column.column_type); + } + } +} + +fn restore_sqlserver_unsafe_column_types(column_types: &mut [String], query: &SqlServerUnsafeTypeQuery) { + restore_sqlserver_spatial_column_types(column_types, &query.spatial_columns); + restore_sqlserver_column_types(column_types, &query.restored_columns); +} + fn columns_from_metadata(metadata: &tiberius::ResultMetadata) -> Vec { metadata.columns().iter().map(|c| c.name().to_string()).collect() } @@ -557,7 +570,7 @@ async fn collect_first_result_limited( start: Instant, max_rows: Option, result_offset: usize, - spatial_columns: &[SqlServerSpatialColumn], + query: &SqlServerUnsafeTypeQuery, ) -> Result { let row_limit = query_result_row_limit(max_rows); let mut remaining_offset = result_offset; @@ -566,7 +579,7 @@ async fn collect_first_result_limited( let mut rows: Vec> = Vec::new(); let mut spatial_values: Vec>> = Vec::new(); let mut spatial_values_builder = - SpatialColumnBuilder::new(spatial_columns.iter().map(|column| column.column_index)); + SpatialColumnBuilder::new(query.spatial_columns.iter().map(|column| column.column_index)); let mut truncated = false; while let Some(item) = stream.try_next().await.map_err(|e| e.to_string())? { @@ -574,7 +587,7 @@ async fn collect_first_result_limited( QueryItem::Metadata(metadata) if metadata.result_index() == 0 => { columns = columns_from_metadata(&metadata); column_types = column_types_from_metadata(&metadata); - restore_sqlserver_spatial_column_types(&mut column_types, spatial_columns); + restore_sqlserver_unsafe_column_types(&mut column_types, query); } QueryItem::Metadata(_) => {} QueryItem::Row(row) if row.result_index() == 0 => { @@ -584,7 +597,7 @@ async fn collect_first_result_limited( } if rows.len() < row_limit { let (values, srids) = - row_to_json_with_spatial_metadata(&row, spatial_columns, |column_index, srid| { + row_to_json_with_spatial_metadata(&row, &query.spatial_columns, |column_index, srid| { spatial_values_builder.observe(column_index, srid); }); rows.push(values); @@ -597,7 +610,7 @@ async fn collect_first_result_limited( } } - restore_sqlserver_spatial_column_types(&mut column_types, spatial_columns); + restore_sqlserver_unsafe_column_types(&mut column_types, query); Ok(QueryResult { columns, @@ -648,15 +661,22 @@ struct SqlServerSpatialColumn { column_type: String, } +#[derive(Debug, Clone, PartialEq, Eq)] +struct SqlServerRestoredColumn { + column_index: usize, + column_type: String, +} + #[derive(Debug, Clone, PartialEq, Eq)] struct SqlServerUnsafeTypeQuery { sql: String, spatial_columns: Vec, + restored_columns: Vec, } impl SqlServerUnsafeTypeQuery { fn plain(sql: &str) -> Self { - Self { sql: sql.to_string(), spatial_columns: Vec::new() } + Self { sql: sql.to_string(), spatial_columns: Vec::new(), restored_columns: Vec::new() } } } @@ -1032,6 +1052,8 @@ fn build_sqlserver_unsafe_type_query( format!( "{quoted_output} = CASE WHEN {value_ref} IS NULL THEN NULL ELSE N'SRID=' + CONVERT(nvarchar(20), {value_ref}.STSrid) + N';' + {value_ref}.AsTextZM() END" ) + } else if is_sqlserver_hierarchyid_column(column) { + format!("{quoted_output} = CASE WHEN {value_ref} IS NULL THEN NULL ELSE {value_ref}.ToString() END") } else if is_sqlserver_variant_column(column) { format!("{quoted_output} = CAST({value_ref} AS NVARCHAR(MAX))") } else { @@ -1049,6 +1071,12 @@ fn build_sqlserver_unsafe_type_query( .map(|column_type| SqlServerSpatialColumn { column_index, column_type: column_type.to_string() }) }) .collect(); + let restored_columns = columns + .iter() + .enumerate() + .filter(|(_, column)| is_sqlserver_hierarchyid_column(column)) + .map(|(column_index, _)| SqlServerRestoredColumn { column_index, column_type: "hierarchyid".to_string() }) + .collect(); // Re-apply the original ORDER BY / OFFSET / FETCH on the outer query so // ordering and pagination survive the derived-table rewrite. When the sort @@ -1068,11 +1096,14 @@ fn build_sqlserver_unsafe_type_query( statement.inner ), spatial_columns, + restored_columns, }) } fn is_sqlserver_unsafe_column(column: &SqlServerDescribedColumn) -> bool { - is_sqlserver_spatial_column(column) || is_sqlserver_variant_column(column) + is_sqlserver_spatial_column(column) + || is_sqlserver_hierarchyid_column(column) + || is_sqlserver_variant_column(column) } fn is_sqlserver_spatial_column(column: &SqlServerDescribedColumn) -> bool { @@ -1092,6 +1123,13 @@ fn sqlserver_spatial_column_type(column: &SqlServerDescribedColumn) -> Option<&' }) } +fn is_sqlserver_hierarchyid_column(column: &SqlServerDescribedColumn) -> bool { + [&column.system_type_name, &column.user_type_name].into_iter().flatten().any(|name| { + let normalized = name.trim().trim_matches(['[', ']']).to_ascii_lowercase(); + normalized == "hierarchyid" || normalized.ends_with(".hierarchyid") + }) +} + fn is_sqlserver_variant_column(column: &SqlServerDescribedColumn) -> bool { [&column.system_type_name, &column.user_type_name].into_iter().flatten().any(|name| { let normalized = name.trim().trim_matches(['[', ']']).to_ascii_lowercase(); @@ -1145,9 +1183,9 @@ fn normalized_sqlserver_select_statement(sql: &str) -> Option Option { // OFFSET/FETCH carry no column references, so keep that tail verbatim and // rebuild only the sort-key list against the outer projection. @@ -1173,8 +1211,7 @@ fn sqlserver_outer_order_by(order_by: &str, columns: &[SqlServerDescribedColumn] let mut parts = Vec::with_capacity(order_exprs.len()); for order_expr in order_exprs { - let output_name = sqlserver_order_by_output_name(columns, &order_expr.expr)?; - let mut part = quote_sqlserver_identifier(&output_name); + let mut part = sqlserver_outer_order_by_expression(columns, &order_expr.expr)?; if order_expr.options.asc == Some(false) { part.push_str(" DESC"); } @@ -1192,11 +1229,10 @@ fn sqlserver_outer_order_by(order_by: &str, columns: &[SqlServerDescribedColumn] Some(outer) } -/// Map an ORDER BY sort key to the outer rewrite's output column name. Returns -/// `None` for keys that do not resolve to a projection column, keys targeting a -/// column the rewrite transforms (spatial/variant), and any non-trivial -/// expression, none of which can be safely re-applied on the outer query. -fn sqlserver_order_by_output_name(columns: &[SqlServerDescribedColumn], expr: &Expr) -> Option { +/// Map an ORDER BY sort key to an expression on the outer rewrite. hierarchyid +/// uses the untouched derived-table value so converting its displayed value to +/// a logical path never changes SQL Server's native depth-first ordering. +fn sqlserver_outer_order_by_expression(columns: &[SqlServerDescribedColumn], expr: &Expr) -> Option { let index = match expr { Expr::Identifier(identifier) => sqlserver_projection_column_index(columns, &identifier.value)?, Expr::CompoundIdentifier(identifiers) => { @@ -1212,10 +1248,15 @@ fn sqlserver_order_by_output_name(columns: &[SqlServerDescribedColumn], expr: &E _ => return None, }; let column = columns.get(index)?; + if is_sqlserver_hierarchyid_column(column) { + let source_alias = quote_sqlserver_identifier("dbx_unsafe_source"); + let source_column = quote_sqlserver_identifier(&sqlserver_source_column_name(index)); + return Some(format!("{source_alias}.{source_column}")); + } if is_sqlserver_unsafe_column(column) { return None; } - Some(sqlserver_output_column_name(column, index)) + Some(quote_sqlserver_identifier(&sqlserver_output_column_name(column, index))) } fn sqlserver_projection_column_index(columns: &[SqlServerDescribedColumn], name: &str) -> Option { @@ -1597,7 +1638,7 @@ pub async fn stream_first_result_set( active_result_index = Some(metadata.result_index()); columns = columns_from_metadata(&metadata); column_types = column_types_from_metadata(&metadata); - restore_sqlserver_spatial_column_types(&mut column_types, &query.spatial_columns); + restore_sqlserver_unsafe_column_types(&mut column_types, &query); on_item(SqlServerStreamItem::Columns { columns: &columns, column_types: &column_types })?; columns_emitted = true; } @@ -1607,7 +1648,7 @@ pub async fn stream_first_result_set( active_result_index = Some(row.result_index()); columns = row.columns().iter().map(|c| c.name().to_string()).collect(); column_types = row.columns().iter().map(sqlserver_column_type_name).collect(); - restore_sqlserver_spatial_column_types(&mut column_types, &query.spatial_columns); + restore_sqlserver_unsafe_column_types(&mut column_types, &query); on_item(SqlServerStreamItem::Columns { columns: &columns, column_types: &column_types })?; columns_emitted = true; } @@ -2839,14 +2880,7 @@ pub async fn execute_query_with_max_rows( }; let (result, messages) = capture_sqlserver_messages(async { let stream = sqlserver_driver_result(client.query(query.sql.as_str(), &[])).await?; - sqlserver_driver_result(collect_first_result_limited( - stream, - start, - max_rows, - result_offset, - &query.spatial_columns, - )) - .await + sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows, result_offset, &query)).await }) .await; let mut result = query_result_with_server_messages(result?, messages); @@ -2932,7 +2966,7 @@ pub(crate) async fn execute_batch_with_max_rows_metadata( start, max_rows, result_offset, - &query.spatial_columns, + &query, )) .await }) @@ -3019,7 +3053,8 @@ async fn execute_simple_batch_first_result_with_max_rows( let result_offset = crate::query_result_sql::sqlserver_result_offset(sql); let (result, messages) = capture_sqlserver_messages(async { let stream = sqlserver_driver_result(client.simple_query(sql)).await?; - sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows, result_offset, &[])).await + let query = SqlServerUnsafeTypeQuery::plain(sql); + sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows, result_offset, &query)).await }) .await; let mut result = query_result_with_server_messages(result?, messages); @@ -3271,16 +3306,17 @@ mod tests { is_sqlserver_legacy_duplicate_probe_error, is_sqlserver_spatial_column, is_sqlserver_variant_column, push_sqlserver_ordered_events, query_result_with_server_messages, query_result_with_server_messages_metadata, requires_simple_query_batch, restore_sqlserver_legacy_probe_output_names, - restore_sqlserver_spatial_column_types, server_messages_query_result, sqlserver_batch_can_use_execute, - sqlserver_bulk_token_row, sqlserver_cell_to_json, sqlserver_columns_sql, sqlserver_completion_assistant_sql, - sqlserver_dml_output_returns_rows, sqlserver_done_trace_event, sqlserver_filter_definition_error, - sqlserver_hidden_schema_names, sqlserver_indexes_sql, sqlserver_legacy_indexes_sql, sqlserver_legacy_probe, - sqlserver_legacy_probe_with_nonce, sqlserver_legacy_wildcard_metadata_query, sqlserver_list_objects_sql, - sqlserver_list_schemas_sql, sqlserver_list_tables_sql, sqlserver_probe_explicit_alias, - sqlserver_query_messages, sqlserver_schema_name_predicate, sqlserver_spatial_marker, - sqlserver_supports_session_database_switch, sqlserver_table_comment_sql, sqlserver_triggers_sql, - sqlserver_visible_object_predicate, strip_dbx_sqlserver_row_number_column, SqlServerDescribedColumn, - SqlServerProbeOutputNameOverride, SqlServerResultSet, SqlServerSpatialColumn, SqlServerTdsEvent, + restore_sqlserver_spatial_column_types, restore_sqlserver_unsafe_column_types, server_messages_query_result, + sqlserver_batch_can_use_execute, sqlserver_bulk_token_row, sqlserver_cell_to_json, sqlserver_columns_sql, + sqlserver_completion_assistant_sql, sqlserver_dml_output_returns_rows, sqlserver_done_trace_event, + sqlserver_filter_definition_error, sqlserver_hidden_schema_names, sqlserver_indexes_sql, + sqlserver_legacy_indexes_sql, sqlserver_legacy_probe, sqlserver_legacy_probe_with_nonce, + sqlserver_legacy_wildcard_metadata_query, sqlserver_list_objects_sql, sqlserver_list_schemas_sql, + sqlserver_list_tables_sql, sqlserver_probe_explicit_alias, sqlserver_query_messages, + sqlserver_schema_name_predicate, sqlserver_spatial_marker, sqlserver_supports_session_database_switch, + sqlserver_table_comment_sql, sqlserver_triggers_sql, sqlserver_visible_object_predicate, + strip_dbx_sqlserver_row_number_column, SqlServerDescribedColumn, SqlServerProbeOutputNameOverride, + SqlServerRestoredColumn, SqlServerResultSet, SqlServerSpatialColumn, SqlServerTdsEvent, SQLSERVER_COMPLETION_CONTEXT_SQL, SQLSERVER_RESULT_TYPE_PROBE_SQL, }; use crate::types::{ @@ -3890,7 +3926,7 @@ mod tests { let first_result = source.split("async fn execute_simple_batch_first_result_with_max_rows").nth(1).unwrap(); let first_result = first_result.split("fn strip_dbx_sqlserver_row_number_column").next().unwrap(); - assert!(first_result.contains("collect_first_result_limited(stream, start, max_rows, result_offset, &[])")); + assert!(first_result.contains("collect_first_result_limited(stream, start, max_rows, result_offset, &query)")); assert!(!first_result.contains("collect_result_sets_limited")); } @@ -4475,6 +4511,74 @@ mod tests { rewritten.spatial_columns, vec![SqlServerSpatialColumn { column_index: 1, column_type: "geometry".to_string() }] ); + assert!(rewritten.restored_columns.is_empty()); + } + + #[test] + fn sqlserver_wraps_hierarchyid_columns_as_logical_paths() { + let rewritten = build_sqlserver_unsafe_type_query( + "SELECT id, path, label FROM dbo.nodes ORDER BY 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("path".to_string()), + system_type_name: Some("hierarchyid".to_string()), + user_type_schema: Some("sys".to_string()), + user_type_name: Some("hierarchyid".to_string()), + }, + SqlServerDescribedColumn { + name: Some("label".to_string()), + system_type_name: Some("nvarchar(40)".to_string()), + user_type_schema: None, + user_type_name: None, + }, + ], + ) + .unwrap(); + + assert_eq!( + rewritten.sql, + "SELECT [id] = [dbx_unsafe_source].[dbx_col_1], [path] = CASE WHEN [dbx_unsafe_source].[dbx_col_2] IS NULL THEN NULL ELSE [dbx_unsafe_source].[dbx_col_2].ToString() END, [label] = [dbx_unsafe_source].[dbx_col_3] FROM (SELECT id, path, label FROM dbo.nodes) AS [dbx_unsafe_source]([dbx_col_1], [dbx_col_2], [dbx_col_3]) ORDER BY [id]" + ); + assert!(rewritten.spatial_columns.is_empty()); + assert_eq!( + rewritten.restored_columns, + vec![SqlServerRestoredColumn { column_index: 1, column_type: "hierarchyid".to_string() }] + ); + + let mut column_types = vec!["int".to_string(), "nvarchar".to_string(), "nvarchar".to_string()]; + restore_sqlserver_unsafe_column_types(&mut column_types, &rewritten); + assert_eq!(column_types, vec!["int", "hierarchyid", "nvarchar"]); + } + + #[test] + fn sqlserver_orders_hierarchyid_by_the_original_server_value() { + let rewritten = build_sqlserver_unsafe_type_query( + "SELECT id, path FROM dbo.nodes ORDER BY path", + &[ + SqlServerDescribedColumn { + name: Some("id".to_string()), + system_type_name: Some("int".to_string()), + user_type_schema: None, + user_type_name: None, + }, + SqlServerDescribedColumn { + name: Some("path".to_string()), + system_type_name: Some("sys.hierarchyid".to_string()), + user_type_schema: Some("sys".to_string()), + user_type_name: Some("hierarchyid".to_string()), + }, + ], + ) + .unwrap(); + + assert!(rewritten.sql.contains("[path] = CASE WHEN")); + assert!(rewritten.sql.ends_with("ORDER BY [dbx_unsafe_source].[dbx_col_2]")); } #[test]