fix(databricks)!: GET_JSON_OBJECT roundtrip (#7728)

Author: geooo109

sqlglot-integration-tests

@@ -10,14 +10,6 @@
 from sqlglot.generators.spark import SparkGenerator
 
 
-def _jsonextract_sql(
-    self: DatabricksGenerator, expression: exp.JSONExtract | exp.JSONExtractScalar
-) -> str:
-    this = self.sql(expression, "this")
-    expr = self.sql(expression, "expression")
-    return f"{this}:{expr}"
-
-
 class DatabricksGenerator(SparkGenerator):
     TABLESAMPLE_SEED_KEYWORD = "REPEATABLE"
     COPY_PARAMS_ARE_WRAPPED = False
@@ -52,9 +44,10 @@ class DatabricksGenerator(SparkGenerator):
                     transforms.any_to_exists,
                 ]
             ),
-            exp.JSONExtract: _jsonextract_sql,
-            exp.JSONExtractScalar: _jsonextract_sql,
-            exp.JSONPathRoot: lambda *_: "",
+            exp.JSONExtract: lambda self, e: f"{self.sql(e, 'this')}:{self.sql(e, 'expression')}",
+            exp.JSONPathRoot: lambda self, e: (
+                "$" if isinstance(e.parent and e.parent.parent, exp.JSONExtractScalar) else ""
+            ),
             exp.ToChar: lambda self, e: (
                 self.cast_sql(exp.Cast(this=e.this, to=exp.DataType(this="STRING")))
                 if e.args.get("is_numeric")
@@ -98,7 +91,12 @@ def columndef_sql(self, expression: exp.ColumnDef, sep: str = " ") -> str:
 
     def jsonpath_sql(self, expression: exp.JSONPath) -> str:
         expression.set("escape", None)
-        return super().jsonpath_sql(expression)
+        path = super().jsonpath_sql(expression)
+
+        if isinstance(expression.parent, exp.JSONExtractScalar):
+            return f"{self.dialect.QUOTE_START}{path}{self.dialect.QUOTE_END}"
+
+        return path
 
     def uniform_sql(self, expression: exp.Uniform) -> str:
         gen = expression.args.get("gen")

sqlglot/generators/databricks.py

@@ -141,9 +141,6 @@ def test_databricks(self):
 
         self.validate_all(
             "SELECT c1:item[1].price",
-            read={
-                "spark": "SELECT GET_JSON_OBJECT(c1, '$.item[1].price')",
-            },
             write={
                 "databricks": "SELECT c1:item[1].price",
                 "spark": "SELECT GET_JSON_OBJECT(c1, '$.item[1].price')",
@@ -152,8 +149,11 @@ def test_databricks(self):
 
         self.validate_all(
             "SELECT GET_JSON_OBJECT(c1, '$.item[1].price')",
+            read={
+                "spark": "SELECT GET_JSON_OBJECT(c1, '$.item[1].price')",
+            },
             write={
-                "databricks": "SELECT c1:item[1].price",
+                "databricks": "SELECT GET_JSON_OBJECT(c1, '$.item[1].price')",
                 "spark": "SELECT GET_JSON_OBJECT(c1, '$.item[1].price')",
             },
         )
@@ -271,7 +271,7 @@ def test_databricks(self):
 
         self.validate_identity(
             """WITH t AS (SELECT '{"x-y": "z"}' AS c) SELECT get_json_object(c, '$.x-y') FROM t""",
-            """WITH t AS (SELECT '{"x-y": "z"}' AS c) SELECT c:["x-y"] FROM t""",
+            """WITH t AS (SELECT '{"x-y": "z"}' AS c) SELECT GET_JSON_OBJECT(c, '$["x-y"]') FROM t""",
         ).selects[0].expression.assert_is(exp.JSONPath)
 
     # https://docs.databricks.com/sql/language-manual/functions/colonsign.html
@@ -296,9 +296,11 @@ def test_json(self):
             """SELECT c1:price FROM VALUES ('{ "price": 5 }') AS T(c1)""",
         )
         self.validate_identity(
-            """SELECT GET_JSON_OBJECT(c1, '$.price') FROM VALUES ('{ "price": 5 }') AS T(c1)""",
-            """SELECT c1:price FROM VALUES ('{ "price": 5 }') AS T(c1)""",
+            """SELECT GET_JSON_OBJECT(c1, '$.price') FROM VALUES ('{ "price": 5 }') AS T(c1)"""
         )
+        self.validate_identity("SELECT GET_JSON_OBJECT(col, path_col)")
+        self.validate_identity("SELECT GET_JSON_OBJECT(col, CONCAT('$.', field_name))")
+        self.validate_identity("SELECT GET_JSON_OBJECT(GET_JSON_OBJECT(col, '$[0]'), '$.a')")
         self.validate_identity(
             """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"]""",

tests/dialects/test_databricks.py

@@ -1020,7 +1020,7 @@ def test_hive(self):
                 "hive": """WITH t AS (SELECT '{"x-y": "z"}' AS c) SELECT GET_JSON_OBJECT(c, '$.x-y') FROM t""",
                 "spark2": """WITH t AS (SELECT '{"x-y": "z"}' AS c) SELECT GET_JSON_OBJECT(c, '$.x-y') FROM t""",
                 "spark": """WITH t AS (SELECT '{"x-y": "z"}' AS c) SELECT GET_JSON_OBJECT(c, '$.x-y') FROM t""",
-                "databricks": """WITH t AS (SELECT '{"x-y": "z"}' AS c) SELECT c:["x-y"] FROM t""",
+                "databricks": """WITH t AS (SELECT '{"x-y": "z"}' AS c) SELECT GET_JSON_OBJECT(c, '$["x-y"]') FROM t""",
             },
         )
 

tests/dialects/test_hive.py

@@ -32,7 +32,7 @@ def test_redshift(self):
             """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 '{ "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')""",
                 "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')""",