Fix: clickhouse dateadd/datediff closes #2108

tobymao · Aug 23, 2023 · f62f35c · f62f35c
1 parent bd531b2
commit f62f35c
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 0 deletions.
diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
@@ -11,6 +11,7 @@
     var_map_sql,
 )
 from sqlglot.errors import ParseError
+from sqlglot.helper import seq_get
 from sqlglot.parser import parse_var_map
 from sqlglot.tokens import Token, TokenType
 
@@ -68,6 +69,18 @@ class Parser(parser.Parser):
         FUNCTIONS = {
             **parser.Parser.FUNCTIONS,
             "ANY": exp.AnyValue.from_arg_list,
+            "DATE_ADD": lambda args: exp.DateAdd(
+                this=seq_get(args, 2), expression=seq_get(args, 1), unit=seq_get(args, 0)
+            ),
+            "DATEADD": lambda args: exp.DateAdd(
+                this=seq_get(args, 2), expression=seq_get(args, 1), unit=seq_get(args, 0)
+            ),
+            "DATE_DIFF": lambda args: exp.DateDiff(
+                this=seq_get(args, 2), expression=seq_get(args, 1), unit=seq_get(args, 0)
+            ),
+            "DATEDIFF": lambda args: exp.DateDiff(
+                this=seq_get(args, 2), expression=seq_get(args, 1), unit=seq_get(args, 0)
+            ),
             "MAP": parse_var_map,
             "MATCH": exp.RegexpLike.from_arg_list,
             "UNIQ": exp.ApproxDistinct.from_arg_list,
@@ -341,6 +354,12 @@ class Generator(generator.Generator):
             exp.ApproxDistinct: rename_func("uniq"),
             exp.Array: inline_array_sql,
             exp.CastToStrType: rename_func("CAST"),
+            exp.DateAdd: lambda self, e: self.func(
+                "DATE_ADD", exp.Literal.string(e.text("unit") or "day"), e.expression, e.this
+            ),
+            exp.DateDiff: lambda self, e: self.func(
+                "DATE_DIFF", exp.Literal.string(e.text("unit") or "day"), e.expression, e.this
+            ),
             exp.Final: lambda self, e: f"{self.sql(e, 'this')} FINAL",
             exp.Map: lambda self, e: _lower_func(var_map_sql(self, e)),
             exp.PartitionedByProperty: lambda self, e: f"PARTITION BY {self.sql(e, 'this')}",

diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
@@ -6,6 +6,31 @@ class TestClickhouse(Validator):
     dialect = "clickhouse"
 
     def test_clickhouse(self):
+        self.validate_all(
+            "DATE_ADD('day', 1, x)",
+            read={
+                "clickhouse": "dateAdd(day, 1, x)",
+                "presto": "DATE_ADD('day', 1, x)",
+            },
+            write={
+                "clickhouse": "DATE_ADD('day', 1, x)",
+                "presto": "DATE_ADD('day', 1, x)",
+                "": "DATE_ADD(x, 1, 'day')",
+            },
+        )
+        self.validate_all(
+            "DATE_DIFF('day', a, b)",
+            read={
+                "clickhouse": "dateDiff('day', a, b)",
+                "presto": "DATE_DIFF('day', a, b)",
+            },
+            write={
+                "clickhouse": "DATE_DIFF('day', a, b)",
+                "presto": "DATE_DIFF('day', a, b)",
+                "": "DATEDIFF(b, a, day)",
+            },
+        )
+
         expr = parse_one("count(x)")
         self.assertEqual(expr.sql(dialect="clickhouse"), "COUNT(x)")
         self.assertIsNone(expr._meta)