Fix(bigquery): improve support for cast to timestamp with format, tim…

…e zone (#1873)

* Fix(bigquery): improve support for cast to timestamp with format, time zone

* Cleanup

* Cleanup

* Cleanup

sqlglot/dialects/bigquery.py

@@ -163,6 +163,12 @@ def _pushdown_cte_column_names(expression: exp.Expression) -> exp.Expression:
     return expression
 
 
+def _parse_timestamp(args: t.List) -> exp.StrToTime:
+    this = format_time_lambda(exp.StrToTime, "bigquery")([seq_get(args, 1), seq_get(args, 0)])
+    this.set("zone", seq_get(args, 2))
+    return this
+
+
 class BigQuery(Dialect):
     UNNEST_COLUMN_ONLY = True
 
@@ -266,9 +272,7 @@ class Parser(parser.Parser):
             "PARSE_DATE": lambda args: format_time_lambda(exp.StrToDate, "bigquery")(
                 [seq_get(args, 1), seq_get(args, 0)]
             ),
-            "PARSE_TIMESTAMP": lambda args: format_time_lambda(exp.StrToTime, "bigquery")(
-                [seq_get(args, 1), seq_get(args, 0)]
-            ),
+            "PARSE_TIMESTAMP": _parse_timestamp,
             "REGEXP_CONTAINS": exp.RegexpLike.from_arg_list,
             "REGEXP_EXTRACT": lambda args: exp.RegexpExtract(
                 this=seq_get(args, 0),
@@ -400,7 +404,9 @@ class Generator(generator.Generator):
                 ]
             ),
             exp.StrToDate: lambda self, e: f"PARSE_DATE({self.format_time(e)}, {self.sql(e, 'this')})",
-            exp.StrToTime: lambda self, e: f"PARSE_TIMESTAMP({self.format_time(e)}, {self.sql(e, 'this')})",
+            exp.StrToTime: lambda self, e: self.func(
+                "PARSE_TIMESTAMP", self.format_time(e), e.this, e.args.get("zone")
+            ),
             exp.TimeAdd: _date_add_sql("TIME", "ADD"),
             exp.TimeSub: _date_add_sql("TIME", "SUB"),
             exp.TimestampAdd: _date_add_sql("TIMESTAMP", "ADD"),
@@ -551,10 +557,15 @@ class Generator(generator.Generator):
         }
 
         def attimezone_sql(self, expression: exp.AtTimeZone) -> str:
-            if not isinstance(expression.parent, exp.Cast):
+            parent = expression.parent
+
+            # BigQuery allows CAST(.. AS {STRING|TIMESTAMP} [FORMAT <fmt> [AT TIME ZONE <tz>]]).
+            # Only the TIMESTAMP one should use the below conversion, when AT TIME ZONE is included.
+            if not isinstance(parent, exp.Cast) or not parent.to.is_type("text"):
                 return self.func(
                     "TIMESTAMP", self.func("DATETIME", expression.this, expression.args.get("zone"))
                 )
+
             return super().attimezone_sql(expression)
 
         def trycast_sql(self, expression: exp.TryCast) -> str:

sqlglot/expressions.py

@@ -4482,7 +4482,7 @@ class StrToDate(Func):
 
 
 class StrToTime(Func):
-    arg_types = {"this": True, "format": True}
+    arg_types = {"this": True, "format": True, "zone": False}
 
 
 # Spark allows unix_timestamp()

sqlglot/parser.py

@@ -3710,21 +3710,27 @@ def _parse_cast(self, strict: bool) -> exp.Expression:
             if self._match(TokenType.CHARACTER_SET):
                 to = self.expression(exp.CharacterSet, this=self._parse_var_or_string())
         elif self._match(TokenType.FORMAT):
-            fmt = self._parse_at_time_zone(self._parse_string())
+            fmt_string = self._parse_string()
+            fmt = self._parse_at_time_zone(fmt_string)
 
             if to.this in exp.DataType.TEMPORAL_TYPES:
-                return self.expression(
+                this = self.expression(
                     exp.StrToDate if to.this == exp.DataType.Type.DATE else exp.StrToTime,
                     this=this,
                     format=exp.Literal.string(
                         format_time(
-                            fmt.this if fmt else "",
+                            fmt_string.this if fmt_string else "",
                             self.FORMAT_MAPPING or self.TIME_MAPPING,
                             self.FORMAT_TRIE or self.TIME_TRIE,
                         )
                     ),
                 )
 
+                if isinstance(fmt, exp.AtTimeZone) and isinstance(this, exp.StrToTime):
+                    this.set("zone", fmt.args["zone"])
+
+                return this
+
         return self.expression(exp.Cast if strict else exp.TryCast, this=this, to=to, format=fmt)
 
     def _parse_concat(self) -> t.Optional[exp.Expression]:

tests/dialects/test_bigquery.py

@@ -29,6 +29,7 @@ def test_bigquery(self):
         with self.assertRaises(ParseError):
             transpile("SELECT * FROM UNNEST(x) AS x(y)", read="bigquery")
 
+        self.validate_identity("SELECT PARSE_TIMESTAMP('%c', 'Thu Dec 25 07:30:00 2008', 'UTC')")
         self.validate_identity("SELECT ANY_VALUE(fruit HAVING MAX sold) FROM fruits")
         self.validate_identity("SELECT ANY_VALUE(fruit HAVING MIN sold) FROM fruits")
         self.validate_identity("SELECT `project-id`.udfs.func(call.dir)")
@@ -105,6 +106,14 @@ def test_bigquery(self):
         self.validate_all("CAST(x AS NVARCHAR)", write={"bigquery": "CAST(x AS STRING)"})
         self.validate_all("CAST(x AS TIMESTAMPTZ)", write={"bigquery": "CAST(x AS TIMESTAMP)"})
         self.validate_all("CAST(x AS RECORD)", write={"bigquery": "CAST(x AS STRUCT)"})
+        self.validate_all(
+            "SELECT CAST('20201225' AS TIMESTAMP FORMAT 'YYYYMMDD' AT TIME ZONE 'America/New_York')",
+            write={"bigquery": "SELECT PARSE_TIMESTAMP('%Y%m%d', '20201225', 'America/New_York')"},
+        )
+        self.validate_all(
+            "SELECT CAST('20201225' AS TIMESTAMP FORMAT 'YYYYMMDD')",
+            write={"bigquery": "SELECT PARSE_TIMESTAMP('%Y%m%d', '20201225')"},
+        )
         self.validate_all(
             "SELECT CAST(TIMESTAMP '2008-12-25 00:00:00+00:00' AS STRING FORMAT 'YYYY-MM-DD HH24:MI:SS TZH:TZM') AS date_time_to_string",
             write={