From f9e35d418164c849a5468acb747bc5dee40f2838 Mon Sep 17 00:00:00 2001
From: vaggelisd <daniasevangelos@gmail.com>
Date: Tue, 8 Oct 2024 13:47:58 +0300
Subject: [PATCH 1/3] feat: Make exp.Update be a DML node

---
 sqlglot/expressions.py     | 2 +-
 sqlglot/optimizer/scope.py | 4 ++--
 tests/test_optimizer.py    | 6 ++++++
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 77a067404b..8bfbfdd874 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -3278,7 +3278,7 @@ class Intersect(SetOperation):
     pass
 
 
-class Update(Expression):
+class Update(DML):
     arg_types = {
         "with": False,
         "this": False,
diff --git a/sqlglot/optimizer/scope.py b/sqlglot/optimizer/scope.py
index fafd61fe15..0bd7ab26f0 100644
--- a/sqlglot/optimizer/scope.py
+++ b/sqlglot/optimizer/scope.py
@@ -562,8 +562,8 @@ def _traverse_scope(scope):
     elif isinstance(expression, exp.DML):
         yield from _traverse_ctes(scope)
         for query in find_all_in_scope(expression, exp.Query):
-            # This check ensures we don't yield the CTE queries twice
-            if not isinstance(query.parent, exp.CTE):
+            # This check ensures we don't yield the CTE/nested queries twice
+            if not isinstance(query.parent, (exp.CTE, exp.Subquery)):
                 yield from _traverse_scope(Scope(query, cte_sources=scope.cte_sources))
         return
     else:
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 2c2015b3d9..b1c2b7bbae 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -762,6 +762,12 @@ 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}"

From 8a9bb9ca9f01b646013b83e67eb1c0a1172ca75f Mon Sep 17 00:00:00 2001
From: vaggelisd <daniasevangelos@gmail.com>
Date: Tue, 8 Oct 2024 14:46:23 +0300
Subject: [PATCH 2/3] Add more tests

---
 tests/test_optimizer.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index b1c2b7bbae..2a51d84cac 100644
--- 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)

From 87d18fb0a2b4c74503018d529dde4192434064d2 Mon Sep 17 00:00:00 2001
From: vaggelisd <daniasevangelos@gmail.com>
Date: Tue, 8 Oct 2024 15:17:26 +0300
Subject: [PATCH 3/3] Add qualify test

---
 tests/fixtures/optimizer/qualify_tables.sql | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/fixtures/optimizer/qualify_tables.sql b/tests/fixtures/optimizer/qualify_tables.sql
index 61d0b9677f..49e07faf9b 100644
--- a/tests/fixtures/optimizer/qualify_tables.sql
+++ b/tests/fixtures/optimizer/qualify_tables.sql
@@ -184,3 +184,6 @@ COPY INTO (SELECT * FROM c.db.x AS x) TO 'data' WITH (FORMAT 'CSV');
 # title: tablesample
 SELECT 1 FROM x TABLESAMPLE SYSTEM (10 PERCENT) CROSS JOIN y TABLESAMPLE SYSTEM (10 PERCENT);
 SELECT 1 FROM c.db.x AS x TABLESAMPLE SYSTEM (10 PERCENT) CROSS JOIN c.db.y AS y TABLESAMPLE SYSTEM (10 PERCENT);
+
+WITH cte_tbl AS (SELECT 1 AS col2) UPDATE y SET col1 = (SELECT * FROM x) WHERE EXISTS(SELECT 1 FROM cte_tbl);
+WITH cte_tbl AS (SELECT 1 AS col2) UPDATE c.db.y SET col1 = (SELECT * FROM c.db.x AS x) WHERE EXISTS(SELECT 1 FROM cte_tbl AS cte_tbl);