feat!: timestamp diff for mysql and databricks

tobymao · Jan 31, 2024 · 617a8c0 · 617a8c0
1 parent fb450f0
commit 617a8c0
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 1 deletion.
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
@@ -611,6 +611,7 @@ class Generator(generator.Generator):
             exp.TimeFromParts: rename_func("TIME"),
             exp.TimeSub: date_add_interval_sql("TIME", "SUB"),
             exp.TimestampAdd: date_add_interval_sql("TIMESTAMP", "ADD"),
+            exp.TimestampDiff: rename_func("TIMESTAMP_DIFF"),
             exp.TimestampSub: date_add_interval_sql("TIMESTAMP", "SUB"),
             exp.TimeStrToTime: timestrtotime_sql,
             exp.Trim: lambda self, e: self.func(f"TRIM", e.this, e.expression),

diff --git a/sqlglot/dialects/databricks.py b/sqlglot/dialects/databricks.py
@@ -22,6 +22,7 @@ class Parser(Spark.Parser):
             "DATEADD": parse_date_delta(exp.DateAdd),
             "DATE_ADD": parse_date_delta(exp.DateAdd),
             "DATEDIFF": parse_date_delta(exp.DateDiff),
+            "TIMESTAMPDIFF": parse_date_delta(exp.TimestampDiff),
         }
 
         FACTOR = {
@@ -48,6 +49,9 @@ class Generator(Spark.Generator):
             exp.DatetimeDiff: lambda self, e: self.func(
                 "TIMESTAMPDIFF", e.text("unit"), e.expression, e.this
             ),
+            exp.TimestampDiff: 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/sqlglot/dialects/mysql.py b/sqlglot/dialects/mysql.py
@@ -19,6 +19,7 @@
     no_pivot_sql,
     no_tablesample_sql,
     no_trycast_sql,
+    parse_date_delta,
     parse_date_delta_with_interval,
     path_to_jsonpath,
     rename_func,
@@ -306,6 +307,7 @@ class Parser(parser.Parser):
                 format=exp.Literal.string("%B"),
             ),
             "STR_TO_DATE": _str_to_date,
+            "TIMESTAMPDIFF": parse_date_delta(exp.TimestampDiff),
             "TO_DAYS": lambda args: exp.paren(
                 exp.DateDiff(
                     this=exp.TsOrDsToDate(this=seq_get(args, 0)),
@@ -673,6 +675,9 @@ class Generator(generator.Generator):
             exp.TableSample: no_tablesample_sql,
             exp.TimeFromParts: rename_func("MAKETIME"),
             exp.TimestampAdd: date_add_interval_sql("DATE", "ADD"),
+            exp.TimestampDiff: lambda self, e: self.func(
+                "TIMESTAMPDIFF", e.text("unit"), e.expression, e.this
+            ),
             exp.TimestampSub: date_add_interval_sql("DATE", "SUB"),
             exp.TimeStrToUnix: rename_func("UNIX_TIMESTAMP"),
             exp.TimeStrToTime: lambda self, e: self.sql(exp.cast(e.this, "datetime", copy=True)),

diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -4731,6 +4731,7 @@ class TimestampSub(Func, TimeUnit):
 
 
 class TimestampDiff(Func, TimeUnit):
+    _sql_names = ["TIMESTAMPDIFF", "TIMESTAMP_DIFF"]
     arg_types = {"this": True, "expression": True, "unit": False}
 
 

diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
@@ -416,10 +416,25 @@ def test_bigquery(self):
             "SELECT TIMESTAMP_DIFF(TIMESTAMP_SECONDS(60), TIMESTAMP_SECONDS(0), minute)",
             write={
                 "bigquery": "SELECT TIMESTAMP_DIFF(TIMESTAMP_SECONDS(60), TIMESTAMP_SECONDS(0), MINUTE)",
+                "databricks": "SELECT TIMESTAMPDIFF(MINUTE, CAST(FROM_UNIXTIME(0) AS TIMESTAMP), CAST(FROM_UNIXTIME(60) AS TIMESTAMP))",
                 "duckdb": "SELECT DATE_DIFF('MINUTE', TO_TIMESTAMP(0), TO_TIMESTAMP(60))",
                 "snowflake": "SELECT TIMESTAMPDIFF(MINUTE, TO_TIMESTAMP(0), TO_TIMESTAMP(60))",
             },
         )
+        self.validate_all(
+            "TIMESTAMP_DIFF(a, b, MONTH)",
+            read={
+                "bigquery": "TIMESTAMP_DIFF(a, b, month)",
+                "databricks": "TIMESTAMPDIFF(month, b, a)",
+                "mysql": "TIMESTAMPDIFF(month, b, a)",
+            },
+            write={
+                "databricks": "TIMESTAMPDIFF(MONTH, b, a)",
+                "mysql": "TIMESTAMPDIFF(MONTH, b, a)",
+                "snowflake": "TIMESTAMPDIFF(MONTH, b, a)",
+            },
+        )
+
         self.validate_all(
             "SELECT TIMESTAMP_MICROS(x)",
             read={

diff --git a/tests/fixtures/identity.sql b/tests/fixtures/identity.sql
@@ -158,7 +158,7 @@ CAST(x AS UUID)
 FILTER(a, x -> x.a.b.c.d.e.f.g)
 FILTER(a, x -> FOO(x.a.b.c.d.e.f.g) + x.a.b.c.d.e.f.g)
 TIMESTAMP_FROM_PARTS(2019, 1, 10, 2, 3, 4, 123456789, 'America/Los_Angeles')
-TIMESTAMP_DIFF(CURRENT_TIMESTAMP(), 1, DAY)
+TIMESTAMPDIFF(CURRENT_TIMESTAMP(), 1, DAY)
 DATETIME_DIFF(CURRENT_DATE, 1, DAY)
 QUANTILE(x, 0.5)
 REGEXP_REPLACE('new york', '(\w)(\w*)', x -> UPPER(x[1]) || LOWER(x[2]))