Fix: fill in missing implementation details for replace(None) (#3166)

tobymao · Mar 19, 2024 · 5e18d49 · 5e18d49
1 parent cdbe39e
commit 5e18d49
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -621,7 +621,7 @@ def replace(self, expression):
             return expression
 
         key = self.arg_key
-        value = parent.args[key]
+        value = parent.args.get(key)
 
         if isinstance(value, list):
             index = self.index
@@ -639,7 +639,7 @@ def replace(self, expression):
                 else:
                     value[index] = expression
                     parent._set_parent(key, expression, index=index)
-        else:
+        elif value is not None:
             if expression is None:
                 parent.args.pop(key)
             else:

diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
@@ -342,6 +342,9 @@ def test_pushdown_projection(self):
     def test_simplify(self):
         self.check_file("simplify", simplify, set_dialect=True)
 
+        expression = parse_one("SELECT a, c, b FROM table1 WHERE 1 = 1")
+        self.assertEqual(simplify(simplify(expression.find(exp.Where))).sql(), "WHERE TRUE")
+
         expression = parse_one("TRUE AND TRUE AND TRUE")
         self.assertEqual(exp.true(), optimizer.simplify.simplify(expression))
         self.assertEqual(exp.true(), optimizer.simplify.simplify(expression.this))