fix: keep authored SQL tokens when formatting

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",
     [