Repository: tobymao/sqlglot

Issue:
Unable to parse `KEEP QUOTES`/`OMIT QUOTES` in Trino
Found an [issue](https://github.com/apache/superset/issues/31768) when trying to parse a Trino query that uses `JSON_QUERY(col, jsonpath OMIT QUOTES)`:

```python
import sqlglot

# https://trino.io/docs/current/functions/json.html#id6
sql = """
SELECT
      id,
      json_query(description, 'strict $.comment' KEEP QUOTES) AS quoted_comment,
      json_query(description, 'strict $.comment' OMIT QUOTES) AS unquoted_comment
FROM customers
"""

sqlglot.parse_one(sql, read="trino")
```

Will raise:

```python
...
  File "/Users/beto/Projects/sqlglot/sqlglot/parser.py", line 5817, in _parse_function_call
    self._match_r_paren(this)
  File "/Users/beto/Projects/sqlglot/sqlglot/parser.py", line 7930, in _match_r_paren
    self.raise_error("Expecting )")
  File "/Users/beto/Projects/sqlglot/sqlglot/parser.py", line 1662, in raise_error
    raise error
sqlglot.errors.ParseError: Expecting ). Line 4, Col: 48.
```

Discussion:

Patch:
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index ac699944..0601a367 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -67,8 +67,9 @@ def _build_datetime(
                     expr.set("safe", safe)
                     return expr
 
-        if kind == exp.DataType.Type.DATE and not int_value:
-            formatted_exp = build_formatted_time(exp.TsOrDsToDate, "snowflake")(args)
+        if kind in (exp.DataType.Type.DATE, exp.DataType.Type.TIME) and not int_value:
+            klass = exp.TsOrDsToDate if kind == exp.DataType.Type.DATE else exp.TsOrDsToTime
+            formatted_exp = build_formatted_time(klass, "snowflake")(args)
             formatted_exp.set("safe", safe)
             return formatted_exp
 
@@ -449,6 +450,7 @@ class Snowflake(Dialect):
             "TIMESTAMP_NTZ_FROM_PARTS": build_timestamp_from_parts,
             "TRY_PARSE_JSON": lambda args: exp.ParseJSON(this=seq_get(args, 0), safe=True),
             "TRY_TO_DATE": _build_datetime("TRY_TO_DATE", exp.DataType.Type.DATE, safe=True),
+            "TRY_TO_TIME": _build_datetime("TRY_TO_TIME", exp.DataType.Type.TIME, safe=True),
             "TRY_TO_TIMESTAMP": _build_datetime(
                 "TRY_TO_TIMESTAMP", exp.DataType.Type.TIMESTAMP, safe=True
             ),
@@ -979,6 +981,9 @@ class Snowflake(Dialect):
             exp.TsOrDsToDate: lambda self, e: self.func(
                 "TRY_TO_DATE" if e.args.get("safe") else "TO_DATE", e.this, self.format_time(e)
             ),
+            exp.TsOrDsToTime: lambda self, e: self.func(
+                "TRY_TO_TIME" if e.args.get("safe") else "TO_TIME", e.this, self.format_time(e)
+            ),
             exp.UnixToTime: rename_func("TO_TIMESTAMP"),
             exp.Uuid: rename_func("UUID_STRING"),
             exp.VarMap: lambda self, e: var_map_sql(self, e, "OBJECT_CONSTRUCT"),
diff --git a/sqlglot/dialects/trino.py b/sqlglot/dialects/trino.py
index adfaa5a5..b8c3201b 100644
--- a/sqlglot/dialects/trino.py
+++ b/sqlglot/dialects/trino.py
@@ -4,6 +4,7 @@ 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
+import typing as t
 
 
 class Trino(Presto):
@@ -33,13 +34,26 @@ class Trino(Presto):
             ),
         }
 
