feat(optimizer)!: annotate type for Snowflake DATEADD function

sqlglot/dialects/snowflake.py

@@ -547,6 +547,40 @@ def _annotate_reverse(self: TypeAnnotator, expression: exp.Reverse) -> exp.Rever
     return expression
 
 
+def _annotate_dateadd(self: TypeAnnotator, expression: exp.DateAdd) -> exp.DateAdd:
+    self._annotate_args(expression)
+
+    this = expression.this
+    unit = expression.unit
+
+    if not this.type or not unit:
+        self._set_type(expression, exp.DataType.Type.UNKNOWN)
+        return expression
+
+    input_type = this.type.this
+    unit_name = unit.name.upper() if hasattr(unit, "name") else str(unit).upper()
+
+    # If input is TIME, return TIME
+    if input_type == exp.DataType.Type.TIME:
+        self._set_type(expression, exp.DataType.Type.TIME)
+    # If input is TIMESTAMP, return TIMESTAMP
+    elif input_type == exp.DataType.Type.TIMESTAMP:
+        self._set_type(expression, exp.DataType.Type.TIMESTAMP)
+    # If input is DATE:
+    elif input_type == exp.DataType.Type.DATE:
+        # If unit is day or larger, return DATE
+        if unit_name in ("YEAR", "QUARTER", "MONTH", "WEEK", "DAY"):
+            self._set_type(expression, exp.DataType.Type.DATE)
+        # If unit is smaller than day, return TIMESTAMP
+        else:
+            self._set_type(expression, exp.DataType.Type.TIMESTAMP)
+    else:
+        # Fallback to default behavior
+        self._set_type(expression, exp.DataType.Type.UNKNOWN)
+
+    return expression
+
+
 class Snowflake(Dialect):
     # https://docs.snowflake.com/en/sql-reference/identifiers-syntax
     NORMALIZATION_STRATEGY = NormalizationStrategy.UPPERCASE
@@ -698,6 +732,7 @@ class Snowflake(Dialect):
             if e.args.get("source_tz")
             else exp.DataType.Type.TIMESTAMPTZ,
         ),
+        exp.DateAdd: _annotate_dateadd,
         exp.Reverse: _annotate_reverse,
     }
 

tests/fixtures/optimizer/annotate_functions.sql

@@ -1695,6 +1695,22 @@ INT;
 DATE_PART('day', tbl.date_col);
 INT;
 
+# dialect: snowflake
+DATEADD(HOUR, 3, TO_TIME('05:00:00'));
+TIME;
+
+# dialect: snowflake
+DATEADD(YEAR, 1, TO_TIMESTAMP('2022-05-08 14:30:00'));
+TIMESTAMP;
+
+# dialect: snowflake
+DATEADD(MONTH, 1, '2023-01-31'::DATE);
+DATE;
+
+# dialect: snowflake
+DATEADD(HOUR, 2, '2022-04-05'::DATE);
+TIMESTAMP;
+
 # dialect: snowflake
 DEGREES(PI()/3);
 DOUBLE;