Add more tests

tobymao · Oct 8, 2024 · 8a9bb9c · 8a9bb9c
1 parent f9e35d4
commit 8a9bb9c
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
@@ -762,12 +762,6 @@ def test_scope(self):
  )
  self.assertEqual(set(scopes[3].sources), {""})
 
- sql = (
- "UPDATE customers SET total_spent = (SELECT 1 FROM t1) WHERE EXISTS (SELECT 1 FROM t2)"
- )
- expression = parse_one(sql)
- self.assertEqual(len(traverse_scope(expression)), 3)
-
  inner_query = "SELECT bar FROM baz"
  for udtf in (f"UNNEST(({inner_query}))", f"LATERAL ({inner_query})"):
  sql = f"SELECT a FROM foo CROSS JOIN {udtf}"
@@ -785,6 +779,18 @@ def test_scope(self):
  self.assertEqual(scopes[2].expression.sql(), f"SELECT a FROM foo CROSS JOIN {udtf}")
  self.assertEqual(set(scopes[2].sources), {"", "foo"})
 
+ # Check DML statement scopes
+ sql = (
+ "UPDATE customers SET total_spent = (SELECT 1 FROM t1) WHERE EXISTS (SELECT 1 FROM t2)"
+ )
+ self.assertEqual(len(traverse_scope(parse_one(sql))), 3)
+
+ sql = "UPDATE tbl1 SET col = 1 WHERE EXISTS (SELECT 1 FROM tbl2 WHERE tbl1.id = tbl2.id)"
+ self.assertEqual(len(traverse_scope(parse_one(sql))), 1)
+
+ sql = "UPDATE tbl1 SET col = 0"
+ self.assertEqual(len(traverse_scope(parse_one(sql))), 0)
+
  @patch("sqlglot.optimizer.scope.logger")
  def test_scope_warning(self, logger):
  self.assertEqual(len(traverse_scope(parse_one("WITH q AS (@y) SELECT * FROM q"))), 1)