Repository: tobymao/sqlglot

Issue:
Sqlglot does not preserve new Databricks's Variant Syntax
The newest Databricks Runtime (15.3) does support a number of fancy syntax for their new Variant Type.
Docs: https://docs.databricks.com/en/semi-structured/variant.html

Those do not get preserved

```python
import sqlglot as sg
print(sg.parse_one("SELECT raw:store.bicycle FROM store_data", dialect="databricks").sql("databricks"))
-- prints "SELECT GET_JSON_OBJECT(raw, '$.store.bicycle') FROM store_data"

print(sg.parse_one("SELECT raw:`zip code`, raw:`fb:testid` FROM store_data", dialect="databricks").sql("databricks"))
-- prints "SELECT GET_JSON_OBJECT(raw, '$[\\'`zip\\'][\\'code`\\']'), GET_JSON_OBJECT(raw, '$[\\'`fb\\'][\\'testid`\\']') FROM store_data"
```







Discussion:
VaggelisD: As discussed in Slack, made `JSONPathTokenizer` a `Dialect` attribute similarly to `Tokenizer`, `Parser`, `Generator` in an attempt to make it parameterizable by each dialect

Patch:
diff --git a/sqlglot/dialects/databricks.py b/sqlglot/dialects/databricks.py
index deddb6e1..fc1772d5 100644
--- a/sqlglot/dialects/databricks.py
+++ b/sqlglot/dialects/databricks.py
@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-from sqlglot import exp, transforms
+import typing as t
+
+from sqlglot import exp, transforms, jsonpath
 from sqlglot.dialects.dialect import (
     date_delta_sql,
     build_date_delta,
@@ -10,16 +12,34 @@ from sqlglot.dialects.spark import Spark
 from sqlglot.tokens import TokenType
 
 
+def _build_json_extract(args: t.List) -> exp.JSONExtract:
+    # Transform GET_JSON_OBJECT(expr, '$.<path>') -> expr:<path>
+    this = args[0]
+    path = args[1].name.lstrip("$.")
+    return exp.JSONExtract(this=this, expression=path)
+
+
 def _timestamp_diff(
     self: Databricks.Generator, expression: exp.DatetimeDiff | exp.TimestampDiff
 ) -> str:
     return self.func("TIMESTAMPDIFF", expression.unit, expression.expression, expression.this)
 
 
+def _jsonextract_sql(
+    self: Databricks.Generator, expression: exp.JSONExtract | exp.JSONExtractScalar
+) -> str:
+    this = self.sql(expression, "this")
+    expr = self.sql(expression, "expression")
+    return f"{this}:{expr}"
+
+
 class Databricks(Spark):
     SAFE_DIVISION = False
     COPY_PARAMS_ARE_CSV = False
 
+    class JSONPathTokenizer(jsonpath.JSONPathTokenizer):
+        IDENTIFIERS = ["`", '"']
+
     class Parser(Spark.Parser):
         LOG_DEFAULTS_TO_LN = True
         STRICT_CAST = True
@@ -31,6 +51,7 @@ class Databricks(Spark):
             "DATE_ADD": build_date_delta(exp.DateAdd),
             "DATEDIFF": build_date_delta(exp.DateDiff),
             "TIMESTAMPDIFF": build_date_delta(exp.TimestampDiff),
+            "GET_JSON_OBJECT": _build_json_extract,
         }
 
         FACTOR = {
@@ -42,6 +63,8 @@ class Databricks(Spark):
         TABLESAMPLE_SEED_KEYWORD = "REPEATABLE"
         COPY_PARAMS_ARE_WRAPPED = False
         COPY_PARAMS_EQ_REQUIRED = True
+        JSON_PATH_SINGLE_QUOTE_ESCAPE = False
+        QUOTE_JSON_PATH = False
 
         TRANSFORMS = {
             **Spark.Generator.TRANSFORMS,
@@ -65,6 +88,9 @@ class Databricks(Spark):
                     transforms.unnest_to_explode,
                 ]
             ),
+            exp.JSONExtract: _jsonextract_sql,
+            exp.JSONExtractScalar: _jsonextract_sql,
+            exp.JSONPathRoot: lambda *_: "",
             exp.ToChar: lambda self, e: self.function_fallback_sql(e),
         }
 
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
index 01042a75..18b35a7d 100644
--- a/sqlglot/dialects/dialect.py
+++ b/sqlglot/dialects/dialect.py
@@ -9,7 +9,7 @@ from sqlglot import exp
 from sqlglot.errors import ParseError
 from sqlglot.generator import Generator
 from sqlglot.helper import AutoName, flatten, is_int, seq_get
