From 6f1bf657263afc79dc8d772d2cb9be207aab2de0 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Thu, 6 Aug 2026 17:33:58 +0800 Subject: [PATCH] fix(sqlserver): support multi-column ALTER TABLE ADD Closes #5484 --- crates/dbx-core/src/sql_analysis.rs | 125 +++++++++++++++++++++++++- crates/dbx-core/tests/sql_analysis.rs | 32 +++++++ 2 files changed, 153 insertions(+), 4 deletions(-) diff --git a/crates/dbx-core/src/sql_analysis.rs b/crates/dbx-core/src/sql_analysis.rs index e44db6655..a92baaa71 100644 --- a/crates/dbx-core/src/sql_analysis.rs +++ b/crates/dbx-core/src/sql_analysis.rs @@ -127,15 +127,125 @@ fn parse_sqlserver(sql: &str) -> Result, ParserError> { Err(error) => { let mut fallback_tokens = Tokenizer::new(&dialect, sql).tokenize_with_location()?; normalize_sqlserver_create_proc_tokens(&mut fallback_tokens); - if remove_sqlserver_query_hint_tokens(&mut fallback_tokens) { - Parser::new(&dialect).with_tokens_with_locations(fallback_tokens).parse_statements() - } else { - Err(error) + let mut changed = normalize_sqlserver_alter_table_add_tokens(&mut fallback_tokens); + changed |= remove_sqlserver_query_hint_tokens(&mut fallback_tokens); + if changed { + if let Ok(statements) = + Parser::new(&dialect).with_tokens_with_locations(fallback_tokens).parse_statements() + { + return Ok(statements); + } } + Err(error) } } } +fn normalize_sqlserver_alter_table_add_tokens(tokens: &mut Vec) -> bool { + let mut insertions = Vec::new(); + let mut index = 0usize; + + while index < tokens.len() { + if unquoted_token_keyword(&tokens[index]) != Some(Keyword::ALTER) { + index += 1; + continue; + } + let Some(table_index) = next_significant_token_index(tokens, index + 1) else { + break; + }; + if unquoted_token_keyword(&tokens[table_index]) != Some(Keyword::TABLE) { + index += 1; + continue; + } + + let statement_end = sqlserver_statement_end_index(tokens, table_index + 1); + let Some(add_index) = sqlserver_alter_table_add_index(tokens, table_index + 1, statement_end) else { + index = statement_end.saturating_add(1); + continue; + }; + let add_token = tokens[add_index].clone(); + let mut depth = 0usize; + for token_index in add_index + 1..statement_end { + match tokens[token_index].token { + Token::LParen => depth += 1, + Token::RParen => depth = depth.saturating_sub(1), + Token::Comma if depth == 0 => { + let Some(next_index) = next_significant_token_index(tokens, token_index + 1) else { + continue; + }; + if next_index >= statement_end { + continue; + } + if is_sqlserver_alter_table_operation_starter(&tokens[next_index]) { + break; + } + insertions.push((next_index, add_token.clone())); + } + _ => {} + } + } + index = statement_end.saturating_add(1); + } + + let changed = !insertions.is_empty(); + for (index, token) in insertions.into_iter().rev() { + tokens.insert(index, token); + } + changed +} + +fn sqlserver_statement_end_index(tokens: &[TokenWithSpan], start: usize) -> usize { + let mut depth = 0usize; + for (index, token) in tokens.iter().enumerate().skip(start) { + match token.token { + Token::LParen => depth += 1, + Token::RParen => depth = depth.saturating_sub(1), + Token::SemiColon | Token::EOF if depth == 0 => return index, + _ => {} + } + } + tokens.len() +} + +fn sqlserver_alter_table_add_index(tokens: &[TokenWithSpan], start: usize, end: usize) -> Option { + let mut depth = 0usize; + for (index, token) in tokens.iter().enumerate().take(end).skip(start) { + match token.token { + Token::LParen => depth += 1, + Token::RParen => depth = depth.saturating_sub(1), + _ if depth == 0 && unquoted_token_keyword(token) == Some(Keyword::ADD) => return Some(index), + _ => {} + } + } + None +} + +fn is_sqlserver_alter_table_operation_starter(token: &TokenWithSpan) -> bool { + let Token::Word(word) = &token.token else { + return false; + }; + if word.quote_style.is_some() { + return false; + } + + matches!( + word.value.to_ascii_uppercase().as_str(), + "ADD" + | "ALTER" + | "DISABLE" + | "DROP" + | "ENABLE" + | "NOCHECK" + | "PARTITION" + | "REBUILD" + | "RENAME" + | "REPLICA" + | "SET" + | "SWAP" + | "SWITCH" + ) +} + fn remove_sqlserver_query_hint_tokens(tokens: &mut Vec) -> bool { let mut depth = 0usize; let mut ranges = Vec::new(); @@ -316,6 +426,13 @@ fn token_keyword(token: &TokenWithSpan) -> Option { } } +fn unquoted_token_keyword(token: &TokenWithSpan) -> Option { + match &token.token { + Token::Word(word) if word.quote_style.is_none() => Some(word.keyword), + _ => None, + } +} + fn starts_with_duckdb_parser_gap_sql(sql: &str) -> bool { starts_with_duckdb_result_sql_keyword(sql) && starts_with_executable_sql_keyword(sql, &["FROM", "SUMMARIZE", "SUMMARISE", "PIVOT", "UNPIVOT"]) diff --git a/crates/dbx-core/tests/sql_analysis.rs b/crates/dbx-core/tests/sql_analysis.rs index 286d99bb1..e5adda136 100644 --- a/crates/dbx-core/tests/sql_analysis.rs +++ b/crates/dbx-core/tests/sql_analysis.rs @@ -410,6 +410,38 @@ fn sqlserver_proc_identifiers_remain_identifiers_outside_create() { assert_eq!(analysis.columns[0].name, "proc"); } +#[test] +fn sqlserver_alter_table_single_add_supports_multiple_columns() { + for sql in [ + "ALTER TABLE dbo.demo\nADD isOldWell BIT NULL,\n isNewWell BIT NULL;", + "ALTER TABLE [dbo].[demo] ADD amount DECIMAL(10, 2) DEFAULT (0), [display_name] NVARCHAR(50) NULL;", + "ALTER TABLE dbo.demo ADD enabled BIT NULL, CHECK (enabled IN (0, 1));", + ] { + analyze_sql_references(sql, Some("sqlserver")) + .unwrap_or_else(|error| panic!("SQL Server single-ADD multi-column ALTER TABLE should analyze: {error}")); + } +} + +#[test] +fn sqlserver_alter_table_add_normalization_preserves_boundaries() { + analyze_sql_references("ALTER TABLE dbo.demo ADD isOldWell BIT NULL, ADD isNewWell BIT NULL;", Some("sqlserver")) + .expect("existing repeated-ADD parser behavior should remain valid"); + + let missing_comma = + analyze_sql_references("ALTER TABLE dbo.demo ADD isOldWell BIT NULL isNewWell BIT NULL;", Some("sqlserver")) + .expect_err("missing column separators must remain invalid"); + assert!(missing_comma.contains("isNewWell")); + + analyze_sql_references( + "ALTER TABLE dbo.demo ADD amount DECIMAL(10, 2) NULL, label NVARCHAR(20) NULL; SELECT label FROM dbo.demo;", + Some("sqlserver"), + ) + .expect("data-type commas and multiple statements should remain parseable"); + + analyze_sql_references("ALTER TABLE dbo.demo ADD first_flag BIT NULL, second_flag BIT NULL;", Some("postgres")) + .expect_err("other dialects must not inherit SQL Server ALTER TABLE normalization"); +} + #[test] fn sqlserver_query_hints_do_not_raise_parser_errors() { let analysis = analyze_sql_references(