From 91abfb747b0329b575f6e1ef566a9f7ab2a6981b Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Sun, 19 Jul 2026 11:11:37 +0800
Subject: [PATCH] fix(sqlserver): recognize CREATE PROC syntax

---
 crates/dbx-core/src/sql_analysis.rs   | 58 +++++++++++++++++++++++++--
 crates/dbx-core/tests/sql_analysis.rs | 32 +++++++++++++++
 2 files changed, 87 insertions(+), 3 deletions(-)

diff --git a/crates/dbx-core/src/sql_analysis.rs b/crates/dbx-core/src/sql_analysis.rs
index 4a1067ae0..9a853229b 100644
--- a/crates/dbx-core/src/sql_analysis.rs
+++ b/crates/dbx-core/src/sql_analysis.rs
@@ -10,8 +10,9 @@ use sqlparser::ast::{
 use sqlparser::dialect::{
     ClickHouseDialect, DuckDbDialect, GenericDialect, MsSqlDialect, MySqlDialect, PostgreSqlDialect, SQLiteDialect,
 };
-use sqlparser::parser::Parser;
-use sqlparser::tokenizer::Span;
+use sqlparser::keywords::Keyword;
+use sqlparser::parser::{Parser, ParserError};
+use sqlparser::tokenizer::{Span, Token, TokenWithSpan, Tokenizer};
 
 use crate::sql::{starts_with_duckdb_result_sql_keyword, starts_with_executable_sql_keyword};
 
@@ -99,7 +100,7 @@ pub fn analyze_sql_references(sql: &str, dialect: Option<&str>) -> Result<SqlRef
         "postgres" => Parser::parse_sql(&PostgreSqlDialect {}, &parser_sql),
         "mysql" => Parser::parse_sql(&MySqlDialect {}, &parser_sql),
         "sqlite" => Parser::parse_sql(&SQLiteDialect {}, &parser_sql),
-        "sqlserver" => Parser::parse_sql(&MsSqlDialect {}, &parser_sql),
+        "sqlserver" => parse_sqlserver(&parser_sql),
         "clickhouse" => Parser::parse_sql(&ClickHouseDialect {}, &parser_sql),
         "duckdb" => Parser::parse_sql(&DuckDbDialect {}, &parser_sql),
         _ => Parser::parse_sql(&GenericDialect {}, &parser_sql),
@@ -114,6 +115,57 @@ pub fn analyze_sql_references(sql: &str, dialect: Option<&str>) -> Result<SqlRef
     Ok(SqlReferenceAnalysis { tables: analyzer.tables, columns: analyzer.columns, scopes: analyzer.scopes })
 }
 
+fn parse_sqlserver(sql: &str) -> Result<Vec<Statement>, ParserError> {
+    let dialect = MsSqlDialect {};
+    let mut tokens = Tokenizer::new(&dialect, sql).tokenize_with_location()?;
+    normalize_sqlserver_create_proc_tokens(&mut tokens);
+    Parser::new(&dialect).with_tokens_with_locations(tokens).parse_statements()
+}
+
+fn normalize_sqlserver_create_proc_tokens(tokens: &mut [TokenWithSpan]) {
+    let significant_indexes: Vec<usize> = tokens
+        .iter()
+        .enumerate()
+        .filter_map(|(index, token)| (!matches!(token.token, Token::Whitespace(_))).then_some(index))
+        .collect();
+
+    for (position, index) in significant_indexes.iter().copied().enumerate() {
+        if token_keyword(&tokens[index]) != Some(Keyword::CREATE) {
+            continue;
+        }
+
+        let mut proc_position = position + 1;
+        if significant_indexes.get(proc_position).and_then(|index| token_keyword(&tokens[*index])) == Some(Keyword::OR)
+        {
+            proc_position += 1;
+            if significant_indexes.get(proc_position).and_then(|index| token_keyword(&tokens[*index]))
+                != Some(Keyword::ALTER)
+            {
+                continue;
+            }
+            proc_position += 1;
+        }
+
+        let Some(proc_index) = significant_indexes.get(proc_position).copied() else {
+            continue;
+        };
+        let Token::Word(word) = &mut tokens[proc_index].token else {
+            continue;
+        };
+        // SQL Server documents PROC as a contextual synonym for PROCEDURE after CREATE.
+        if word.quote_style.is_none() && word.value.eq_ignore_ascii_case("proc") {
+            word.keyword = Keyword::PROCEDURE;
+        }
+    }
+}
+
+fn token_keyword(token: &TokenWithSpan) -> Option<Keyword> {
+    match &token.token {
+        Token::Word(word) => 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 f32acfd94..f4912e79d 100644
--- a/crates/dbx-core/tests/sql_analysis.rs
+++ b/crates/dbx-core/tests/sql_analysis.rs
@@ -91,6 +91,38 @@ fn extracts_unqualified_order_by_columns_for_sqlserver_queries() {
     assert_eq!(columns, vec![(None, "PDReceiveDatePartInfo")]);
 }
 
+#[test]
+fn sqlserver_create_proc_and_procedure_are_equivalent() {
+    for sql in ["CREATE PROC test\nAS\n", "CREATE PROCEDURE test\nAS\n", "CREATE PROC test AS SELECT 1;"] {
+        let analysis = analyze_sql_references(sql, Some("sqlserver"))
+            .unwrap_or_else(|error| panic!("SQL Server procedure declaration should analyze: {error}"));
+        assert!(analysis.tables.is_empty());
+        assert!(analysis.columns.is_empty());
+    }
+}
+
+#[test]
+fn sqlserver_create_or_alter_proc_is_supported() {
+    analyze_sql_references("CREATE OR ALTER PROC test AS SELECT 1;", Some("sqlserver"))
+        .unwrap_or_else(|error| panic!("SQL Server CREATE OR ALTER PROC should analyze: {error}"));
+}
+
+#[test]
+fn create_proc_remains_invalid_outside_sqlserver() {
+    let error = analyze_sql_references("CREATE PROC test AS SELECT 1", Some("postgres"))
+        .expect_err("PostgreSQL must not inherit SQL Server's PROC synonym");
+
+    assert!(error.contains("an object type after CREATE"));
+}
+
+#[test]
+fn sqlserver_proc_identifiers_remain_identifiers_outside_create() {
+    let analysis = analyze_sql_references("SELECT proc FROM jobs", Some("sqlserver")).unwrap();
+
+    assert_eq!(analysis.tables[0].name, "jobs");
+    assert_eq!(analysis.columns[0].name, "proc");
+}
+
 #[test]
 fn duckdb_parser_gap_queries_do_not_raise_syntax_errors() {
     for sql in ["FROM users;", "SUMMARIZE users;", "SUMMARISE users;"] {