fix #212: support bq typed array literals

tconbeer · Sep 30, 2022 · 0fe9901 · 0fe9901
1 parent a220c49
commit 0fe9901
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,7 @@ All notable changes to this project will be documented in this file.
 
 ### Formatting Changes + Bug Fixes
 
+-   BigQuery typed array literals like `array<float64>[1, 2]` are now supported, and spaces will no longer be inserted around `<` and `>` ([#212](https://github.com/tconbeer/sqlfmt/issues/212))
 -   SparkSQL-specific keywords `tablesample`, `cluster by`, `distribute by`, `sort by`, and `lateral view` are now supported by the polyglot dialect ([#264](https://github.com/tconbeer/sqlfmt/issues/264))
 -   `pivot` and `unpivot` are now supported as word operators, and will have a space between the keyword and the following parentheses
 -   `values` is now supported as an unterminated keyword; tuples of values will be indented from the `values` keyword if they span more than one line ([#263](https://github.com/tconbeer/sqlfmt/issues/263))

diff --git a/src/sqlfmt/dialect.py b/src/sqlfmt/dialect.py
@@ -410,6 +410,17 @@ def __init__(self) -> None:
                     + group(r"\W", r"$"),
                     action=actions.handle_set_operator,
                 ),
+                Rule(
+                    name="bq_typed_array",
+                    priority=3000,
+                    pattern=group(
+                        r"array<\w+>",
+                    )
+                    + group(r"\[", r"$"),
+                    action=partial(
+                        actions.add_node_to_buffer, token_type=TokenType.NAME
+                    ),
+                ),
                 Rule(
                     name="name",
                     priority=5000,

diff --git a/tests/data/unformatted/124_bq_typed_array_literals.sql b/tests/data/unformatted/124_bq_typed_array_literals.sql
@@ -0,0 +1,10 @@
+SELECT ARRAY<FLOAT64>[1, 2, 3] as floats,
+STRUCT("Nathan" as name, ARRAY<FLOAT64>[] as laps),
+array<INT64>[3, 4, 5, 6, 1000000, 20000000, 30000000, 409000000, 5000000, 60000000, 700000] as ints
+)))))__SQLFMT_OUTPUT__(((((
+select
+    array<float64>[1, 2, 3] as floats,
+    struct("Nathan" as name, array<float64>[] as laps),
+    array<int64>[
+        3, 4, 5, 6, 1000000, 20000000, 30000000, 409000000, 5000000, 60000000, 700000
+    ] as ints
diff --git a/tests/functional_tests/test_general_formatting.py b/tests/functional_tests/test_general_formatting.py
@@ -36,6 +36,7 @@
         "unformatted/121_stubborn_merge_edge_cases.sql",
         "unformatted/122_values.sql",
         "unformatted/123_spark_keywords.sql",
+        "unformatted/124_bq_typed_array_literals.sql",
         "unformatted/200_base_model.sql",
         "unformatted/201_basic_snapshot.sql",
         "unformatted/202_unpivot_macro.sql",

diff --git a/tests/unit_tests/test_dialect.py b/tests/unit_tests/test_dialect.py
@@ -190,6 +190,7 @@ def test_rule_props_are_unique(self, dialect: Polyglot) -> None:
             ("main", "set_operator", "intersect"),
             ("main", "set_operator", "minus"),
             ("main", "set_operator", "except"),
+            ("main", "bq_typed_array", "array<INT64>"),
             ("main", "name", "my_table_45"),
             ("main", "name", "replace"),
             ("main", "other_identifiers", "$2"),
@@ -266,6 +267,11 @@ def test_regex_exact_match(
             ("main", "unterm_keyword", "for"),
             ("main", "star_replace_exclude", "replace"),
             ("main", "unterm_keyword", "selection"),
+            (
+                "main",
+                "bq_typed_array",
+                "array < something and int64 > something_else[0]",
+            ),
             ("jinja", "jinja_set_block_start", "{% set foo = 'baz' %}"),
         ],
     )