Repository: tobymao/sqlglot

Issue:
bigquery: json_value transpilation gives incorrect result or fails due to escaping
noticed via sqlmesh, which recently updated `sqlglot (25.19.0 -> 25.31.4)`.  

We have an expression like `select json_value('{"fu bar": 42}', "$.fu bar")` which runs correctly on bigquery yielding 42.   Previously this was transpiled unchanged, but now results in an incorrect selection (note the added `.`s):

```
>>> sqlglot.__version__
'25.31.4'
>>> print(sqlglot.transpile("""select json_value('{"fu bar": 42}', "$.fu bar")""", read="bigquery", write="bigquery")[0])
SELECT json_value('{"fu bar": 42}', '$.fu.bar')
```
This query now yields `null` instead of 42 in bigquery.

The space-containing-key can be escaped in bigquery with double-quotes for [json_value](https://cloud.google.com/bigquery/docs/reference/standard-sql/json_functions#json_value) which is apparently different from the bracketed escaping used by [json_extract](https://cloud.google.com/bigquery/docs/reference/standard-sql/json_functions#json_extract)... 🤯    So this also works successfully in bigquery: `SELECT json_value('{"fu bar": 42}', '$."fu bar"')`.

But it now gets transpiled to sql which throws an error, namely `Invalid token in JSONPath at: ['fu bar']`:

```
>>> print(sqlglot.transpile("""SELECT json_value('{"fu bar": 42}', '$."fu bar"')""", read="bigquery", write="bigquery")[0])
SELECT json_value('{"fu bar": 42}', '$[\'fu bar\']')
```

a workaround is apparently to use the deprecated json_extract[_scalar] functions.  i.e. the query `SELECT json_extract('{"fu bar": 42}', '$[\'fu bar\']')` works correctly and is transpiled unchanged:

```
>>> print(sqlglot.transpile("""SELECT json_extract('{"fu bar": 42}', '$[\\'fu bar\\']')""", read="bigquery", write="bigquery")[0])
SELECT json_extract('{"fu bar": 42}', '$[\'fu bar\']')
```

Hints from issue comments:
@tobymao: this is tricky, there are a couple of things we should do @VaggelisD / @georgesittas 

1. make it so json tokenizer is sqlglot only, we can do that with a new constant USE_RS_TOKENIZER = USE_RS_TOKENIZER inside the tokenizer
2. Override _scan_var so that it doesn't do a strip() on the peek so it can accept spaces
3. modify _scan so that it doesn't strip white spaces

this is legal in bigquery... json_value('...', '$. a b c '), notice the leading and trailing spaces

we also need to override the escaping rules. for json_value, it uses double quotes to escape, and for json_extract it uses [ to escape.@schmee-td: Ran into a variant of this, works as expected before 25.29.0 but gives incorrect output after:

`select JSON_VALUE(foo, '$.attributes."$message"') from table`

gets turned into

`SELECT JSON_VALUE(foo, '$.attributes[\'$message\']') FROM table`

which is not valid BQ syntax.

EDIT: On a closer read, this was already covered in @tobymao's reply: "we also need to override the escaping rules. for json_value, it uses double quotes to escape, and for json_extract it uses [ to escape." 🙏 

Solution patch:
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
index 5eb02a1492..2f3ac53efb 100644
--- a/sqlglot/dialects/bigquery.py
+++ b/sqlglot/dialects/bigquery.py
@@ -40,6 +40,11 @@
 logger = logging.getLogger("sqlglot")
 
 
+JSON_EXTRACT_TYPE = t.Union[exp.JSONExtract, exp.JSONExtractScalar, exp.JSONExtractArray]
+
+DQUOTES_ESCAPING_JSON_FUNCTIONS = ("JSON_QUERY", "JSON_VALUE", "JSON_QUERY_ARRAY")
+
+
 def _derived_table_values_to_unnest(self: BigQuery.Generator, expression: exp.Values) -> str:
     if not expression.find_ancestor(exp.From, exp.Join):
         return self.values_sql(expression)
@@ -324,6 +329,23 @@ def _build_contains_substring(args: t.List) -> exp.Contains | exp.Anonymous:
     return exp.Contains(this=this, expression=expr)
 
 
+def _json_extract_sql(self: BigQuery.Generator, expression: JSON_EXTRACT_TYPE) -> str:
+    name = (expression._meta and expression.meta.get("name")) or expression.sql_name()
+    upper = name.upper()
+
+    dquote_escaping = upper in DQUOTES_ESCAPING_JSON_FUNCTIONS
+
+    if dquote_escaping:
+        self._quote_json_path_key_using_brackets = False
+
+    sql = rename_func(upper)(self, expression)
+
+    if dquote_escaping:
+        self._quote_json_path_key_using_brackets = True
+
+    return sql
+
+
 class BigQuery(Dialect):
     WEEK_OFFSET = -1
     UNNEST_COLUMN_ONLY = True
@@ -869,6 +891,9 @@ class Generator(generator.Generator):
             exp.ILike: no_ilike_sql,
             exp.IntDiv: rename_func("DIV"),
             exp.Int64: rename_func("INT64"),
+            exp.JSONExtract: _json_extract_sql,
+            exp.JSONExtractArray: _json_extract_sql,
+            exp.JSONExtractScalar: _json_extract_sql,
             exp.JSONFormat: rename_func("TO_JSON_STRING"),
             exp.Levenshtein: _levenshtein_sql,
             exp.Max: max_or_greatest,
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 5e8e1dccba..842df5a753 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -658,6 +658,7 @@ class Generator(metaclass=_Generator):
         "_next_name",
         "_identifier_start",
         "_identifier_end",
+        "_quote_json_path_key_using_brackets",
     )
 
     def __init__(
@@ -706,6 +707,8 @@ def __init__(
         self._identifier_start = self.dialect.IDENTIFIER_START
         self._identifier_end = self.dialect.IDENTIFIER_END
 
+        self._quote_json_path_key_using_brackets = True
+
     def generate(self, expression: exp.Expression, copy: bool = True) -> str:
         """
         Generates the SQL string corresponding to the given syntax tree.
@@ -2871,7 +2874,7 @@ def json_path_part(self, expression: int | str | exp.JSONPathPart) -> str:
         if isinstance(expression, int):
             return str(expression)
 
-        if self.JSON_PATH_SINGLE_QUOTE_ESCAPE:
+        if self._quote_json_path_key_using_brackets and self.JSON_PATH_SINGLE_QUOTE_ESCAPE:
             escaped = expression.replace("'", "\\'")
             escaped = f"\\'{expression}\\'"
         else:
@@ -4072,7 +4075,11 @@ def _jsonpathkey_sql(self, expression: exp.JSONPathKey) -> str:
             return f".{this}"
 
         this = self.json_path_part(this)
-        return f"[{this}]" if self.JSON_PATH_BRACKETED_KEY_SUPPORTED else f".{this}"
+        return (
+            f"[{this}]"
+            if self._quote_json_path_key_using_brackets and self.JSON_PATH_BRACKETED_KEY_SUPPORTED
+            else f".{this}"
+        )
 
     def _jsonpathsubscript_sql(self, expression: exp.JSONPathSubscript) -> str:
         this = self.json_path_part(expression.this)
diff --git a/sqlglot/jsonpath.py b/sqlglot/jsonpath.py
index 911debe4b5..115bd15941 100644
--- a/sqlglot/jsonpath.py
+++ b/sqlglot/jsonpath.py
@@ -146,6 +146,28 @@ def _parse_bracket() -> exp.JSONPathPart:
 
         return node
 
+    def _parse_var_text() -> str:
+        """
+        Consumes & returns the text for a var. In BigQuery it's valid to have a key with spaces
+        in it, e.g JSON_QUERY(..., '$. a b c ') should produce a single JSONPathKey(' a b c ').
+        This is done by merging "consecutive" vars until a key separator is found (dot, colon etc)
+        or the path string is exhausted.
+        """
+        prev_index = i - 2
+
+        while _match(TokenType.VAR):
+            pass
+
+        start = 0 if prev_index < 0 else tokens[prev_index].end + 1
+
+        if i >= len(tokens):
+            # This key is the last token for the path, so it's text is the remaining path
+            text = path[start:]
+        else:
+            text = path[start : tokens[i].start]
+
+        return text
+
     # We canonicalize the JSON path AST so that it always starts with a
     # "root" element, so paths like "field" will be generated as "$.field"
     _match(TokenType.DOLLAR)
@@ -155,8 +177,10 @@ def _parse_bracket() -> exp.JSONPathPart:
         if _match(TokenType.DOT) or _match(TokenType.COLON):
             recursive = _prev().text == ".."
 
-            if _match(TokenType.VAR) or _match(TokenType.IDENTIFIER):
-                value: t.Optional[str | exp.JSONPathWildcard] = _prev().text
+            if _match(TokenType.VAR):
+                value: t.Optional[str | exp.JSONPathWildcard] = _parse_var_text()
+            elif _match(TokenType.IDENTIFIER):
+                value = _prev().text
             elif _match(TokenType.STAR):
                 value = exp.JSONPathWildcard()
             else:
@@ -170,7 +194,9 @@ def _parse_bracket() -> exp.JSONPathPart:
                 raise ParseError(_error("Expected key name or * after DOT"))
         elif _match(TokenType.L_BRACKET):
             expressions.append(_parse_bracket())
-        elif _match(TokenType.VAR) or _match(TokenType.IDENTIFIER):
+        elif _match(TokenType.VAR):
+            expressions.append(exp.JSONPathKey(this=_parse_var_text()))
+        elif _match(TokenType.IDENTIFIER):
             expressions.append(exp.JSONPathKey(this=_prev().text))
         elif _match(TokenType.STAR):
             expressions.append(exp.JSONPathWildcard())


Test patch:
diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
index 366cadedbe..26b12a15fa 100644
--- a/tests/dialects/test_bigquery.py
+++ b/tests/dialects/test_bigquery.py
@@ -1574,14 +1574,6 @@ def test_bigquery(self):
                 "snowflake": "IFF((y) <> 0, (x) / (y), NULL)",
             },
         )
-        self.validate_all(
-            """SELECT JSON_QUERY('{"class": {"students": []}}', '$.class')""",
-            write={
-                "bigquery": """SELECT JSON_QUERY('{"class": {"students": []}}', '$.class')""",
-                "duckdb": """SELECT '{"class": {"students": []}}' -> '$.class'""",
-                "snowflake": """SELECT GET_PATH(PARSE_JSON('{"class": {"students": []}}'), 'class')""",
-            },
-        )
         self.validate_all(
             """SELECT JSON_VALUE_ARRAY('{"arr": [1, "a"]}', '$.arr')""",
             write={
@@ -2139,7 +2131,16 @@ def test_null_ordering(self):
                 },
             )
 
-    def test_json_extract_scalar(self):
+    def test_json_extract(self):
+        self.validate_all(
+            """SELECT JSON_QUERY('{"class": {"students": []}}', '$.class')""",
+            write={
+                "bigquery": """SELECT JSON_QUERY('{"class": {"students": []}}', '$.class')""",
+                "duckdb": """SELECT '{"class": {"students": []}}' -> '$.class'""",
+                "snowflake": """SELECT GET_PATH(PARSE_JSON('{"class": {"students": []}}'), 'class')""",
+            },
+        )
+
         for func in ("JSON_EXTRACT_SCALAR", "JSON_VALUE"):
             with self.subTest(f"Testing BigQuery's {func}"):
                 self.validate_all(
@@ -2164,6 +2165,18 @@ def test_json_extract_scalar(self):
                     self.parse_one(sql).sql("bigquery", normalize_functions="upper"), sql
                 )
 
+        # Test double quote escaping
+        for func in ("JSON_VALUE", "JSON_QUERY", "JSON_QUERY_ARRAY"):
+            self.validate_identity(
+                f"{func}(doc, '$. a b c .d')", f"""{func}(doc, '$." a b c ".d')"""
+            )
+
+        # Test single quote & bracket escaping
+        for func in ("JSON_EXTRACT", "JSON_EXTRACT_SCALAR", "JSON_EXTRACT_ARRAY"):
+            self.validate_identity(
+                f"{func}(doc, '$. a b c .d')", f"""{func}(doc, '$[\\' a b c \\'].d')"""
+            )
+
     def test_json_extract_array(self):
         for func in ("JSON_QUERY_ARRAY", "JSON_EXTRACT_ARRAY"):
             with self.subTest(f"Testing BigQuery's {func}"):