diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index 56b6adcb5..de8bcf299 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -23,6 +23,48 @@ pub const SQLSERVER_DRIVER_PANIC_ERROR_PREFIX: &str = "SQL Server driver panic:" pub const SQLSERVER_LEGACY_DRIVER_PROFILE: &str = "sqlserver-legacy"; pub const SQLSERVER_LEGACY_DRIVER_LABEL: &str = "SQL Server legacy compatibility component"; const SQLSERVER_DEFAULT_PORT: u16 = 1433; +const SQLSERVER_UNSAFE_PROBE_BLOCK_ERROR_PREFIX: &str = "SQL Server unsafe result type:"; +const SQLSERVER_RESULT_TYPE_PROBE_SQL: &str = "\ + DECLARE @dbx_use_describe_dmv bit = CASE \ + WHEN @P2 = 0 AND OBJECT_ID(N'sys.dm_exec_describe_first_result_set') IS NOT NULL THEN 1 ELSE 0 END; \ + SELECT @dbx_use_describe_dmv AS dbx_use_describe_dmv; \ + IF @dbx_use_describe_dmv = 1 \ + BEGIN \ + EXEC sys.sp_executesql \ + N'SELECT name, system_type_name, user_type_schema, user_type_name \ + FROM sys.dm_exec_describe_first_result_set(@sql, NULL, 0) \ + WHERE error_number IS NULL AND is_hidden = 0 \ + ORDER BY column_ordinal', \ + N'@sql nvarchar(max)', @sql = @P1 \ + END \ + ELSE \ + BEGIN \ + DECLARE @dbx_probe_table sysname = N'##dbx_result_type_probe_' + CONVERT(nvarchar(12), @@SPID); \ + DECLARE @dbx_probe_object nvarchar(258) = N'tempdb..' + QUOTENAME(@dbx_probe_table); \ + DECLARE @dbx_probe_sql nvarchar(max); \ + BEGIN TRY \ + SET @dbx_probe_sql = N'SELECT TOP (0) * INTO ' + QUOTENAME(@dbx_probe_table) + \ + N' FROM (' + @P3 + N') AS dbx_probe_source'; \ + EXEC sys.sp_executesql @dbx_probe_sql; \ + SELECT c.name, TYPE_NAME(c.system_type_id) AS system_type_name, \ + SCHEMA_NAME(t.schema_id) AS user_type_schema, t.name AS user_type_name \ + FROM tempdb.sys.columns c \ + JOIN tempdb.sys.types t ON c.user_type_id = t.user_type_id \ + WHERE c.object_id = OBJECT_ID(@dbx_probe_object) \ + ORDER BY c.column_id; \ + SET @dbx_probe_sql = N'DROP TABLE ' + QUOTENAME(@dbx_probe_table); \ + EXEC sys.sp_executesql @dbx_probe_sql; \ + END TRY \ + BEGIN CATCH \ + IF OBJECT_ID(@dbx_probe_object) IS NOT NULL \ + BEGIN \ + SET @dbx_probe_sql = N'DROP TABLE ' + QUOTENAME(@dbx_probe_table); \ + EXEC sys.sp_executesql @dbx_probe_sql; \ + END \ + DECLARE @dbx_probe_error nvarchar(2048) = ERROR_MESSAGE(); \ + RAISERROR(@dbx_probe_error, 16, 1); \ + END CATCH \ + END"; const SIMPLE_QUERY_MODULE_KEYWORDS: &[&str] = &["FUNCTION", "PROC", "PROCEDURE", "TRIGGER", "VIEW"]; // Match JDBC/tiberius `encrypt=false`: encrypt only login, then drop back to raw TDS. const SQLSERVER_LEGACY_ENCRYPTION_LEVEL: tiberius::EncryptionLevel = tiberius::EncryptionLevel::Off; @@ -398,31 +440,76 @@ pub fn is_driver_panic_error(error: &str) -> bool { async fn describe_sqlserver_result_set( client: &mut SqlServerClient, sql: &str, + legacy_sql: &str, ) -> Result, String> { - let describe_sql = "\ - SELECT name, system_type_name, user_type_schema, user_type_name \ - FROM sys.dm_exec_describe_first_result_set(@P1, NULL, 0) \ - WHERE error_number IS NULL AND is_hidden = 0 \ - ORDER BY column_ordinal"; - let stream = sqlserver_driver_result(client.query(describe_sql, &[&sql])).await?; - let rows = sqlserver_driver_result(stream.into_first_result()).await?; + describe_sqlserver_result_set_with_mode(client, sql, legacy_sql, false).await +} - Ok(rows - .iter() - .map(|row| SqlServerDescribedColumn { - name: row.try_get::<&str, _>(0).ok().flatten().map(str::to_string), - system_type_name: row.try_get::<&str, _>(1).ok().flatten().map(str::to_string), - user_type_schema: row.try_get::<&str, _>(2).ok().flatten().map(str::to_string), - user_type_name: row.try_get::<&str, _>(3).ok().flatten().map(str::to_string), - }) - .collect()) +async fn describe_sqlserver_result_set_with_mode( + client: &mut SqlServerClient, + sql: &str, + legacy_sql: &str, + force_legacy: bool, +) -> Result, String> { + // SQL Server 2008 has no first-result-set DMV. Keep one probe round trip by + // selecting the modern DMV path server-side and using metadata-only execution otherwise. + let force_legacy = i32::from(force_legacy); + let mut stream = + sqlserver_driver_result(client.query(SQLSERVER_RESULT_TYPE_PROBE_SQL, &[&sql, &force_legacy, &legacy_sql])) + .await?; + let mut active_result_index = None; + let mut uses_describe_dmv = None; + let mut rows = Vec::new(); + + loop { + let item = match sqlserver_driver_result(stream.try_next()).await { + Ok(item) => item, + Err(error) if uses_describe_dmv == Some(false) => { + return Err(format!("{SQLSERVER_UNSAFE_PROBE_BLOCK_ERROR_PREFIX} {error}")); + } + Err(error) => return Err(error), + }; + let Some(item) = item else { + break; + }; + match item { + QueryItem::Metadata(result_metadata) => { + active_result_index = Some(result_metadata.result_index()); + } + QueryItem::Row(row) if active_result_index == Some(0) => { + uses_describe_dmv = row.try_get::(0).ok().flatten(); + } + QueryItem::Row(row) => rows.push(row), + } + } + + if uses_describe_dmv.is_none() { + return Err("SQL Server result type probe did not report its compatibility mode".to_string()); + } + Ok(rows.iter().map(sqlserver_described_column_from_row).collect()) +} + +fn sqlserver_described_column_from_row(row: &Row) -> SqlServerDescribedColumn { + SqlServerDescribedColumn { + name: row.try_get::<&str, _>(0).ok().flatten().map(str::to_string), + system_type_name: row.try_get::<&str, _>(1).ok().flatten().map(str::to_string), + user_type_schema: row.try_get::<&str, _>(2).ok().flatten().map(str::to_string), + user_type_name: row.try_get::<&str, _>(3).ok().flatten().map(str::to_string), + } +} + +fn is_blocking_sqlserver_unsafe_probe_error(error: &str) -> bool { + error.starts_with(SQLSERVER_UNSAFE_PROBE_BLOCK_ERROR_PREFIX) } async fn sqlserver_unsafe_type_query(client: &mut SqlServerClient, sql: &str) -> Result, String> { if !is_single_sqlserver_select(sql) { return Ok(None); } - let columns = describe_sqlserver_result_set(client, sql).await?; + let Some(legacy_sql) = normalized_sqlserver_select_statement(sql) else { + return Ok(None); + }; + let columns = describe_sqlserver_result_set(client, sql, &legacy_sql).await?; Ok(build_sqlserver_unsafe_type_query(sql, &columns)) } @@ -742,7 +829,9 @@ pub async fn stream_first_result_set( ) -> Result { let query_sql = match sqlserver_unsafe_type_query(client, sql).await { Ok(Some(sql)) => sql, - Ok(None) | Err(_) => sql.to_string(), + Ok(None) => sql.to_string(), + Err(error) if is_blocking_sqlserver_unsafe_probe_error(&error) => return Err(error), + Err(_) => sql.to_string(), }; let mut stream = sqlserver_driver_result(client.query(query_sql.as_str(), &[])).await?; let mut active_result_index: Option = None; @@ -1832,7 +1921,9 @@ pub async fn execute_query_with_max_rows( if starts_with_executable_sql_keyword(sql, &["SELECT", "EXEC", "WITH", "TABLE"]) { let query_sql = match sqlserver_unsafe_type_query(client, sql).await { Ok(Some(sql)) => sql, - Ok(None) | Err(_) => sql.to_string(), + Ok(None) => sql.to_string(), + Err(error) if is_blocking_sqlserver_unsafe_probe_error(&error) => return Err(error), + Err(_) => sql.to_string(), }; let (result, messages) = capture_sqlserver_messages(async { let stream = sqlserver_driver_result(client.query(query_sql.as_str(), &[])).await?; @@ -1913,17 +2004,21 @@ pub async fn execute_batch_with_max_rows( } if is_single_sqlserver_select(sql) { - if let Ok(Some(query_sql)) = sqlserver_unsafe_type_query(client, sql).await { - 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)).await - }) - .await; - return result.map(|result| { - let mut result = query_result_with_server_messages(result, messages); - strip_dbx_sqlserver_row_number_column(&mut result, sql); - vec![result] - }); + match sqlserver_unsafe_type_query(client, sql).await { + Ok(Some(query_sql)) => { + 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)).await + }) + .await; + return result.map(|result| { + let mut result = query_result_with_server_messages(result, messages); + strip_dbx_sqlserver_row_number_column(&mut result, sql); + vec![result] + }); + } + Err(error) if is_blocking_sqlserver_unsafe_probe_error(&error) => return Err(error), + Ok(None) | Err(_) => {} } } execute_simple_batch_with_max_rows(client, sql, max_rows).await @@ -2133,12 +2228,14 @@ fn first_sql_tokens(sql: &str, limit: usize) -> Vec { mod tests { use super::{ build_sqlserver_unsafe_type_query, capture_sqlserver_messages, format_sqlserver_numeric, - is_sqlserver_spatial_column, is_sqlserver_variant_column, query_result_with_server_messages, - requires_simple_query_batch, sqlserver_batch_can_use_execute, sqlserver_cell_to_json, sqlserver_columns_sql, - sqlserver_completion_assistant_sql, sqlserver_dml_output_returns_rows, sqlserver_hidden_schema_names, - sqlserver_indexes_sql, sqlserver_list_objects_sql, sqlserver_list_schemas_sql, sqlserver_list_tables_sql, + is_blocking_sqlserver_unsafe_probe_error, is_sqlserver_spatial_column, is_sqlserver_variant_column, + query_result_with_server_messages, requires_simple_query_batch, sqlserver_batch_can_use_execute, + sqlserver_cell_to_json, sqlserver_columns_sql, sqlserver_completion_assistant_sql, + sqlserver_dml_output_returns_rows, sqlserver_hidden_schema_names, sqlserver_indexes_sql, + sqlserver_list_objects_sql, sqlserver_list_schemas_sql, sqlserver_list_tables_sql, sqlserver_schema_name_predicate, sqlserver_table_comment_sql, sqlserver_visible_object_predicate, strip_dbx_sqlserver_row_number_column, SqlServerDescribedColumn, SqlServerResultSet, + SQLSERVER_RESULT_TYPE_PROBE_SQL, }; use crate::types::{ CompletionAssistantMatchMode, CompletionAssistantObjectKind, CompletionAssistantRequest, QueryResult, @@ -3014,6 +3111,24 @@ mod tests { })); } + #[test] + fn sqlserver_result_type_probe_keeps_modern_and_legacy_paths_in_one_round_trip() { + assert!(SQLSERVER_RESULT_TYPE_PROBE_SQL.contains("sys.dm_exec_describe_first_result_set")); + assert!(SQLSERVER_RESULT_TYPE_PROBE_SQL.contains("##dbx_result_type_probe_")); + assert!(SQLSERVER_RESULT_TYPE_PROBE_SQL.contains("SELECT TOP (0) * INTO")); + assert!(SQLSERVER_RESULT_TYPE_PROBE_SQL.contains("FROM tempdb.sys.columns")); + assert!(!SQLSERVER_RESULT_TYPE_PROBE_SQL.contains("FMTONLY")); + assert_eq!(SQLSERVER_RESULT_TYPE_PROBE_SQL.matches("SELECT @dbx_use_describe_dmv").count(), 1); + } + + #[test] + fn sqlserver_legacy_probe_errors_are_blocking_only_when_marked_unsafe() { + assert!(is_blocking_sqlserver_unsafe_probe_error( + "SQL Server unsafe result type: legacy schema capture failed" + )); + assert!(!is_blocking_sqlserver_unsafe_probe_error("Invalid object name")); + } + #[test] fn sqlserver_wraps_sql_variant_columns_as_nvarchar() { let rewritten = build_sqlserver_unsafe_type_query( @@ -3090,4 +3205,66 @@ mod tests { assert!(rewritten.contains("AS NVARCHAR(MAX))")); assert!(rewritten.contains("FROM dbo.t")); } + + #[tokio::test] + #[ignore = "requires DBX_LIVE_SQLSERVER_HOST/PORT/USER/PASSWORD"] + async fn live_sqlserver_legacy_probe_casts_variant_and_keeps_connection_usable() { + let host = std::env::var("DBX_LIVE_SQLSERVER_HOST").expect("DBX_LIVE_SQLSERVER_HOST"); + let port = std::env::var("DBX_LIVE_SQLSERVER_PORT") + .expect("DBX_LIVE_SQLSERVER_PORT") + .parse::() + .expect("valid DBX_LIVE_SQLSERVER_PORT"); + let user = std::env::var("DBX_LIVE_SQLSERVER_USER").expect("DBX_LIVE_SQLSERVER_USER"); + let password = std::env::var("DBX_LIVE_SQLSERVER_PASSWORD").expect("DBX_LIVE_SQLSERVER_PASSWORD"); + let mut client = + super::connect(&host, port, &user, &password, Some("tempdb"), None, std::time::Duration::from_secs(10)) + .await + .unwrap(); + + let setup = "\ + IF OBJECT_ID('tempdb..#dbx_issue_4002') IS NOT NULL DROP TABLE #dbx_issue_4002; \ + CREATE TABLE #dbx_issue_4002 (id int NOT NULL, payload sql_variant NULL); \ + INSERT INTO #dbx_issue_4002 (id, payload) VALUES (1, CAST(N'legacy' AS nvarchar(20)))"; + client.simple_query(setup).await.unwrap().into_results().await.unwrap(); + + let sql = "SELECT id, payload FROM #dbx_issue_4002"; + let ordinary_columns = super::describe_sqlserver_result_set_with_mode( + &mut client, + "SELECT 42 AS answer", + "SELECT 42 AS answer", + true, + ) + .await + .unwrap(); + assert_eq!(ordinary_columns[0].system_type_name.as_deref(), Some("int")); + + let legacy_columns = super::describe_sqlserver_result_set_with_mode(&mut client, sql, sql, true).await.unwrap(); + assert_eq!(legacy_columns.len(), 2); + assert!(is_sqlserver_variant_column(&legacy_columns[1])); + + let rewritten = build_sqlserver_unsafe_type_query(sql, &legacy_columns).unwrap(); + let legacy_rows = client.query(rewritten, &[]).await.unwrap().into_first_result().await.unwrap(); + assert_eq!(legacy_rows[0].get::(0), Some(1)); + assert_eq!(legacy_rows[0].get::<&str, _>(1), Some("legacy")); + + let ordinary = super::execute_query(&mut client, "SELECT CAST(42 AS int) AS answer").await.unwrap(); + assert_eq!(ordinary.rows, vec![vec![serde_json::json!(42)]]); + 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 boundary_error = super::describe_sqlserver_result_set_with_mode( + &mut client, + "SELECT 1 AS duplicate_name, 2 AS duplicate_name", + "SELECT 1 AS duplicate_name, 2 AS duplicate_name", + true, + ) + .await + .unwrap_err(); + assert!(is_blocking_sqlserver_unsafe_probe_error(&boundary_error)); + + let continued = super::execute_query(&mut client, "SELECT CAST(7 AS int) AS still_connected").await.unwrap(); + assert_eq!(continued.rows, vec![vec![serde_json::json!(7)]]); + + client.simple_query("DROP TABLE #dbx_issue_4002").await.unwrap().into_results().await.unwrap(); + } }