From 30e3e2fe7c77d122bac02fd69ddf1b344d414eb3 Mon Sep 17 00:00:00 2001 From: Filipe Reis <20783380+fireis@users.noreply.github.com> Date: Fri, 1 Sep 2023 12:45:35 -0300 Subject: [PATCH] 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 --- sqlglot/dialects/databricks.py | 15 ++++++++++++++- tests/dialects/test_bigquery.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/sqlglot/dialects/databricks.py b/sqlglot/dialects/databricks.py index 2149aca73d..6ec04873a3 100644 --- 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 index 000725d914..c0ebf5ae16 100644 --- 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)"})