diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py index 77a067404..8bfbfdd87 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 fafd61fe1..0bd7ab26f 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/fixtures/optimizer/qualify_tables.sql b/tests/fixtures/optimizer/qualify_tables.sql index 61d0b9677..49e07faf9 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); diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index 2c2015b3d..2a51d84ca 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -779,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)