From 04bf374c1c706171d4e9cc5833d9cd4651496465 Mon Sep 17 00:00:00 2001
From: Ted Conbeer <tconbeer@users.noreply.github.com>
Date: Fri, 24 Jan 2025 15:19:28 -0700
Subject: [PATCH] fix/#637/db variant (#661)

* fix #637: no space between a colon and square bracket

* chore: update changelog
---
 CHANGELOG.md                                      | 3 +++
 src/sqlfmt/node_manager.py                        | 9 +++++++++
 tests/data/unformatted/136_databricks_variant.sql | 6 ++++++
 tests/functional_tests/test_general_formatting.py | 1 +
 tests/unit_tests/test_node_manager.py             | 1 +
 5 files changed, 20 insertions(+)
 create mode 100644 tests/data/unformatted/136_databricks_variant.sql

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1ae2acc..8437065 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,9 @@ All notable changes to this project will be documented in this file.
 ### Formatting Changes and Bug Fixes
 
 - sqlfmt no longer adds a space between the `*` and `columns` in DuckDB `*columns` expressions ([#657](https://github.com/tconbeer/sqlfmt/issues/657) - thank you [@aersam](https://github.com/aersam)!)
+- sqlfmt no longer adds a space between the `:` and `['name']` in a DataBricks escaped variant expression like `foo:['bar.baz']` ([#637](https://github.com/tconbeer/sqlfmt/issues/637) - thank you [@aersam](https://github.com/aersam)!)
+
+### Testing
 - sqlfmt is now tested against and fully supports Python 3.13
 
 ## [0.24.0] - 2024-11-22
diff --git a/src/sqlfmt/node_manager.py b/src/sqlfmt/node_manager.py
index 1d0ccb1..88e95e5 100644
--- a/src/sqlfmt/node_manager.py
+++ b/src/sqlfmt/node_manager.py
@@ -217,6 +217,15 @@ def whitespace(
             )
         ):
             return NO_SPACE
+        # open square brackets that follow colons are escaped databricks
+        # variant cols
+        elif (
+            token.type is TokenType.BRACKET_OPEN
+            and token.token == "["
+            and previous_token
+            and previous_token.type is TokenType.COLON
+        ):
+            return NO_SPACE
         # need a space before any other open bracket
         elif token.type is TokenType.BRACKET_OPEN:
             return SPACE
diff --git a/tests/data/unformatted/136_databricks_variant.sql b/tests/data/unformatted/136_databricks_variant.sql
new file mode 100644
index 0000000..8eada68
--- /dev/null
+++ b/tests/data/unformatted/136_databricks_variant.sql
@@ -0,0 +1,6 @@
+-- https://github.com/tconbeer/sqlfmt/issues/637
+SELECT e.col:['attr_ID']::string as id
+FROM expl e
+)))))__SQLFMT_OUTPUT__(((((
+-- https://github.com/tconbeer/sqlfmt/issues/637
+select e.col:['attr_ID']::string as id from expl e
diff --git a/tests/functional_tests/test_general_formatting.py b/tests/functional_tests/test_general_formatting.py
index 605b4eb..5a2a7fe 100644
--- a/tests/functional_tests/test_general_formatting.py
+++ b/tests/functional_tests/test_general_formatting.py
@@ -56,6 +56,7 @@
         "unformatted/133_for_else.sql",
         "unformatted/134_databricks_type_hints.sql",
         "unformatted/135_star_columns.sql",
+        "unformatted/136_databricks_variant.sql",
         "unformatted/200_base_model.sql",
         "unformatted/201_basic_snapshot.sql",
         "unformatted/202_unpivot_macro.sql",
diff --git a/tests/unit_tests/test_node_manager.py b/tests/unit_tests/test_node_manager.py
index 1622dde..f3dd4d5 100644
--- a/tests/unit_tests/test_node_manager.py
+++ b/tests/unit_tests/test_node_manager.py
@@ -314,6 +314,7 @@ def test_identifier_whitespace(default_mode: Mode, source_string: str) -> None:
         "([])\n",
         "()[] + foo()[offset(1)]\n",
         "using (id)\n",
+        "foo:['bar.baz']\n",
     ],
 )
 def test_bracket_whitespace(default_mode: Mode, source_string: str) -> None: