github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-19 16:16:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(query): schema-qualify unqualified names for SQL Server

Browse Source
This commit is contained in:
zipg
2026-09-13 01:51:42 +08:00
committed by GitHub
parent 7472f16066
commit 579008ffbd
1 changed files with 153 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/query.rs
+153 -14
View File
@@ -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<String> {
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<String> {
changed.then(|| statements.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))
}
fn qualify_sqlserver_unqualified_dml(sql: &str, schema: &str) -> Option<String> {
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::<Vec<_>>().join("; "))
}
fn qualify_kingbase_unqualified_relations(sql: &str, schema: &str, identifier_quote: Option<&str>) -> Option<String> {
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::<Vec<_>>().join("; "))
}
struct KingbaseRelationQualifier<'a> {
struct SchemaRelationQualifier<'a> {
schema: &'a str,
identifier_quote: char,
cte_names: &'a HashSet<String>,
table_aliases: &'a HashSet<String>,
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<Self::Break> {
@@ -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<String>) -> bool {
qualify_unqualified_relation_name_with_quote(name, schema, cte_names, '"')
fn qualify_unqualified_relation_name(
name: &mut ObjectName,
schema: &str,
cte_names: &HashSet<String>,
table_aliases: &HashSet<String>,
) -> 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<String>,
table_aliases: &HashSet<String>,
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<St
}
}
struct TableAliasCollector<'a> {
names: &'a mut HashSet<String>,
}
impl VisitorMut for TableAliasCollector<'_> {
type Break = ();
fn post_visit_table_factor(&mut self, table_factor: &mut TableFactor) -> ControlFlow<Self::Break> {
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<String> {
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<DatabaseType>) -> 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<CancellationToken>,
options: QueryExecutionOptions,
) -> Result<Vec<ExecuteMultiResult>, 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!(