-        def _parse_json_query(self):
+        def _parse_json_query_quote(self) -> t.Optional[exp.JSONExtractQuote]:
+            if not (
+                self._match_text_seq("KEEP", "QUOTES") or self._match_text_seq("OMIT", "QUOTES")
+            ):
+                return None
+
+            return self.expression(
+                exp.JSONExtractQuote,
+                option=self._tokens[self._index - 2].text.upper(),
+                scalar=self._match_text_seq("ON", "SCALAR", "STRING"),
+            )
+
+        def _parse_json_query(self) -> exp.JSONExtract:
             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,
+                quote=self._parse_json_query_quote(),
             )
 
     class Generator(Presto.Generator):
@@ -68,7 +82,10 @@ class Trino(Presto):
             option = self.sql(expression, "option")
             option = f" {option}" if option else ""
 
-            return self.func("JSON_QUERY", expression.this, json_path + option)
+            quote = self.sql(expression, "quote")
+            quote = f" {quote}" if quote else ""
+
+            return self.func("JSON_QUERY", expression.this, json_path + option + quote)
 
         def groupconcat_sql(self, expression: exp.GroupConcat) -> str:
             this = expression.this
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index f52f4147..7cba5b4d 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -6156,6 +6156,7 @@ class JSONExtract(Binary, Func):
         "variant_extract": False,
         "json_query": False,
         "option": False,
+        "quote": False,
     }
     _sql_names = ["JSON_EXTRACT"]
     is_var_len_args = True
@@ -6165,6 +6166,14 @@ class JSONExtract(Binary, Func):
         return self.expression.output_name if not self.expressions else ""
 
 
+# https://trino.io/docs/current/functions/json.html#json-query
+class JSONExtractQuote(Expression):
+    arg_types = {
+        "option": True,
+        "scalar": False,
+    }
+
+
 class JSONExtractArray(Func):
     arg_types = {"this": True, "expression": False}
     _sql_names = ["JSON_EXTRACT_ARRAY"]
@@ -6654,7 +6663,7 @@ class TsOrDsToDatetime(Func):
 
 
 class TsOrDsToTime(Func):
-    pass
+    arg_types = {"this": True, "format": False, "safe": False}
 
 
 class TsOrDsToTimestamp(Func):
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 3ceee215..d7d52b55 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -3934,6 +3934,16 @@ class Generator(metaclass=_Generator):
 
     def tsordstotime_sql(self, expression: exp.TsOrDsToTime) -> str:
         this = expression.this
+        time_format = self.format_time(expression)
+
+        if time_format:
+            return self.sql(
+                exp.cast(
+                    exp.StrToTime(this=this, format=expression.args["format"]),
+                    exp.DataType.Type.TIME,
+                )
+            )
+
         if isinstance(this, exp.TsOrDsToTime) or this.is_type(exp.DataType.Type.TIME):
             return self.sql(this)
 
@@ -4421,6 +4431,10 @@ class Generator(metaclass=_Generator):
 
         return f"{empty}{error}{null}"
 
+    def jsonextractquote_sql(self, expression: exp.JSONExtractQuote) -> str:
+        scalar = " ON SCALAR STRING" if expression.args.get("scalar") else ""
+        return f"{self.sql(expression, 'option')} QUOTES{scalar}"
+
     def jsonexists_sql(self, expression: exp.JSONExists) -> str:
         this = self.sql(expression, "this")
         path = self.sql(expression, "path")


Test patch:
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index e1f43882..cb38af75 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -36,7 +36,6 @@ class TestSnowflake(Validator):
         self.validate_identity("SELECT CAST(obj AS OBJECT(x CHAR) RENAME FIELDS)")
         self.validate_identity("SELECT CAST(obj AS OBJECT(x CHAR, y VARCHAR) ADD FIELDS)")
         self.validate_identity("SELECT TO_TIMESTAMP(123.4)").selects[0].assert_is(exp.Anonymous)