-from sqlglot.jsonpath import parse as parse_json_path
+from sqlglot.jsonpath import JSONPathTokenizer, parse as parse_json_path
 from sqlglot.parser import Parser
 from sqlglot.time import TIMEZONES, format_time
 from sqlglot.tokens import Token, Tokenizer, TokenType
@@ -125,12 +125,16 @@ class _Dialect(type):
 
         base = seq_get(bases, 0)
         base_tokenizer = (getattr(base, "tokenizer_class", Tokenizer),)
+        base_jsonpath_tokenizer = (getattr(base, "jsonpath_tokenizer_class", JSONPathTokenizer),)
         base_parser = (getattr(base, "parser_class", Parser),)
         base_generator = (getattr(base, "generator_class", Generator),)
 
         klass.tokenizer_class = klass.__dict__.get(
             "Tokenizer", type("Tokenizer", base_tokenizer, {})
         )
+        klass.jsonpath_tokenizer_class = klass.__dict__.get(
+            "JSONPathTokenizer", type("JSONPathTokenizer", base_jsonpath_tokenizer, {})
+        )
         klass.parser_class = klass.__dict__.get("Parser", type("Parser", base_parser, {}))
         klass.generator_class = klass.__dict__.get(
             "Generator", type("Generator", base_generator, {})
@@ -324,6 +328,7 @@ class Dialect(metaclass=_Dialect):
     # --- Autofilled ---
 
     tokenizer_class = Tokenizer
+    jsonpath_tokenizer_class = JSONPathTokenizer
     parser_class = Parser
     generator_class = Generator
 
@@ -621,9 +626,8 @@ class Dialect(metaclass=_Dialect):
             path_text = path.name
             if path.is_number:
                 path_text = f"[{path_text}]"
-
             try:
-                return parse_json_path(path_text)
+                return parse_json_path(path_text, self)
             except ParseError as e:
                 logger.warning(f"Invalid JSON path syntax. {str(e)}")
 
@@ -655,6 +659,12 @@ class Dialect(metaclass=_Dialect):
             self._tokenizer = self.tokenizer_class(dialect=self)
         return self._tokenizer
 
+    @property
+    def jsonpath_tokenizer(self) -> JSONPathTokenizer:
+        if not hasattr(self, "_jsonpath_tokenizer"):
+            self._jsonpath_tokenizer = self.jsonpath_tokenizer_class(dialect=self)
+        return self._jsonpath_tokenizer
+
     def parser(self, **opts) -> Parser:
         return self.parser_class(dialect=self, **opts)
 
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index babfe340..c55aa62c 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -5457,7 +5457,7 @@ class OpenJSON(Func):
     arg_types = {"this": True, "path": False, "expressions": False}
 
 
-class JSONBContains(Binary):
+class JSONBContains(Binary, Func):
     _sql_names = ["JSONB_CONTAINS"]
 
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 3fef2a14..74ae5466 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -369,6 +369,9 @@ class Generator(metaclass=_Generator):
     # The keywords to use when prefixing & separating WITH based properties
     WITH_PROPERTIES_PREFIX = "WITH"
 
+    # Whether to quote the generated expression of exp.JsonPath
+    QUOTE_JSON_PATH = True
+
     TYPE_MAPPING = {
         exp.DataType.Type.NCHAR: "CHAR",
         exp.DataType.Type.NVARCHAR: "VARCHAR",
@@ -2657,7 +2660,10 @@ class Generator(metaclass=_Generator):
 
     def jsonpath_sql(self, expression: exp.JSONPath) -> str:
         path = self.expressions(expression, sep="", flat=True).lstrip(".")
-        return f"{self.dialect.QUOTE_START}{path}{self.dialect.QUOTE_END}"
+        if self.QUOTE_JSON_PATH:
+            path = f"{self.dialect.QUOTE_START}{path}{self.dialect.QUOTE_END}"
+
+        return path
 
     def json_path_part(self, expression: int | str | exp.JSONPathPart) -> str:
         if isinstance(expression, exp.JSONPathPart):
diff --git a/sqlglot/jsonpath.py b/sqlglot/jsonpath.py
index 129a4e6f..911debe4 100644
--- a/sqlglot/jsonpath.py
+++ b/sqlglot/jsonpath.py
@@ -8,6 +8,7 @@ from sqlglot.tokens import Token, Tokenizer, TokenType
 
 if t.TYPE_CHECKING:
     from sqlglot._typing import Lit
+    from sqlglot.dialects.dialect import DialectType
 
 
 class JSONPathTokenizer(Tokenizer):
@@ -36,9 +37,12 @@ class JSONPathTokenizer(Tokenizer):
     STRING_ESCAPES = ["\\"]
 
 
-def parse(path: str) -> exp.JSONPath:
+def parse(path: str, dialect: DialectType = None) -> exp.JSONPath:
     """Takes in a JSON path string and parses it into a JSONPath expression."""
-    tokens = JSONPathTokenizer().tokenize(path)
+    from sqlglot.dialects import Dialect
+
+    jsonpath_tokenizer = Dialect.get_or_raise(dialect).jsonpath_tokenizer
+    tokens = jsonpath_tokenizer.tokenize(path)
     size = len(tokens)
 
     i = 0
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 0d2cc2b9..0aae2e40 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -4486,10 +4486,22 @@ class Parser(metaclass=_Parser):
 
     def _parse_struct_types(self, type_required: bool = False) -> t.Optional[exp.Expression]:
         index = self._index
-        this = (
-            self._parse_type(parse_interval=False, fallback_to_identifier=True)
-            or self._parse_id_var()
-        )
+
+        if (
+            self._curr
+            and self._next
+            and self._curr.token_type in self.TYPE_TOKENS
+            and self._next.token_type in self.TYPE_TOKENS
+        ):
+            # Takes care of special cases like `STRUCT<list ARRAY<...>>` where the identifier is also a
+            # type token. Without this, the list will be parsed as a type and we'll eventually crash
+            this = self._parse_id_var()
+        else:
+            this = (
+                self._parse_type(parse_interval=False, fallback_to_identifier=True)
+                or self._parse_id_var()
+            )
+
         self._match(TokenType.COLON)
 
         if (


Test patch:
diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
index 42cb27f7..a7affe68 100644
--- a/tests/dialects/test_bigquery.py
+++ b/tests/dialects/test_bigquery.py
@@ -103,6 +103,7 @@ LANGUAGE js AS
         select_with_quoted_udf = self.validate_identity("SELECT `p.d.UdF`(data) FROM `p.d.t`")
         self.assertEqual(select_with_quoted_udf.selects[0].name, "p.d.UdF")
 
+        self.validate_identity("CAST(x AS STRUCT<list ARRAY<INT64>>)")
         self.validate_identity("assert.true(1 = 1)")
         self.validate_identity("SELECT ARRAY_TO_STRING(list, '--') AS text")
         self.validate_identity("SELECT jsondoc['some_key']")
diff --git a/tests/dialects/test_databricks.py b/tests/dialects/test_databricks.py
index 9ef3b86e..c88679dc 100644
--- a/tests/dialects/test_databricks.py
+++ b/tests/dialects/test_databricks.py
@@ -96,33 +96,30 @@ class TestDatabricks(Validator):
 
     # https://docs.databricks.com/sql/language-manual/functions/colonsign.html
     def test_json(self):
+        self.validate_identity("SELECT c1:price, c1:price.foo, c1:price.bar[1]")
         self.validate_identity(
-            """SELECT c1 : price FROM VALUES ('{ "price": 5 }') AS T(c1)""",
-            """SELECT GET_JSON_OBJECT(c1, '$.price') FROM VALUES ('{ "price": 5 }') AS T(c1)""",
+            """SELECT c1:item[1].price FROM VALUES ('{ "item": [ { "model" : "basic", "price" : 6.12 }, { "model" : "medium", "price" : 9.24 } ] }') AS T(c1)"""
         )
         self.validate_identity(
-            """SELECT c1:['price'] FROM VALUES('{ "price": 5 }') AS T(c1)""",
-            """SELECT GET_JSON_OBJECT(c1, '$.price') FROM VALUES ('{ "price": 5 }') AS T(c1)""",
+            """SELECT c1:item[*].price FROM VALUES ('{ "item": [ { "model" : "basic", "price" : 6.12 }, { "model" : "medium", "price" : 9.24 } ] }') AS T(c1)"""
         )
         self.validate_identity(
-            """SELECT c1:item[1].price FROM VALUES('{ "item": [ { "model" : "basic", "price" : 6.12 }, { "model" : "medium", "price" : 9.24 } ] }') AS T(c1)""",
-            """SELECT GET_JSON_OBJECT(c1, '$.item[1].price') FROM VALUES ('{ "item": [ { "model" : "basic", "price" : 6.12 }, { "model" : "medium", "price" : 9.24 } ] }') AS T(c1)""",
+            """SELECT FROM_JSON(c1:item[*].price, 'ARRAY<DOUBLE>')[0] FROM VALUES ('{ "item": [ { "model" : "basic", "price" : 6.12 }, { "model" : "medium", "price" : 9.24 } ] }') AS T(c1)"""
         )
         self.validate_identity(
-            """SELECT c1:item[*].price FROM VALUES('{ "item": [ { "model" : "basic", "price" : 6.12 }, { "model" : "medium", "price" : 9.24 } ] }') AS T(c1)""",
-            """SELECT GET_JSON_OBJECT(c1, '$.item[*].price') FROM VALUES ('{ "item": [ { "model" : "basic", "price" : 6.12 }, { "model" : "medium", "price" : 9.24 } ] }') AS T(c1)""",
+            """SELECT INLINE(FROM_JSON(c1:item[*], 'ARRAY<STRUCT<model STRING, price DOUBLE>>')) FROM VALUES ('{ "item": [ { "model" : "basic", "price" : 6.12 }, { "model" : "medium", "price" : 9.24 } ] }') AS T(c1)"""
         )
         self.validate_identity(
-            """SELECT from_json(c1:item[*].price, 'ARRAY<DOUBLE>')[0] FROM VALUES('{ "item": [ { "model" : "basic", "price" : 6.12 }, { "model" : "medium", "price" : 9.24 } ] }') AS T(c1)""",
-            """SELECT FROM_JSON(GET_JSON_OBJECT(c1, '$.item[*].price'), 'ARRAY<DOUBLE>')[0] FROM VALUES ('{ "item": [ { "model" : "basic", "price" : 6.12 }, { "model" : "medium", "price" : 9.24 } ] }') AS T(c1)""",
+            """SELECT c1:['price'] FROM VALUES ('{ "price": 5 }') AS T(c1)""",
+            """SELECT c1:price FROM VALUES ('{ "price": 5 }') AS T(c1)""",
         )
         self.validate_identity(
-            """SELECT inline(from_json(c1:item[*], 'ARRAY<STRUCT<model STRING, price DOUBLE>>')) FROM VALUES('{ "item": [ { "model" : "basic", "price" : 6.12 }, { "model" : "medium", "price" : 9.24 } ] }') AS T(c1)""",
-            """SELECT INLINE(FROM_JSON(GET_JSON_OBJECT(c1, '$.item[*]'), 'ARRAY<STRUCT<model STRING, price DOUBLE>>')) FROM VALUES ('{ "item": [ { "model" : "basic", "price" : 6.12 }, { "model" : "medium", "price" : 9.24 } ] }') AS T(c1)""",
+            """SELECT GET_JSON_OBJECT(c1, '$.price') FROM VALUES ('{ "price": 5 }') AS T(c1)""",
+            """SELECT c1:price FROM VALUES ('{ "price": 5 }') AS T(c1)""",
         )
         self.validate_identity(
-            "SELECT c1 : price",
-            "SELECT GET_JSON_OBJECT(c1, '$.price')",
+            """SELECT raw:`zip code`, raw:`fb:testid`, raw:store['bicycle'], raw:store["zip code"]""",
+            """SELECT raw:["zip code"], raw:["fb:testid"], raw:store.bicycle, raw:store["zip code"]""",
         )
 
     def test_datediff(self):
diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
index 071677db..d9d56754 100644
--- a/tests/dialects/test_postgres.py
+++ b/tests/dialects/test_postgres.py
@@ -8,6 +8,14 @@ class TestPostgres(Validator):
     dialect = "postgres"
 
     def test_postgres(self):
+        self.validate_all(
+            "x ? y",
+            write={
+                "": "JSONB_CONTAINS(x, y)",
+                "postgres": "x ? y",
+            },
+        )
+
         self.validate_identity("SHA384(x)")
         self.validate_identity(
             'CREATE TABLE x (a TEXT COLLATE "de_DE")', "CREATE TABLE x (a TEXT COLLATE de_DE)"
diff --git a/tests/dialects/test_redshift.py b/tests/dialects/test_redshift.py
index 87076fe7..26d87d21 100644
--- a/tests/dialects/test_redshift.py
+++ b/tests/dialects/test_redshift.py
@@ -28,7 +28,7 @@ class TestRedshift(Validator):
             """SELECT JSON_EXTRACT_PATH_TEXT('{ "farm": {"barn": { "color": "red", "feed stocked": true }}}', 'farm', 'barn', 'color')""",
             write={
                 "bigquery": """SELECT JSON_EXTRACT_SCALAR('{ "farm": {"barn": { "color": "red", "feed stocked": true }}}', '$.farm.barn.color')""",
-                "databricks": """SELECT GET_JSON_OBJECT('{ "farm": {"barn": { "color": "red", "feed stocked": true }}}', '$.farm.barn.color')""",
+                "databricks": """SELECT '{ "farm": {"barn": { "color": "red", "feed stocked": true }}}':farm.barn.color""",
                 "duckdb": """SELECT '{ "farm": {"barn": { "color": "red", "feed stocked": true }}}' ->> '$.farm.barn.color'""",
                 "postgres": """SELECT JSON_EXTRACT_PATH_TEXT('{ "farm": {"barn": { "color": "red", "feed stocked": true }}}', 'farm', 'barn', 'color')""",
                 "presto": """SELECT JSON_EXTRACT_SCALAR('{ "farm": {"barn": { "color": "red", "feed stocked": true }}}', '$.farm.barn.color')""",
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index 02615f24..9bdbf3ec 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -337,7 +337,7 @@ WHERE
             """SELECT PARSE_JSON('{"fruit":"banana"}'):fruit""",
             write={
                 "bigquery": """SELECT JSON_EXTRACT(PARSE_JSON('{"fruit":"banana"}'), '$.fruit')""",
-                "databricks": """SELECT GET_JSON_OBJECT('{"fruit":"banana"}', '$.fruit')""",
+                "databricks": """SELECT '{"fruit":"banana"}':fruit""",
                 "duckdb": """SELECT JSON('{"fruit":"banana"}') -> '$.fruit'""",
                 "mysql": """SELECT JSON_EXTRACT('{"fruit":"banana"}', '$.fruit')""",
                 "presto": """SELECT JSON_EXTRACT(JSON_PARSE('{"fruit":"banana"}'), '$.fruit')""",
diff --git a/tests/test_jsonpath.py b/tests/test_jsonpath.py
index 4daf3c1c..c939c523 100644
--- a/tests/test_jsonpath.py
+++ b/tests/test_jsonpath.py
@@ -2,8 +2,9 @@ import json
 import os
 import unittest
 
-from sqlglot import exp, jsonpath
+from sqlglot import exp
 from sqlglot.errors import ParseError, TokenError
+from sqlglot.jsonpath import parse
 from tests.helpers import FIXTURES_DIR
 
 
@@ -25,7 +26,7 @@ class TestJsonpath(unittest.TestCase):
             exp.JSONPathSelector(this=exp.JSONPathScript(this="@.x)")),
         ]
         self.assertEqual(
-            jsonpath.parse("$.*.a[0]['x'][*, 'y', 1].z[?(@.a == 'b'), 1:][1:5][1,?@.a][(@.x)]"),
+            parse("$.*.a[0]['x'][*, 'y', 1].z[?(@.a == 'b'), 1:][1:5][1,?@.a][(@.x)]"),
             exp.JSONPath(expressions=expected_expressions),
         )
 
@@ -36,7 +37,7 @@ class TestJsonpath(unittest.TestCase):
             ("$[((@.length-1))]", "$[((@.length-1))]"),
         ):
             with self.subTest(f"{selector} -> {expected}"):
-                self.assertEqual(jsonpath.parse(selector).sql(), f"'{expected}'")
+                self.assertEqual(parse(selector).sql(), f"'{expected}'")
 
     def test_cts_file(self):
         with open(os.path.join(FIXTURES_DIR, "jsonpath", "cts.json")) as file:
@@ -131,9 +132,9 @@ class TestJsonpath(unittest.TestCase):
             with self.subTest(f"{selector.strip()} /* {test['name']} */"):
                 if test.get("invalid_selector"):
                     try:
-                        jsonpath.parse(selector)
+                        parse(selector)
                     except (ParseError, TokenError):
                         pass
                 else:
-                    path = jsonpath.parse(selector)
+                    path = parse(selector)
                     self.assertEqual(path.sql(), f"'{overrides.get(selector, selector)}'")