Optimizer: add datetime simplification, fix date (#922)

sqlglot/optimizer/simplify.py

@@ -361,15 +361,15 @@ def _simplify_binary(expression, a, b):
             return boolean
     elif isinstance(a, exp.Cast) and isinstance(b, exp.Interval):
         a, b = extract_date(a), extract_interval(b)
-        if b:
+        if a and b:
             if isinstance(expression, exp.Add):
                 return date_literal(a + b)
             if isinstance(expression, exp.Sub):
                 return date_literal(a - b)
     elif isinstance(a, exp.Interval) and isinstance(b, exp.Cast):
         a, b = extract_interval(a), extract_date(b)
         # you cannot subtract a date from an interval
-        if a and isinstance(expression, exp.Add):
+        if a and b and isinstance(expression, exp.Add):
             return date_literal(a + b)
 
     return None
@@ -424,9 +424,15 @@ def eval_boolean(expression, a, b):
 
 
 def extract_date(cast):
-    if cast.args["to"].this == exp.DataType.Type.DATE:
-        return datetime.date.fromisoformat(cast.name)
-    return None
+    # The "fromisoformat" conversion could fail if the cast is used on an identifier,
+    # so in that case we can't extract the date.
+    try:
+        if cast.args["to"].this == exp.DataType.Type.DATE:
+            return datetime.date.fromisoformat(cast.name)
+        if cast.args["to"].this == exp.DataType.Type.DATETIME:
+            return datetime.datetime.fromisoformat(cast.name)
+    except ValueError:
+        return None
 
 
 def extract_interval(interval):
@@ -450,7 +456,8 @@ def extract_interval(interval):
 
 
 def date_literal(date):
-    return exp.Cast(this=exp.Literal.string(date), to=exp.DataType.build("DATE"))
+    expr_type = exp.DataType.build("DATETIME" if isinstance(date, datetime.datetime) else "DATE")
+    return exp.Cast(this=exp.Literal.string(date), to=expr_type)
 
 
 def boolean_literal(condition):

tests/fixtures/optimizer/simplify.sql

@@ -375,6 +375,18 @@ CAST('1998-12-01' AS DATE) - INTERVAL '90' foo;
 date '1998-12-01' + interval '90' foo;
 CAST('1998-12-01' AS DATE) + INTERVAL '90' foo;
 
+CAST(x AS DATE) + interval '1' week;
+CAST(x AS DATE) + INTERVAL '1' week;
+
+CAST('2008-11-11' AS DATETIME) + INTERVAL '5' MONTH;
+CAST('2009-04-11 00:00:00' AS DATETIME);
+
+datetime '1998-12-01' - interval '90' day;
+CAST('1998-09-02 00:00:00' AS DATETIME);
+
+CAST(x AS DATETIME) + interval '1' week;
+CAST(x AS DATETIME) + INTERVAL '1' week;
+
 --------------------------------------
 -- Comparisons
 --------------------------------------