-        self.validate_identity("SELECT TO_TIME(x) FROM t")
         self.validate_identity("SELECT TO_TIMESTAMP(x) FROM t")
         self.validate_identity("SELECT TO_TIMESTAMP_NTZ(x) FROM t")
         self.validate_identity("SELECT TO_TIMESTAMP_LTZ(x) FROM t")
@@ -730,13 +729,6 @@ class TestSnowflake(Validator):
                 "spark": "SELECT CAST('2013-04-05 01:02:03' AS TIMESTAMP)",
             },
         )
-        self.validate_all(
-            "SELECT TO_TIME('12:05:00')",
-            write={
-                "bigquery": "SELECT CAST('12:05:00' AS TIME)",
-                "snowflake": "SELECT CAST('12:05:00' AS TIME)",
-            },
-        )
         self.validate_all(
             "SELECT TO_TIMESTAMP('04/05/2013 01:02:03', 'mm/DD/yyyy hh24:mi:ss')",
             read={
@@ -1296,6 +1288,37 @@ class TestSnowflake(Validator):
             },
         )
 
+        self.validate_identity("SELECT TO_TIME(x) FROM t")
+        self.validate_all(
+            "SELECT TO_TIME('12:05:00')",
+            write={
+                "bigquery": "SELECT CAST('12:05:00' AS TIME)",
+                "snowflake": "SELECT CAST('12:05:00' AS TIME)",
+                "duckdb": "SELECT CAST('12:05:00' AS TIME)",
+            },
+        )
+        self.validate_all(
+            "SELECT TO_TIME(CONVERT_TIMEZONE('UTC', 'US/Pacific', '2024-08-06 09:10:00.000')) AS pst_time",
+            write={
+                "snowflake": "SELECT TO_TIME(CONVERT_TIMEZONE('UTC', 'US/Pacific', '2024-08-06 09:10:00.000')) AS pst_time",
+                "duckdb": "SELECT CAST(CAST('2024-08-06 09:10:00.000' AS TIMESTAMP) AT TIME ZONE 'UTC' AT TIME ZONE 'US/Pacific' AS TIME) AS pst_time",
+            },
+        )
+        self.validate_all(
+            "SELECT TO_TIME('11.15.00', 'hh24.mi.ss')",
+            write={
+                "snowflake": "SELECT TO_TIME('11.15.00', 'hh24.mi.ss')",
+                "duckdb": "SELECT CAST(STRPTIME('11.15.00', '%H.%M.%S') AS TIME)",
+            },
+        )
+        self.validate_all(
+            "SELECT TRY_TO_TIME('11.15.00', 'hh24.mi.ss')",
+            write={
+                "snowflake": "SELECT TRY_TO_TIME('11.15.00', 'hh24.mi.ss')",
+                "duckdb": "SELECT CAST(STRPTIME('11.15.00', '%H.%M.%S') AS TIME)",
+            },
+        )
+
     def test_semi_structured_types(self):
         self.validate_identity("SELECT CAST(a AS VARIANT)")
         self.validate_identity("SELECT CAST(a AS ARRAY)")
diff --git a/tests/dialects/test_trino.py b/tests/dialects/test_trino.py
index 3b22f0dc..44ff1ea0 100644
--- a/tests/dialects/test_trino.py
+++ b/tests/dialects/test_trino.py
@@ -9,6 +9,13 @@ class TestTrino(Validator):
         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)")
+        self.validate_identity("JSON_QUERY(description, 'strict $.comment' KEEP QUOTES)")
+        self.validate_identity(
+            "JSON_QUERY(description, 'strict $.comment' OMIT QUOTES ON SCALAR STRING)"
+        )
+        self.validate_identity(
+            "JSON_QUERY(content, 'strict $.HY.*' WITH UNCONDITIONAL WRAPPER KEEP QUOTES)"
+        )
 
     def test_listagg(self):
         self.validate_identity(