From b26b8d88af14f72d90c0019ec332d268a23b078f Mon Sep 17 00:00:00 2001
From: Jo <46752250+GeorgeSittas@users.noreply.github.com>
Date: Thu, 12 Jan 2023 19:56:20 +0200
Subject: [PATCH] Optimizer: add datetime simplification, fix date (#922)

---
 sqlglot/optimizer/simplify.py         | 19 +++++++++++++------
 tests/fixtures/optimizer/simplify.sql | 12 ++++++++++++
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/sqlglot/optimizer/simplify.py b/sqlglot/optimizer/simplify.py
index c0719f250d..f560760f82 100644
--- a/sqlglot/optimizer/simplify.py
+++ b/sqlglot/optimizer/simplify.py
@@ -361,7 +361,7 @@ 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):
@@ -369,7 +369,7 @@ def _simplify_binary(expression, 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):
diff --git a/tests/fixtures/optimizer/simplify.sql b/tests/fixtures/optimizer/simplify.sql
index cf4195d5ad..4e9e70c1be 100644
--- a/tests/fixtures/optimizer/simplify.sql
+++ b/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
 --------------------------------------