fix(bigquery): Parse timezone for DATE_TRUNC (#4347)

tobymao · Nov 5, 2024 · 1689dc7 · 1689dc7
1 parent 8e16abe
commit 1689dc7
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
@@ -427,6 +427,7 @@ class Parser(parser.Parser):
             "DATE_TRUNC": lambda args: exp.DateTrunc(
                 unit=exp.Literal.string(str(seq_get(args, 1))),
                 this=seq_get(args, 0),
+                zone=seq_get(args, 2),
             ),
             "DATETIME": _build_datetime,
             "DATETIME_ADD": build_date_delta_with_interval(exp.DatetimeAdd),
@@ -732,7 +733,9 @@ class Generator(generator.Generator):
             exp.DateSub: date_add_interval_sql("DATE", "SUB"),
             exp.DatetimeAdd: date_add_interval_sql("DATETIME", "ADD"),
             exp.DatetimeSub: date_add_interval_sql("DATETIME", "SUB"),
-            exp.DateTrunc: lambda self, e: self.func("DATE_TRUNC", e.this, e.text("unit")),
+            exp.DateTrunc: lambda self, e: self.func(
+                "DATE_TRUNC", e.this, e.text("unit"), e.args.get("zone")
+            ),
             exp.FromTimeZone: lambda self, e: self.func(
                 "DATETIME", self.func("TIMESTAMP", e.this, e.args.get("zone")), "'UTC'"
             ),

diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
@@ -165,6 +165,7 @@ def test_bigquery(self):
         self.validate_identity("SELECT * FROM foo.bar.25ab c", "SELECT * FROM foo.bar.`25ab` AS c")
         self.validate_identity("x <> ''")
         self.validate_identity("DATE_TRUNC(col, WEEK(MONDAY))")
+        self.validate_identity("DATE_TRUNC(col, MONTH, 'UTC+8')")
         self.validate_identity("SELECT b'abc'")
         self.validate_identity("SELECT AS STRUCT 1 AS a, 2 AS b")
         self.validate_identity("SELECT DISTINCT AS STRUCT 1 AS a, 2 AS b")