fix: keep authored SQL tokens when formatting

sqb format regenerated SQL from the parse tree, so any spelling the SQL
library canonicalised was written back into the user's file. In Snowflake
that renamed STARTSWITH to STARTS_WITH, which Snowflake rejects, and also
rewrote SUBSTR, TIMESTAMPDIFF and TRY_TO_DECIMAL to synonyms. The formatter's
syntax-tree check could not see this because both spellings parse to the same
node.

After generating the layout, walk the authored and generated token streams
together and write the authored text of every aligned token into the new
layout. Only whitespace, the case of unquoted words, an added AS, and dropped
optional terminators may differ; any other structural difference refuses the
format and leaves the file unchanged. The same check runs on the idempotence
pass.

On a 1,337-file Snowflake project this removes all 82 function renames, with
no new refusals.

Author: kvlonge

crates/sqlbuild-rules-native/src/sql_lint/_helpers/formatter.rs

@@ -14,6 +14,8 @@ use crate::sql_lint::models::{FormatRequest, FormatResponse};
 
 const COMMENT_ATTACHMENT_FAILURE: &str =
     "native formatter could not preserve comment token attachments";
+const TOKEN_PRESERVATION_FAILURE: &str =
+    "native formatter would change authored SQL tokens, not only layout";
 const UNSUPPORTED_SQL_FAILURE: &str =
     "native formatter could not safely restore literal or cast spelling";
 
@@ -80,7 +82,11 @@ fn try_format_sql(request: FormatRequest) -> Result<FormatResponse, String> {
         original_tokens: &semantic_tokens,
         comments: &comments,
     };
-    let formatted = format_once(&neutral, &formatted_context)?;
+    let formatted = preserve_authored_tokens(
+        &neutral,
+        format_once(&neutral, &formatted_context)?,
+        &dialect,
+    )?;
     let formatted_tokens = dialect
         .tokenize(&formatted)
         .map_err(|error| error.to_string())?;
@@ -103,14 +109,105 @@ fn try_format_sql(request: FormatRequest) -> Result<FormatResponse, String> {
         original_tokens: &second_semantic_tokens,
         comments: &second_comments,
     };
-    let second_pass = format_once(&second_neutral, &second_context)?;
+    let second_pass = preserve_authored_tokens(
+        &second_neutral,
+        format_once(&second_neutral, &second_context)?,
+        &dialect,
+    )?;
     if second_pass != formatted {
         return Err("native formatter output is not idempotent".to_string());
     }
     let changed = formatted != original;
     response(formatted, changed, true, None)
 }
 
