Repository: tobymao/sqlglot

Issue:
can't parse JSON query when using `with` keyword
**Before you file an issue**
- Make sure you specify the "read" dialect eg. `parse_one(sql, read="spark")`
- Make sure you specify the "write" dialect eg. `ast.sql(dialect="duckdb")`
- Check if the issue still exists on main

**Fully reproducible code snippet**
```python
from sqlglot import parse_one
parse_one("json_query(content, 'lax $.HY.*' with unconditional wrapper)",dialect="trino")
```

**Official Documentation**
Please include links to official SQL documentation related to your issue.
https://trino.io/docs/current/functions/json.html#json-query

Discussion:

Solution patch:
diff --git a/sqlglot/dialects/presto.py b/sqlglot/dialects/presto.py
index e837052563..6cce854bb0 100644
--- a/sqlglot/dialects/presto.py
+++ b/sqlglot/dialects/presto.py
@@ -166,35 +166,6 @@ def _unix_to_time_sql(self: Presto.Generator, expression: exp.UnixToTime) -> str
     return f"FROM_UNIXTIME(CAST({timestamp} AS DOUBLE) / POW(10, {scale}))"
 
 
-def _jsonextract_sql(self: Presto.Generator, expression: exp.JSONExtract) -> str:
-    is_json_extract = self.dialect.settings.get("variant_extract_is_json_extract", True)
-
-    # Generate JSON_EXTRACT unless the user has configured that a Snowflake / Databricks
-    # VARIANT extract (e.g. col:x.y) should map to dot notation (i.e ROW access) in Presto/Trino
-    if not expression.args.get("variant_extract") or is_json_extract:
-        return self.func(
-            "JSON_EXTRACT", expression.this, expression.expression, *expression.expressions
-        )
-
-    this = self.sql(expression, "this")
-
-    # Convert the JSONPath extraction `JSON_EXTRACT(col, '$.x.y) to a ROW access col.x.y
-    segments = []
-    for path_key in expression.expression.expressions[1:]:
-        if not isinstance(path_key, exp.JSONPathKey):
-            # Cannot transpile subscripts, wildcards etc to dot notation
-            self.unsupported(f"Cannot transpile JSONPath segment '{path_key}' to ROW access")
-            continue
-        key = path_key.this
-        if not exp.SAFE_IDENTIFIER_RE.match(key):
-            key = f'"{key}"'
-        segments.append(f".{key}")
-
-    expr = "".join(segments)
-
-    return f"{this}{expr}"
-
-
 def _to_int(self: Presto.Generator, expression: exp.Expression) -> exp.Expression:
     if not expression.type:
         from sqlglot.optimizer.annotate_types import annotate_types
@@ -435,7 +406,7 @@ class Generator(generator.Generator):
             exp.If: if_sql(),
             exp.ILike: no_ilike_sql,
             exp.Initcap: _initcap_sql,
-            exp.JSONExtract: _jsonextract_sql,
+            exp.JSONExtract: lambda self, e: self.jsonextract_sql(e),
             exp.Last: _first_last_sql,
             exp.LastValue: _first_last_sql,
             exp.LastDay: lambda self, e: self.func("LAST_DAY_OF_MONTH", e.this),
@@ -693,3 +664,33 @@ def delete_sql(self, expression: exp.Delete) -> str:
                     expression = t.cast(exp.Delete, expression.transform(unqualify_columns))
 
             return super().delete_sql(expression)
+
+        def jsonextract_sql(self, expression: exp.JSONExtract) -> str:
+            is_json_extract = self.dialect.settings.get("variant_extract_is_json_extract", True)
+
+            # Generate JSON_EXTRACT unless the user has configured that a Snowflake / Databricks
+            # VARIANT extract (e.g. col:x.y) should map to dot notation (i.e ROW access) in Presto/Trino
+            if not expression.args.get("variant_extract") or is_json_extract:
+                return self.func(
+                    "JSON_EXTRACT", expression.this, expression.expression, *expression.expressions
+                )
+
+            this = self.sql(expression, "this")
+
+            # Convert the JSONPath extraction `JSON_EXTRACT(col, '$.x.y) to a ROW access col.x.y
+            segments = []
+            for path_key in expression.expression.expressions[1:]:
+                if not isinstance(path_key, exp.JSONPathKey):
+                    # Cannot transpile subscripts, wildcards etc to dot notation
+                    self.unsupported(
+                        f"Cannot transpile JSONPath segment '{path_key}' to ROW access"
+                    )
+                    continue
+                key = path_key.this
+                if not exp.SAFE_IDENTIFIER_RE.match(key):
+                    key = f'"{key}"'
+                segments.append(f".{key}")
+
+            expr = "".join(segments)
+
+            return f"{this}{expr}"
diff --git a/sqlglot/dialects/trino.py b/sqlglot/dialects/trino.py
index ccdbb710d7..290fc51c6b 100644
--- a/sqlglot/dialects/trino.py
+++ b/sqlglot/dialects/trino.py
@@ -1,8 +1,9 @@
 from __future__ import annotations
 
