fix: lex *columns as a name before lexing *

tconbeer · Jan 24, 2025 · 1bee229 · 1bee229
1 parent 6644340
commit 1bee229
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 0 deletions.
diff --git a/src/sqlfmt/rules/__init__.py b/src/sqlfmt/rules/__init__.py
@@ -21,6 +21,12 @@
 
 MAIN = [
     *CORE,
+    Rule(
+        name="star_columns",
+        priority=409,  # star is 410
+        pattern=group(r"\*columns") + group(r"\(", r"$"),
+        action=partial(actions.add_node_to_buffer, token_type=TokenType.NAME),
+    ),
     Rule(
         name="statement_start",
         priority=1000,

diff --git a/tests/data/unformatted/135_star_columns.sql b/tests/data/unformatted/135_star_columns.sql
@@ -0,0 +1,5 @@
+-- https://github.com/tconbeer/sqlfmt/issues/657
+SELECT coalesce(*COLUMNS(['a', 'b', 'c'])) AS result
+)))))__SQLFMT_OUTPUT__(((((
+-- https://github.com/tconbeer/sqlfmt/issues/657
+select coalesce(*columns(['a', 'b', 'c'])) as result
diff --git a/tests/functional_tests/test_general_formatting.py b/tests/functional_tests/test_general_formatting.py
@@ -55,6 +55,7 @@
         "unformatted/132_spark_number_literals.sql",
         "unformatted/133_for_else.sql",
         "unformatted/134_databricks_type_hints.sql",
+        "unformatted/135_star_columns.sql",
         "unformatted/200_base_model.sql",
         "unformatted/201_basic_snapshot.sql",
         "unformatted/202_unpivot_macro.sql",

diff --git a/tests/unit_tests/test_rule.py b/tests/unit_tests/test_rule.py
@@ -238,6 +238,7 @@ def get_rule(ruleset: List[Rule], rule_name: str) -> Rule:
         (MAIN, "other_identifiers", "%s"),
         (MAIN, "other_identifiers", "%(name)s"),
         (MAIN, "other_identifiers", "%(anything! else!)s"),
+        (MAIN, "star_columns", "*columns"),
         (MAIN, "newline", "\n"),
         (JINJA, "jinja_comment", "{# my comment #}"),
         (JINJA, "jinja_comment", "{#-my comment -#}"),