+/// Restores authored token text into generated layout, refusing any structural token change.
+pub(crate) fn preserve_authored_tokens(
+    authored_sql: &str,
+    mut formatted: String,
+    dialect: &Dialect,
+) -> Result<String, String> {
+    let authored = dialect
+        .tokenize(authored_sql)
+        .map_err(|error| error.to_string())?;
+    let generated = dialect
+        .tokenize(&formatted)
+        .map_err(|error| error.to_string())?;
+    let mut replacements: Vec<(usize, usize, String)> = Vec::new();
+    let mut authored_index = 0;
+    let mut generated_index = 0;
+    while authored_index < authored.len() || generated_index < generated.len() {
+        let before = authored.get(authored_index);
+        let after = generated.get(generated_index);
+        match (before, after) {
+            (before, Some(after))
+                if after.token_type == TokenType::As
+                    && before.is_none_or(|token| token.token_type != TokenType::As) =>
+            {
+                generated_index += 1;
+            }
+            (Some(before), Some(after)) if tokens_align(before, after) => {
+                let authored_text = raw_token_text(authored_sql, before)?;
+                let generated_text = raw_token_text(&formatted, after)?;
+                if authored_text != generated_text
+                    && !(is_unquoted_word(&authored_text)
+                        && authored_text.eq_ignore_ascii_case(&generated_text))
+                {
+                    replacements.push((after.span.start, after.span.end, authored_text));
+                }
+                authored_index += 1;
+                generated_index += 1;
+            }
+            (Some(before), after) if is_optional_terminator(&authored, authored_index, before) => {
+                if after.is_some_and(|token| token.token_type == before.token_type) {
+                    generated_index += 1;
+                }
+                authored_index += 1;
+            }
+            (None, Some(after)) if after.token_type == TokenType::Semicolon => {
+                generated_index += 1;
+            }
+            _ => return Err(TOKEN_PRESERVATION_FAILURE.to_string()),
+        }
+    }
+    for (start, end, text) in replacements.into_iter().rev() {
+        let start = char_to_byte(&formatted, start)?;
+        let end = char_to_byte(&formatted, end)?;
+        formatted.replace_range(start..end, &text);
+    }
+    Ok(formatted)
+}
+
+fn tokens_align(before: &Token, after: &Token) -> bool {
+    before.token_type == after.token_type || (is_word_token(before) && is_word_token(after))
+}
+
+fn is_word_token(token: &Token) -> bool {
+    is_unquoted_word(&token.text)
+}
+
+fn is_unquoted_word(text: &str) -> bool {
+    let mut characters = text.chars();
+    characters
+        .next()
+        .is_some_and(|first| first.is_ascii_alphabetic() || first == '_')
+        && characters
+            .all(|character| character.is_ascii_alphanumeric() || matches!(character, '_' | '$'))
+}
+
+fn is_optional_terminator(tokens: &[Token], index: usize, token: &Token) -> bool {
+    token.token_type == TokenType::Semicolon
+        || (token.token_type == TokenType::Comma
+            && tokens
+                .get(index + 1)
+                .is_some_and(|next| next.token_type == TokenType::From))
+}
+
+fn raw_token_text(sql: &str, token: &Token) -> Result<String, String> {
+    char_slice(sql, token.span.start, token.span.end)
+        .ok_or_else(|| TOKEN_PRESERVATION_FAILURE.to_string())
+}
+
 #[derive(Debug, Clone)]
 struct Comment {
     start: usize,

crates/sqlbuild-rules-native/src/sql_lint/tests.rs

@@ -1,3 +1,5 @@
+#[path = "tests/test_authored_tokens.rs"]
+mod authored_tokens;
 #[path = "tests/test_batch_lint.rs"]
 mod batch_lint;
 #[path = "tests/test_comment_formatter.rs"]

crates/sqlbuild-rules-native/src/sql_lint/tests/test_authored_tokens.rs

@@ -0,0 +1,97 @@
+use serde_json::{Value, json};
+
+use crate::sql_lint::_helpers::formatter::preserve_authored_tokens;
+use crate::sql_lint::main::formatter::format_json;
+use crate::sql_lint::tests::test_types;
+
+#[test]
+fn given_generator_respellings_when_formatting_then_authored_spellings_are_kept()
+-> Result<(), String> {
+    let test_cases = [
+        test_types::FormatTestCase {
+            description: "Snowflake STARTSWITH keeps its authored name",
+            sql: "select startswith(order_code, 'EU') as is_eu_order from orders",
+            expected_sql: "SELECT\n  startswith(order_code, 'EU') AS is_eu_order\nFROM orders",
+            expected_changed: true,
+        },
+        test_types::FormatTestCase {
+            description: "function synonyms keep their authored names",
+            sql: "select SUBSTR(order_code, 1, 2) as region, TIMESTAMPDIFF(day, ordered_at, shipped_at) as days_to_ship, TRY_TO_DECIMAL(amount_text, 10, 2) as amount from orders",
+            expected_sql: "SELECT\n  SUBSTR(order_code, 1, 2) AS region,\n  TIMESTAMPDIFF(DAY, ordered_at, shipped_at) AS days_to_ship,\n  TRY_TO_DECIMAL(amount_text, 10, 2) AS amount\nFROM orders",
+            expected_changed: true,
+        },
+        test_types::FormatTestCase {
+            description: "operator and function spellings stay authored",
+            sql: "select IFNULL(status, 'open') as status from orders where quantity != 0",
+            expected_sql: "SELECT\n  IFNULL(status, 'open') AS status\nFROM orders\nWHERE\n  quantity != 0",
+            expected_changed: true,
+        },
+    ];
+    for test_case in &test_cases {
+        let response = format_json(
+            &json!({"version": 1, "sql": test_case.sql, "dialect": "snowflake"}).to_string(),
+        )?;
+        let payload: Value = serde_json::from_str(&response).map_err(|error| error.to_string())?;
+        assert_eq!(
+            payload["sql"], test_case.expected_sql,
+            "{}",
+            test_case.description
+        );
+        assert_eq!(
+            payload["changed"], test_case.expected_changed,
+            "{}",
+            test_case.description
+        );
+    }
+    Ok(())
+}
+
+#[test]
+fn given_generated_token_changes_when_preserving_authored_tokens_then_only_layout_survives() {
+    let dialect = polyglot_sql::Dialect::get(polyglot_sql::DialectType::Snowflake);
+    let test_cases = [
+        test_types::AuthoredTokenTestCase {
+            description: "renamed function is restored to its authored spelling",
+            authored: "select startswith(code, 'EU') from orders",
+            generated: "SELECT\n  STARTS_WITH(code, 'EU')\nFROM orders",
+            expected_sql: Some("SELECT\n  startswith(code, 'EU')\nFROM orders"),
+        },
+        test_types::AuthoredTokenTestCase {
+            description: "swapped arguments are restored to authored order",
+            authored: "select datediff(day, a, b) from orders",
+            generated: "SELECT\n  DATEDIFF(DAY, b, a)\nFROM orders",
+            expected_sql: Some("SELECT\n  DATEDIFF(DAY, a, b)\nFROM orders"),
+        },
+        test_types::AuthoredTokenTestCase {
+            description: "an explicit alias keyword may be added",
+            authored: "select a b from orders",
+            generated: "SELECT\n  a AS b\nFROM orders",
+            expected_sql: Some("SELECT\n  a AS b\nFROM orders"),
+        },
+        test_types::AuthoredTokenTestCase {
+            description: "a structural change refuses the format",
+            authored: "select a from orders where a <> 1",
+            generated: "SELECT\n  a\nFROM orders\nWHERE\n  NOT a = 1",
+            expected_sql: None,
+        },
+        test_types::AuthoredTokenTestCase {
+            description: "a dropped token refuses the format",
+            authored: "select a from orders left outer join items using (id)",
+            generated: "SELECT\n  a\nFROM orders\nLEFT JOIN items USING (id)",
+            expected_sql: None,
+        },
+    ];
+    for test_case in test_cases {
+        let actual = preserve_authored_tokens(
+            test_case.authored,
+            test_case.generated.to_string(),
+            &dialect,
+        );
+        assert_eq!(
+            actual.ok().as_deref(),
+            test_case.expected_sql,
+            "{}",
+            test_case.description
+        );
+    }
+}

crates/sqlbuild-rules-native/src/sql_lint/tests/test_types.rs

@@ -104,3 +104,11 @@ pub(crate) struct FunctionDepthFailureTestCase {
     pub expected_code: &'static str,
     pub expected_limit: &'static str,
 }
+
+pub(crate) struct AuthoredTokenTestCase {
+    pub description: &'static str,
+    pub authored: &'static str,
+    pub generated: &'static str,
+    /// `None` when the formatter must refuse the change.
+    pub expected_sql: Option<&'static str>,
+}

tests/integration/src/sqlbuild/cli/commands/main/_test_types.py

@@ -149,6 +149,16 @@ class FormatCompileIntegrationTestCase:
     expected_literal: str
 
 
+@dataclass(frozen=True)
+class AuthoredSpellingFormatIntegrationTestCase:
+    """One authored function or operator spelling that formatting must keep."""
+
+    description: str
+    adapter: str
+    authored_expression: str
+    expected_expression: str
+
+
 @dataclass(frozen=True)
 class BacktickDialectFormatIntegrationTestCase:
     """One backtick-identifier adapter whose macro call must survive formatting."""

tests/integration/src/sqlbuild/cli/commands/main/test_format.py

@@ -10,6 +10,7 @@
 
 from sqlbuild.cli.commands.main.entrypoint.entry import main
 from tests.integration.src.sqlbuild.cli.commands.main._test_types import (
+    AuthoredSpellingFormatIntegrationTestCase,
     BacktickDialectFormatIntegrationTestCase,
     CanonicalFixtureFormatIntegrationTestCase,
     DescriptionFormatIntegrationTestCase,
@@ -32,6 +33,65 @@
 _UNFORMATTED_ORDERS_SQL: str = "MODEL (materialized table);\nselect   1 as order_id\n"
 
 
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        AuthoredSpellingFormatIntegrationTestCase(
+            "Snowflake STARTSWITH is not renamed",
+            "snowflake",
+            "STARTSWITH(order_code, 'EU')",
+            "STARTSWITH(order_code, 'EU')",
+        ),
+        AuthoredSpellingFormatIntegrationTestCase(
+            "Snowflake SUBSTR synonym is kept",
+            "snowflake",
+            "SUBSTR(order_code, 1, 2)",
+            "SUBSTR(order_code, 1, 2)",
+        ),
+        AuthoredSpellingFormatIntegrationTestCase(
+            "Snowflake TIMESTAMPDIFF synonym is kept",
+            "snowflake",
+            "TIMESTAMPDIFF(day, ordered_at, shipped_at)",
+            "TIMESTAMPDIFF(DAY, ordered_at, shipped_at)",
+        ),
+        AuthoredSpellingFormatIntegrationTestCase(
+            "Snowflake TRY_TO_DECIMAL synonym is kept",
+            "snowflake",
+            "TRY_TO_DECIMAL(amount_text, 10, 2)",
+            "TRY_TO_DECIMAL(amount_text, 10, 2)",
+        ),
+        AuthoredSpellingFormatIntegrationTestCase(
+            "DuckDB IFNULL and != spellings are kept",
+            "duckdb",
+            "IFNULL(status, 'open') != 'x'",
+            "IFNULL(status, 'open') != 'x'",
+        ),
+    ],
+    ids=lambda case: case.description,
+)
+def test_given_authored_function_spelling_when_formatting_then_only_layout_changes(
+    test_case: AuthoredSpellingFormatIntegrationTestCase,
+    tmp_path: Path,
+    capsys: pytest.CaptureFixture[str],
+) -> None:
+    """Formatting must never rewrite a function name, including to one the warehouse rejects."""
+    (tmp_path / "sqlbuild_project.toml").write_text(
+        f'name = "orders"\nadapter = "{test_case.adapter}"\n'
+    )
+    model: Path = tmp_path / "models" / "orders.sql"
+    model.parent.mkdir()
+    model.write_text(
+        'MODEL (description "Order flags");\n'
+        f'select   {test_case.authored_expression} as order_flag from __source("raw_orders")\n'
+    )
+    assert main(["--project-dir", str(tmp_path), "format", "--json"]) == 0
+    capsys.readouterr()
+    formatted: str = model.read_text()
+    assert f"  {test_case.expected_expression} AS order_flag\n" in formatted
+    assert main(["--project-dir", str(tmp_path), "format", "--check", "--json"]) == 0
+    assert model.read_text() == formatted
+
+
 @pytest.mark.parametrize(
     "test_case",
     [