feat(tsql): add support for DATETRUNC #4531 (#4537)

* feat(tsql): add support for DATETRUNC * refactor(tsql): minor refactoring of _parse_datetrunc method * refactor(tsql): minor refactoring removed comment * PR Feedback 1
tobymao · Dec 19, 2024 · 7fe9f7f · 7fe9f7f
1 parent 6992c18
commit 7fe9f7f
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 0 deletions.
diff --git a/sqlglot/dialects/tsql.py b/sqlglot/dialects/tsql.py
@@ -370,6 +370,16 @@ def _timestrtotime_sql(self: TSQL.Generator, expression: exp.TimeStrToTime):
     return sql
 
 
+def _build_datetrunc(args: t.List) -> exp.TimestampTrunc:
+    unit = seq_get(args, 0)
+    this = seq_get(args, 1)
+
+    if this and this.is_string:
+        this = exp.cast(this, exp.DataType.Type.DATETIME2)
+
+    return exp.TimestampTrunc(this=this, unit=unit)
+
+
 class TSQL(Dialect):
     SUPPORTS_SEMI_ANTI_JOIN = False
     LOG_BASE_FIRST = False
@@ -570,6 +580,7 @@ class Parser(parser.Parser):
             "SUSER_SNAME": exp.CurrentUser.from_arg_list,
             "SYSTEM_USER": exp.CurrentUser.from_arg_list,
             "TIMEFROMPARTS": _build_timefromparts,
+            "DATETRUNC": _build_datetrunc,
         }
 
         JOIN_HINTS = {"LOOP", "HASH", "MERGE", "REMOTE"}
@@ -936,6 +947,7 @@ class Generator(generator.Generator):
             exp.Trim: trim_sql,
             exp.TsOrDsAdd: date_delta_sql("DATEADD", cast=True),
             exp.TsOrDsDiff: date_delta_sql("DATEDIFF"),
+            exp.TimestampTrunc: lambda self, e: self.func("DATETRUNC", e.unit, e.this),
         }
 
         TRANSFORMS.pop(exp.ReturnsProperty)

diff --git a/tests/dialects/test_tsql.py b/tests/dialects/test_tsql.py
@@ -2090,3 +2090,27 @@ def test_next_value_for(self):
                 "oracle": "SELECT NEXT VALUE FOR db.schema.sequence_name",
             },
         )
+
+    # string literals in the DATETRUNC are casted as DATETIME2
+    def test_datetrunc(self):
+        self.validate_all(
+            "SELECT DATETRUNC(month, 'foo')",
+            write={
+                "duckdb": "SELECT DATE_TRUNC('MONTH', CAST('foo' AS TIMESTAMP))",
+                "tsql": "SELECT DATETRUNC(MONTH, CAST('foo' AS DATETIME2))",
+            },
+        )
+        self.validate_all(
+            "SELECT DATETRUNC(month, foo)",
+            write={
+                "duckdb": "SELECT DATE_TRUNC('MONTH', foo)",
+                "tsql": "SELECT DATETRUNC(MONTH, foo)",
+            },
+        )
+        self.validate_all(
+            "SELECT DATETRUNC(year, CAST('foo1' AS date))",
+            write={
+                "duckdb": "SELECT DATE_TRUNC('YEAR', CAST('foo1' AS DATE))",
+                "tsql": "SELECT DATETRUNC(YEAR, CAST('foo1' AS DATE))",
+            },
+        )