-from sqlglot import exp
+from sqlglot import exp, parser
 from sqlglot.dialects.dialect import merge_without_target_sql, trim_sql, timestrtotime_sql
 from sqlglot.dialects.presto import Presto
+from sqlglot.tokens import TokenType
 
 
 class Trino(Presto):
@@ -13,8 +14,30 @@ class Parser(Presto.Parser):
         FUNCTION_PARSERS = {
             **Presto.Parser.FUNCTION_PARSERS,
             "TRIM": lambda self: self._parse_trim(),
+            "JSON_QUERY": lambda self: self._parse_json_query(),
         }
 
+        JSON_QUERY_OPTIONS: parser.OPTIONS_TYPE = {
+            **dict.fromkeys(
+                ("WITH", "WITHOUT"),
+                (
+                    ("CONDITIONAL", "WRAPPER"),
+                    ("CONDITIONAL", "ARRAY", "WRAPPED"),
+                    ("UNCONDITIONAL", "WRAPPER"),
+                    ("UNCONDITIONAL", "ARRAY", "WRAPPER"),
+                ),
+            ),
+        }
+
+        def _parse_json_query(self):
+            return self.expression(
+                exp.JSONExtract,
+                this=self._parse_bitwise(),
+                expression=self._match(TokenType.COMMA) and self._parse_bitwise(),
+                option=self._parse_var_from_options(self.JSON_QUERY_OPTIONS, raise_unmatched=False),
+                json_query=True,
+            )
+
     class Generator(Presto.Generator):
         TRANSFORMS = {
             **Presto.Generator.TRANSFORMS,
@@ -23,6 +46,7 @@ class Generator(Presto.Generator):
             exp.Merge: merge_without_target_sql,
             exp.TimeStrToTime: lambda self, e: timestrtotime_sql(self, e, include_precision=True),
             exp.Trim: trim_sql,
+            exp.JSONExtract: lambda self, e: self.jsonextract_sql(e),
         }
 
         SUPPORTED_JSON_PATH_PARTS = {
@@ -31,5 +55,15 @@ class Generator(Presto.Generator):
             exp.JSONPathSubscript,
         }
 
+        def jsonextract_sql(self, expression: exp.JSONExtract) -> str:
+            if not expression.args.get("json_query"):
+                return super().jsonextract_sql(expression)
+
+            json_path = self.sql(expression, "expression")
+            option = self.sql(expression, "option")
+            option = f" {option}" if option else ""
+
+            return self.func("JSON_QUERY", expression.this, json_path + option)
+
     class Tokenizer(Presto.Tokenizer):
         HEX_STRINGS = [("X'", "'")]
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index f77c5d97eb..59de679f31 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -5738,6 +5738,8 @@ class JSONExtract(Binary, Func):
         "only_json_types": False,
         "expressions": False,
         "variant_extract": False,
+        "json_query": False,
+        "option": False,
     }
     _sql_names = ["JSON_EXTRACT"]
     is_var_len_args = True


Test patch:
diff --git a/tests/dialects/test_trino.py b/tests/dialects/test_trino.py
index 0ebe7491fd..8c73ec1c84 100644
--- a/tests/dialects/test_trino.py
+++ b/tests/dialects/test_trino.py
@@ -4,6 +4,12 @@
 class TestTrino(Validator):
     dialect = "trino"
 
+    def test_trino(self):
+        self.validate_identity("JSON_EXTRACT(content, json_path)")
+        self.validate_identity("JSON_QUERY(content, 'lax $.HY.*')")
+        self.validate_identity("JSON_QUERY(content, 'strict $.HY.*' WITH UNCONDITIONAL WRAPPER)")
+        self.validate_identity("JSON_QUERY(content, 'strict $.HY.*' WITHOUT CONDITIONAL WRAPPER)")
+
     def test_trim(self):
         self.validate_identity("SELECT TRIM('!' FROM '!foo!')")
         self.validate_identity("SELECT TRIM(BOTH '$' FROM '$var$')")