fix(parser)!: JSON path key quote granularity (#7758)

Author: geooo109

sqlglot/dialects/dialect.py

@@ -2134,10 +2134,9 @@ def _json_extract_segments(self: Generator, expression: JSON_EXTRACT_TYPE) -> st
         if not isinstance(path, exp.JSONPath):
             return rename_func(name)(self, expression)
 
-        escape = path.args.get("escape")
-
         segments = []
         for segment in path.expressions:
+            escape = segment.args.get("quoted")
             path = self.sql(segment)
             if path:
                 if isinstance(segment, exp.JSONPathPart) and (

sqlglot/expressions/query.py

@@ -1967,7 +1967,7 @@ class JSON(Expression):
 
 
 class JSONPath(Expression):
-    arg_types = {"expressions": True, "escape": False}
+    arg_types = {"expressions": True}
 
     @property
     def output_name(self) -> str:
@@ -1984,7 +1984,7 @@ class JSONPathFilter(JSONPathPart):
 
 
 class JSONPathKey(JSONPathPart):
-    arg_types = {"this": True}
+    arg_types = {"this": True, "quoted": False}
 
 
 class JSONPathRecursive(JSONPathPart):

sqlglot/generator.py

@@ -489,6 +489,12 @@ class Generator:
     # Whether to escape keys using single quotes in JSON paths
     JSON_PATH_SINGLE_QUOTE_ESCAPE = False
 
+    # Whether a quoted JSON path key (e.g. from a quoted identifier or ['key'] bracket) must be
+    # rendered in bracket form to preserve its case-sensitivity, even if it would otherwise match
+    # SAFE_JSON_PATH_KEY_RE and render as a bare dotted key. Needed for dialects like Databricks
+    # where a bare colon key is case-insensitive but a bracketed key is case-sensitive.
+    JSON_PATH_KEY_QUOTED_FORCES_BRACKETS = False
+
     # The JSONPathPart expressions supported by this dialect
     SUPPORTED_JSON_PATH_PARTS: t.ClassVar = ALL_JSON_PATH_PARTS.copy()
 
@@ -3721,9 +3727,6 @@ def jsonkeyvalue_sql(self, expression: exp.JSONKeyValue) -> str:
     def jsonpath_sql(self, expression: exp.JSONPath) -> str:
         path = self.expressions(expression, sep="", flat=True).lstrip(".")
 
-        if expression.args.get("escape"):
-            path = self.escape_str(path)
-
         if self.QUOTE_JSON_PATH:
             path = f"{self.dialect.QUOTE_START}{path}{self.dialect.QUOTE_END}"
 
@@ -5238,10 +5241,20 @@ def _jsonpathkey_sql(self, expression: exp.JSONPathKey) -> str:
             this = self.json_path_part(this)
             return f".{this}" if this else ""
 
-        if self.SAFE_JSON_PATH_KEY_RE.match(this):
+        quoted = expression.args.get("quoted")
+        if not (
+            quoted and self.JSON_PATH_KEY_QUOTED_FORCES_BRACKETS
+        ) and self.SAFE_JSON_PATH_KEY_RE.match(this):
             return f".{this}"
 
         this = self.json_path_part(this)
+
+        if quoted and self.QUOTE_JSON_PATH:
+            # The whole path is rendered as a single quoted string literal, so the bracketed key
+            # (which may itself contain backslash-escaped quotes, e.g. ["x \"y\"z"]) must be
+            # escaped again for the outer string literal (-> ["x \\"y\\"z"]).
+            this = self.escape_str(this)
+
         return (
             f"[{this}]"
             if self._quote_json_path_key_using_brackets and self.JSON_PATH_BRACKETED_KEY_SUPPORTED

sqlglot/generators/databricks.py

@@ -15,6 +15,7 @@ class DatabricksGenerator(SparkGenerator):
     COPY_PARAMS_ARE_WRAPPED = False
     COPY_PARAMS_EQ_REQUIRED = True
     JSON_PATH_SINGLE_QUOTE_ESCAPE = False
+    JSON_PATH_KEY_QUOTED_FORCES_BRACKETS = True
     SAFE_JSON_PATH_KEY_RE = exp.SAFE_IDENTIFIER_RE
     QUOTE_JSON_PATH = False
     PARSE_JSON_NAME: str | None = "PARSE_JSON"

sqlglot/parser.py

@@ -6669,13 +6669,12 @@ def _build_json_extract(
         self,
         this: exp.Expr | None,
         path_parts: list[exp.JSONPathPart],
-        escape: bool | None,
     ) -> tuple[exp.Expr | None, list[exp.JSONPathPart]]:
         if len(path_parts) > 1:
             this = self.expression(
                 exp.JSONExtract(
                     this=this,
-                    expression=exp.JSONPath(expressions=path_parts, escape=escape),
+                    expression=exp.JSONPath(expressions=path_parts),
                     variant_extract=True,
                     requires_json=self.JSON_EXTRACT_REQUIRES_JSON_EXPRESSION,
                 )
@@ -6686,28 +6685,24 @@ def _build_json_extract(
 
     def _parse_colon_as_variant_extract(self, this: exp.Expr | None) -> exp.Expr | None:
         path_parts: list[exp.JSONPathPart] = [exp.JSONPathRoot()]
-        escape = None
 
         while self._match(TokenType.COLON):
             if not self.COLON_CHAIN_IS_SINGLE_EXTRACT:
-                this, path_parts = self._build_json_extract(this, path_parts, escape)
-                escape = None
+                this, path_parts = self._build_json_extract(this, path_parts)
 
             key = self._parse_id_var(any_token=True, tokens=(TokenType.SELECT,))
 
             if key:
-                if isinstance(key, exp.Identifier) and key.quoted:
-                    escape = True
-                path_parts.append(exp.JSONPathKey(this=key.name))
+                quoted = isinstance(key, exp.Identifier) and key.quoted
+                path_parts.append(exp.JSONPathKey(this=key.name, quoted=quoted))
 
             while True:
                 if self._match(TokenType.DOT):
                     next_key = self._parse_id_var(any_token=True, tokens=(TokenType.SELECT,))
 
                     if next_key:
-                        if isinstance(next_key, exp.Identifier) and next_key.quoted:
-                            escape = True
-                        path_parts.append(exp.JSONPathKey(this=next_key.name))
+                        quoted = isinstance(next_key, exp.Identifier) and next_key.quoted
+                        path_parts.append(exp.JSONPathKey(this=next_key.name, quoted=quoted))
                 elif self._match(TokenType.L_BRACKET):
                     bracket_expr = self._parse_bracket_key_value()
 
@@ -6716,15 +6711,13 @@ def _parse_colon_as_variant_extract(self, this: exp.Expr | None) -> exp.Expr | N
 
                     if bracket_expr:
                         if bracket_expr.is_string:
-                            path_parts.append(exp.JSONPathKey(this=bracket_expr.name))
-                            escape = True
+                            path_parts.append(exp.JSONPathKey(this=bracket_expr.name, quoted=True))
                         elif bracket_expr.is_star:
                             path_parts.append(exp.JSONPathSubscript(this=exp.JSONPathWildcard()))
                         elif bracket_expr.is_number:
                             path_parts.append(exp.JSONPathSubscript(this=bracket_expr.to_py()))
                         else:
-                            this, path_parts = self._build_json_extract(this, path_parts, escape)
-                            escape = None
+                            this, path_parts = self._build_json_extract(this, path_parts)
 
                             this = self.expression(
                                 exp.Bracket(
@@ -6733,8 +6726,7 @@ def _parse_colon_as_variant_extract(self, this: exp.Expr | None) -> exp.Expr | N
                             )
 
                 elif self._match(TokenType.DCOLON):
-                    this, path_parts = self._build_json_extract(this, path_parts, escape)
-                    escape = None
+                    this, path_parts = self._build_json_extract(this, path_parts)
 
                     cast_type = self._parse_types()
                     if cast_type:
@@ -6744,7 +6736,7 @@ def _parse_colon_as_variant_extract(self, this: exp.Expr | None) -> exp.Expr | N
                 else:
                     break
 
-        this, _ = self._build_json_extract(this, path_parts, escape)
+        this, _ = self._build_json_extract(this, path_parts)
 
         return this
 

tests/dialects/test_databricks.py

@@ -294,7 +294,7 @@ def test_json(self):
         )
         self.validate_identity(
             """SELECT c1:['price'] FROM VALUES ('{ "price": 5 }') 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 GET_JSON_OBJECT(c1, '$.price') FROM VALUES ('{ "price": 5 }') AS T(c1)"""
@@ -304,7 +304,7 @@ def test_json(self):
         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"]""",
+            """SELECT raw:["zip code"], raw:["fb:testid"], raw:store["bicycle"], raw:store["zip code"]""",
         )
         self.validate_all(
             "SELECT col:`fr'uit`",