From 579008ffbdb2e69aef283f581f31dfc4fd6b5c48 Mon Sep 17 00:00:00 2001 From: zipg Date: Sun, 13 Sep 2026 01:51:42 +0800 Subject: [PATCH] fix(query): schema-qualify unqualified names for SQL Server --- crates/dbx-core/src/query.rs | 167 ++++++++++++++++++++++++++++++++--- 1 file changed, 153 insertions(+), 14 deletions(-) diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index ba085371e7..54cc8fe48f 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -4,7 +4,7 @@ use serde::{Deserialize, Serialize}; use sqlparser::ast::{ visit_relations_mut, Ident, ObjectName, ObjectNamePart, ObjectType, Statement, TableFactor, VisitMut, VisitorMut, }; -use sqlparser::dialect::{GenericDialect, PostgreSqlDialect}; +use sqlparser::dialect::{GenericDialect, MsSqlDialect, PostgreSqlDialect}; use sqlparser::parser::Parser; use std::collections::{HashMap, HashSet}; use std::future::Future; @@ -857,6 +857,9 @@ fn sql_for_execution_context_with_identifier_quote( }; match db_type { Some(DatabaseType::Iris) => qualify_iris_unqualified_dml(sql, schema).unwrap_or_else(|| sql.to_string()), + Some(DatabaseType::SqlServer) => { + qualify_sqlserver_unqualified_dml(sql, schema).unwrap_or_else(|| sql.to_string()) + } Some(DatabaseType::Kingbase) => { qualify_kingbase_unqualified_relations(sql, schema, identifier_quote).unwrap_or_else(|| sql.to_string()) } @@ -877,8 +880,9 @@ fn qualify_iris_unqualified_dml(sql: &str, schema: &str) -> Option { continue; } let cte_names = statement_cte_names(statement); + let table_aliases = statement_table_aliases(statement); let _ = visit_relations_mut(statement, |name| { - if qualify_unqualified_relation_name(name, schema, &cte_names) { + if qualify_unqualified_relation_name(name, schema, &cte_names, &table_aliases) { changed = true; } ControlFlow::<()>::Continue(()) @@ -888,6 +892,35 @@ fn qualify_iris_unqualified_dml(sql: &str, schema: &str) -> Option { changed.then(|| statements.iter().map(ToString::to_string).collect::>().join("; ")) } +fn qualify_sqlserver_unqualified_dml(sql: &str, schema: &str) -> Option { + let dialect = MsSqlDialect {}; + let mut statements = Parser::parse_sql(&dialect, sql).ok()?; + if statements.is_empty() { + return None; + } + + let mut changed = false; + for statement in &mut statements { + if !statement_uses_schema_context(statement) { + continue; + } + let cte_names = statement_cte_names(statement); + let table_aliases = statement_table_aliases(statement); + let mut qualifier = SchemaRelationQualifier { + schema, + identifier_quote: '[', + cte_names: &cte_names, + table_aliases: &table_aliases, + parameterized_table_depth: 0, + changed: false, + }; + let _ = statement.visit(&mut qualifier); + changed |= qualifier.changed; + } + + changed.then(|| statements.iter().map(ToString::to_string).collect::>().join("; ")) +} + fn qualify_kingbase_unqualified_relations(sql: &str, schema: &str, identifier_quote: Option<&str>) -> Option { let dialect = PostgreSqlDialect {}; let mut statements = Parser::parse_sql(&dialect, sql).ok()?; @@ -901,10 +934,12 @@ fn qualify_kingbase_unqualified_relations(sql: &str, schema: &str, identifier_qu continue; } let cte_names = statement_cte_names(statement); - let mut qualifier = KingbaseRelationQualifier { + let table_aliases = statement_table_aliases(statement); + let mut qualifier = SchemaRelationQualifier { schema, identifier_quote: identifier_quote_char(identifier_quote), cte_names: &cte_names, + table_aliases: &table_aliases, parameterized_table_depth: 0, changed: false, }; @@ -915,15 +950,16 @@ fn qualify_kingbase_unqualified_relations(sql: &str, schema: &str, identifier_qu changed.then(|| statements.iter().map(ToString::to_string).collect::>().join("; ")) } -struct KingbaseRelationQualifier<'a> { +struct SchemaRelationQualifier<'a> { schema: &'a str, identifier_quote: char, cte_names: &'a HashSet, + table_aliases: &'a HashSet, parameterized_table_depth: usize, changed: bool, } -impl VisitorMut for KingbaseRelationQualifier<'_> { +impl VisitorMut for SchemaRelationQualifier<'_> { type Break = (); fn pre_visit_table_factor(&mut self, table_factor: &mut TableFactor) -> ControlFlow { @@ -946,6 +982,7 @@ impl VisitorMut for KingbaseRelationQualifier<'_> { relation, self.schema, self.cte_names, + self.table_aliases, self.identifier_quote, ) { @@ -966,20 +1003,30 @@ fn statement_uses_schema_context(statement: &Statement) -> bool { ) } -fn qualify_unqualified_relation_name(name: &mut ObjectName, schema: &str, cte_names: &HashSet) -> bool { - qualify_unqualified_relation_name_with_quote(name, schema, cte_names, '"') +fn qualify_unqualified_relation_name( + name: &mut ObjectName, + schema: &str, + cte_names: &HashSet, + table_aliases: &HashSet, +) -> bool { + qualify_unqualified_relation_name_with_quote(name, schema, cte_names, table_aliases, '"') } fn qualify_unqualified_relation_name_with_quote( name: &mut ObjectName, schema: &str, cte_names: &HashSet, + table_aliases: &HashSet, identifier_quote: char, ) -> bool { let [ObjectNamePart::Identifier(table)] = name.0.as_slice() else { return false; }; - if cte_names.contains(&table.value.to_ascii_uppercase()) { + let upper_name = table.value.to_ascii_uppercase(); + if cte_names.contains(&upper_name) || table_aliases.contains(&upper_name) { + return false; + } + if table.value.starts_with('@') || table.value.starts_with('#') { return false; } @@ -1026,8 +1073,32 @@ fn collect_query_cte_names(query: &sqlparser::ast::Query, names: &mut HashSet { + names: &'a mut HashSet, +} + +impl VisitorMut for TableAliasCollector<'_> { + type Break = (); + + fn post_visit_table_factor(&mut self, table_factor: &mut TableFactor) -> ControlFlow { + if let TableFactor::Table { alias: Some(alias), .. } = table_factor { + self.names.insert(alias.name.value.to_ascii_uppercase()); + } + ControlFlow::Continue(()) + } +} + +/// FROM-clause aliases resolve to their table, so a single-part relation that +/// matches one must never be schema-qualified (e.g. `UPDATE p ... FROM products p`). +fn statement_table_aliases(statement: &mut Statement) -> HashSet { + let mut names = HashSet::new(); + let mut collector = TableAliasCollector { names: &mut names }; + let _ = statement.visit(&mut collector); + names +} + fn qualifies_unqualified_agent_relations(db_type: Option) -> bool { - matches!(db_type, Some(DatabaseType::Iris | DatabaseType::Kingbase)) + matches!(db_type, Some(DatabaseType::Iris | DatabaseType::Kingbase | DatabaseType::SqlServer)) } #[derive(Clone, Copy, Debug, Default, Deserialize, PartialEq, Eq)] @@ -2006,6 +2077,7 @@ async fn do_execute_typed( let client = client.clone(); let max_rows = options.max_rows; let execution_mode = options.execution_mode; + let sql = sql_for_execution_context_with_identifier_quote(pool_db_type, sql, schema, Some("[")); let (mut client, lock_wait_ms) = match lock_shared_client_with_wait(&client, cancel_token.clone(), None).await { Ok(value) => value, @@ -2014,10 +2086,10 @@ async fn do_execute_typed( let execution = async { if execution_mode == QueryExecutionMode::Simple { let mut results = - db::sqlserver::execute_simple_batch_with_max_rows(&mut client, sql, max_rows).await?; + db::sqlserver::execute_simple_batch_with_max_rows(&mut client, &sql, max_rows).await?; Ok(results.remove(0)) } else { - db::sqlserver::execute_query_with_max_rows(&mut client, sql, max_rows).await + db::sqlserver::execute_query_with_max_rows(&mut client, &sql, max_rows).await } }; let result = wait_for_query_opt(cancel_token, query_timeout, execution) @@ -3012,7 +3084,7 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed( } if is_sqlserver { - return execute_multi_sqlserver(state, &pool_key, sql, cancel_token, options).await.map_err(Into::into); + return execute_multi_sqlserver(state, &pool_key, sql, schema, cancel_token, options).await.map_err(Into::into); } let is_http_sqlite = { @@ -3651,6 +3723,7 @@ async fn execute_multi_sqlserver( state: &AppState, pool_key: &str, sql: &str, + schema: Option<&str>, cancel_token: Option, options: QueryExecutionOptions, ) -> Result, String> { @@ -3704,11 +3777,14 @@ async fn execute_multi_sqlserver( break; } + let execution_sql = + sql_for_execution_context_with_identifier_quote(Some(DatabaseType::SqlServer), batch, schema, Some("[")); let execution = async { if execution_mode == QueryExecutionMode::Simple { - db::sqlserver::execute_simple_batch_with_max_rows_metadata(&mut client_guard, batch, max_rows).await + db::sqlserver::execute_simple_batch_with_max_rows_metadata(&mut client_guard, &execution_sql, max_rows) + .await } else { - db::sqlserver::execute_batch_with_max_rows_metadata(&mut client_guard, batch, max_rows).await + db::sqlserver::execute_batch_with_max_rows_metadata(&mut client_guard, &execution_sql, max_rows).await } }; let result = wait_for_result_opt(cancel_token.clone(), query_timeout, execution).await; @@ -9469,6 +9545,69 @@ for line in sys.stdin: ); } + #[test] + fn sqlserver_execution_context_qualifies_unqualified_dml_tables() { + assert_eq!( + sql_for_execution_context(Some(DatabaseType::SqlServer), "SELECT TOP 1 * FROM products", Some("core")), + "SELECT TOP 1 * FROM [core].products" + ); + + let qualified_join = sql_for_execution_context( + Some(DatabaseType::SqlServer), + "SELECT p.id FROM products p JOIN customers c ON c.id = p.customer_id", + Some("sales"), + ); + assert!(qualified_join.contains("FROM [sales].products p"), "{qualified_join}"); + assert!(qualified_join.contains("JOIN [sales].customers c"), "{qualified_join}"); + + assert_eq!( + sql_for_execution_context( + Some(DatabaseType::SqlServer), + "UPDATE products SET status = 'done' WHERE id IN (SELECT product_id FROM audit_products)", + Some("core") + ), + "UPDATE [core].products SET status = 'done' WHERE id IN (SELECT product_id FROM [core].audit_products)" + ); + } + + #[test] + fn sqlserver_execution_context_skips_update_target_aliases() { + let qualified = sql_for_execution_context( + Some(DatabaseType::SqlServer), + "UPDATE p SET p.status = 'done' FROM products p JOIN customers c ON c.id = p.customer_id", + Some("core"), + ); + assert!(qualified.starts_with("UPDATE p SET"), "{qualified}"); + assert!(qualified.contains("FROM [core].products p"), "{qualified}"); + assert!(qualified.contains("JOIN [core].customers c ON c.id = p.customer_id"), "{qualified}"); + assert!(!qualified.contains("UPDATE [core].p"), "{qualified}"); + assert!(!qualified.contains("[core].c ON"), "{qualified}"); + } + + #[test] + fn sqlserver_execution_context_preserves_ctes_qualified_tables_and_temp_tables() { + assert_eq!( + sql_for_execution_context( + Some(DatabaseType::SqlServer), + "WITH recent AS (SELECT * FROM products) SELECT * FROM recent", + Some("core") + ), + "WITH recent AS (SELECT * FROM [core].products) SELECT * FROM recent" + ); + assert_eq!( + sql_for_execution_context( + Some(DatabaseType::SqlServer), + "SELECT * FROM [archive].products UNION ALL SELECT * FROM #staged_products", + Some("core") + ), + "SELECT * FROM [archive].products UNION ALL SELECT * FROM #staged_products" + ); + assert_eq!( + sql_for_execution_context(Some(DatabaseType::SqlServer), "CREATE TABLE products (id INT)", Some("core")), + "CREATE TABLE products (id INT)" + ); + } + #[test] fn kingbase_execution_context_qualifies_only_unqualified_relations() { assert_eq!(