Feat: add support for transpiling some datetime functions from bq to …

…databricks (#2142) * add bq to databricks support for datetime functions * Update sqlglot/dialects/databricks.py * Update sqlglot/dialects/databricks.py * Update sqlglot/dialects/databricks.py --------- Co-authored-by: Toby Mao <[email protected]>
tobymao · Sep 1, 2023 · 30e3e2f · 30e3e2f
1 parent 39bce6d
commit 30e3e2f
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 1 deletion.
diff --git a/sqlglot/dialects/databricks.py b/sqlglot/dialects/databricks.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from sqlglot import exp, transforms
-from sqlglot.dialects.dialect import parse_date_delta
+from sqlglot.dialects.dialect import parse_date_delta, timestamptrunc_sql
 from sqlglot.dialects.spark import Spark
 from sqlglot.dialects.tsql import generate_date_delta_with_unit_sql
 from sqlglot.tokens import TokenType
@@ -28,6 +28,19 @@ class Generator(Spark.Generator):
             **Spark.Generator.TRANSFORMS,
             exp.DateAdd: generate_date_delta_with_unit_sql,
             exp.DateDiff: generate_date_delta_with_unit_sql,
+            exp.DatetimeAdd: lambda self, e: self.func(
+                "TIMESTAMPADD", e.text("unit"), e.expression, e.this
+            ),
+            exp.DatetimeSub: lambda self, e: self.func(
+                "TIMESTAMPADD",
+                e.text("unit"),
+                exp.Mul(this=e.expression.copy(), expression=exp.Literal.number(-1)),
+                e.this,
+            ),
+            exp.DatetimeDiff: lambda self, e: self.func(
+                "TIMESTAMPDIFF", e.text("unit"), e.expression, e.this
+            ),
+            exp.DatetimeTrunc: timestamptrunc_sql,
             exp.JSONExtract: lambda self, e: self.binary(e, ":"),
             exp.Select: transforms.preprocess(
                 [

diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
@@ -132,6 +132,34 @@ def test_bigquery(self):
         self.validate_all('x <> """"""', write={"bigquery": "x <> ''"})
         self.validate_all("x <> ''''''", write={"bigquery": "x <> ''"})
         self.validate_all("CAST(x AS DATETIME)", read={"": "x::timestamp"})
+        self.validate_all(
+            "SELECT DATETIME_DIFF('2023-01-01T00:00:00', '2023-01-01T05:00:00', MILLISECOND)",
+            write={
+                "bigquery": "SELECT DATETIME_DIFF('2023-01-01T00:00:00', '2023-01-01T05:00:00', MILLISECOND)",
+                "databricks": "SELECT TIMESTAMPDIFF(MILLISECOND, '2023-01-01T05:00:00', '2023-01-01T00:00:00')",
+            },
+        ),
+        self.validate_all(
+            "SELECT DATETIME_ADD('2023-01-01T00:00:00', INTERVAL 1 MILLISECOND)",
+            write={
+                "bigquery": "SELECT DATETIME_ADD('2023-01-01T00:00:00', INTERVAL 1 MILLISECOND)",
+                "databricks": "SELECT TIMESTAMPADD(MILLISECOND, 1, '2023-01-01T00:00:00')",
+            },
+        ),
+        self.validate_all(
+            "SELECT DATETIME_SUB('2023-01-01T00:00:00', INTERVAL 1 MILLISECOND)",
+            write={
+                "bigquery": "SELECT DATETIME_SUB('2023-01-01T00:00:00', INTERVAL 1 MILLISECOND)",
+                "databricks": "SELECT TIMESTAMPADD(MILLISECOND, 1 * -1, '2023-01-01T00:00:00')",
+            },
+        ),
+        self.validate_all(
+            "SELECT DATETIME_TRUNC('2023-01-01T01:01:01', HOUR)",
+            write={
+                "bigquery": "SELECT DATETIME_TRUNC('2023-01-01T01:01:01', HOUR)",
+                "databricks": "SELECT DATE_TRUNC('HOUR', '2023-01-01T01:01:01')",
+            },
+        ),
         self.validate_all("LEAST(x, y)", read={"sqlite": "MIN(x, y)"})
         self.validate_all("CAST(x AS CHAR)", write={"bigquery": "CAST(x AS STRING)"})
         self.validate_all("CAST(x AS NCHAR)", write={"bigquery": "CAST(x AS STRING)"})