Repository: tobymao/sqlglot
Pull request: Refactor!(optimizer): only create scopes for the DDL queries
We used to create Scopes for both Query and DDL expressions, but the logic that handled the latter was brittle, incomplete and resulted in another expression type that users of the Scope module had to take into account, thus making the interface more complex to use.

This makes it so that we'll only optimize the _queries_ of statements such as CTAS and INSERT INTO, instead of the whole thing. Same holds for lineage inference - we'll only analyze the DDL queries.

Fixes #3086

Problem statement:
Lineage  - not downstreaming to subquery expression
I'm trying to generate column level lineage from given INSERT statement.

I'm exploring column RUN_DATE,  and expect to get expression for last available downstream.

I expect to get same expression of last step in walk for sql_nok as in sql_ok.


Name: sqlglot
Version: 22.2.1

# Python code

```python
from sqlglot.lineage import lineage

sql_nok = """
insert /*+ HINT1*/ 
into	AML_NWU.DLT_TAXDAILY 
( 
	ORG_UNIT_CODE,
	RUN_DATE 
)
select
	nd.ORG_UNIT_CODE,
	nd.RUN_DATE
from
(select /*+ HINT2*/ 
	MAP_TAXDAILY.ORG_UNIT_CODE as ORG_UNIT_CODE,
	to_date('2023-12-19', 'YYYY-MM-DD') as RUN_DATE
from	AML_NWU.MAP_TAXDAILY MAP_TAXDAILY
where (1=1)
and RUN_DATE = to_date('2023-12-19', 'YYYY-MM-DD')
) nd
"""

sql_ok = """
insert /*+ HINT1*/ 
into	AML_NWU.DLT_TAXDAILY 
( 
	ORG_UNIT_CODE,
	RUN_DATE 
)
select
	MAP_TAXDAILY.ORG_UNIT_CODE as ORG_UNIT_CODE,
	to_date('2023-12-19', 'YYYY-MM-DD') as RUN_DATE
from	AML_NWU.MAP_TAXDAILY MAP_TAXDAILY
where (1=1)
and RUN_DATE = to_date('2023-12-19', 'YYYY-MM-DD')

"""

#sql = sql_nok
sql = sql_ok

column = 'RUN_DATE'
lin = lineage(column, sql, dialect='oracle')

for n in lin.walk():
  print(str(n.name) + ' - ' + str(n.expression))
```

Hints from issue discussion:
Hey @sunrutcon, we'll need more time to figure out how we want to deal with lineage on non-queries, because it will require a bit more work than initially expected. Until then, I'm marking this one as not planned. You can work around this issue by intercepting `DDL` expressions with queries (`expression`) and then computing the lineage on the queries instead.
What do you mean by non-query?
Could you provide me more details on workaround?
`INSERT` is not a query but rather a DML statement. It's not clear to us yet how we want to deal with such statements when it comes to lineage inference. In this particular case, fixing the issue will require refactoring some parts of the codebase, and so we'll have to think more about this before proceeding.

>  Could you provide me more details on workaround?

Sure, this is what I had in mind:

```python
from sqlglot import exp, parse_one
from sqlglot.lineage import lineage

sql = "...."
expr = parse_one(sql, dialect="oracle")

if isinstance(expr, exp.DDL) and isinstance(expr.expression, exp.Query):
    expxr = expr.expression

for n in lineage("...", expr, dialect='oracle').walk():
  print(str(n.name) + ' - ' + str(n.expression))
```
Tnx, I've tested workaround, and it is working great.
I've removed INSERT part and applied lineage on  just SELECT, and it is working as expected.

Commit message: Refactor!(optimizer): don't create scopes for DDLs

We used to create Scopes for both Query and DDL expressions, but the
logic that handled the latter was brittle, incomplete and resulted in
another expression type that users of the Scope module had to take
into account, thus making the interface more complex to use.

This makes it so that we'll only optimize the _queries_ of statements
such as CTAS and INSERT INTO, instead of the whole thing. Same holds
for lineage inference - we'll only analyze the DDL queries.

Patch to review:
diff --git a/setup.py b/setup.py
index ba18f646b2..74ca975679 100644
--- a/setup.py
+++ b/setup.py
@@ -32,6 +32,7 @@ def sqlglotrs_version():
             "duckdb>=0.6",
             "mypy",
             "pandas",
+            "pandas-stubs",
             "pyspark",
             "python-dateutil",
             "pdoc",
diff --git a/sqlglot/lineage.py b/sqlglot/lineage.py
index 6d72a2ef59..0a353c945d 100644
--- a/sqlglot/lineage.py
+++ b/sqlglot/lineage.py
@@ -8,6 +8,7 @@
 from sqlglot import Schema, exp, maybe_parse
 from sqlglot.errors import SqlglotError
 from sqlglot.optimizer import Scope, build_scope, find_all_in_scope, normalize_identifiers, qualify
+from sqlglot.transforms import extract_ddl_query
 
 if t.TYPE_CHECKING:
     from sqlglot.dialects.dialect import DialectType
@@ -106,7 +107,7 @@ def lineage(
         **{"validate_qualify_columns": False, "identify": False, **kwargs},  # type: ignore
     )
 
-    scope = build_scope(qualified)
+    scope = build_scope(extract_ddl_query(qualified))
 
     if not scope:
         raise SqlglotError("Cannot build lineage, sql must be SELECT")
diff --git a/sqlglot/optimizer/optimizer.py b/sqlglot/optimizer/optimizer.py
index 1c96e9503d..deb7081095 100644
--- a/sqlglot/optimizer/optimizer.py
+++ b/sqlglot/optimizer/optimizer.py
@@ -20,6 +20,7 @@
 from sqlglot.optimizer.simplify import simplify
 from sqlglot.optimizer.unnest_subqueries import unnest_subqueries
 from sqlglot.schema import ensure_schema
+from sqlglot.transforms import extract_ddl_query
 
 RULES = (
     qualify,
@@ -69,7 +70,8 @@ def optimize(
         **kwargs: If a rule has a keyword argument with a same name in **kwargs, it will be passed in.
 
     Returns:
-        The optimized expression.
+        The optimized expression. If the original expression is a DDL with a `Query` child, then
+        only the child will be optimized and that will be part of the final DDL expression.
     """
     schema = ensure_schema(schema or sqlglot.schema, dialect=dialect)
     possible_kwargs = {
@@ -82,13 +84,19 @@ def optimize(
         **kwargs,
     }
 
-    expression = exp.maybe_parse(expression, dialect=dialect, copy=True)
+    original = exp.maybe_parse(expression, dialect=dialect, copy=True)
+    optimized = extract_ddl_query(original)
+
     for rule in rules:
         # Find any additional rule parameters, beyond `expression`
         rule_params = rule.__code__.co_varnames
         rule_kwargs = {
             param: possible_kwargs[param] for param in rule_params if param in possible_kwargs
         }
-        expression = rule(expression, **rule_kwargs)
+        optimized = rule(optimized, **rule_kwargs)
+
+    if isinstance(original, exp.DDL) and isinstance(original.expression, exp.Query):
+        original.set("expression", optimized)
+        optimized = original
 
-    return t.cast(exp.Expression, expression)
+    return optimized
diff --git a/sqlglot/optimizer/qualify.py b/sqlglot/optimizer/qualify.py
index e4f8b57589..0d3434e742 100644
--- a/sqlglot/optimizer/qualify.py
+++ b/sqlglot/optimizer/qualify.py
@@ -14,6 +14,7 @@
 )
 from sqlglot.optimizer.qualify_tables import qualify_tables
 from sqlglot.schema import Schema, ensure_schema
+from sqlglot.transforms import extract_ddl_query
 
 
 def qualify(
@@ -62,9 +63,13 @@ def qualify(
         identify: If True, quote all identifiers, else only necessary ones.
 
     Returns:
-        The qualified expression.
+        The qualified expression. If the original expression is a DDL with a `Query` child, then
+        only the child will be qualified and that will be part of the final DDL expression.
     """
     schema = ensure_schema(schema, dialect=dialect)
+
+    original = expression
+    expression = extract_ddl_query(expression)
     expression = normalize_identifiers(expression, dialect=dialect)
     expression = qualify_tables(expression, db=db, catalog=catalog, schema=schema)
 
@@ -89,4 +94,8 @@ def qualify(
     if validate_qualify_columns:
         validate_qualify_columns_func(expression)
 
+    if isinstance(original, exp.DDL) and isinstance(original.expression, exp.Query):
+        original.set("expression", expression)
+        expression = original
+
     return expression
diff --git a/sqlglot/optimizer/qualify_tables.py b/sqlglot/optimizer/qualify_tables.py
index 214ac0ae8d..fabed5e332 100644
--- a/sqlglot/optimizer/qualify_tables.py
+++ b/sqlglot/optimizer/qualify_tables.py
@@ -8,6 +8,7 @@
 from sqlglot.helper import csv_reader, name_sequence
 from sqlglot.optimizer.scope import Scope, traverse_scope
 from sqlglot.schema import Schema
+from sqlglot.transforms import extract_ddl_query
 
 if t.TYPE_CHECKING:
     from sqlglot._typing import E
@@ -60,7 +61,7 @@ def _qualify(table: exp.Table) -> None:
             if isinstance(node, exp.Table):
                 _qualify(node)
 
-    for scope in traverse_scope(expression):
+    for scope in traverse_scope(extract_ddl_query(expression)):
         for derived_table in itertools.chain(scope.ctes, scope.derived_tables):
             if isinstance(derived_table, exp.Subquery):
                 unnested = derived_table.unnest()
diff --git a/sqlglot/optimizer/scope.py b/sqlglot/optimizer/scope.py
index 595131c9db..3b61ae5159 100644
--- a/sqlglot/optimizer/scope.py
+++ b/sqlglot/optimizer/scope.py
@@ -8,7 +8,7 @@
 
 from sqlglot import exp
 from sqlglot.errors import OptimizeError
-from sqlglot.helper import ensure_collection, find_new_name
+from sqlglot.helper import ensure_collection, find_new_name, seq_get
 
 logger = logging.getLogger("sqlglot")
 
@@ -122,7 +122,8 @@ def _collect(self):
         for node, parent, _ in self.walk(bfs=False):
             if node is self.expression:
                 continue
-            elif isinstance(node, exp.Column) and not isinstance(node.this, exp.Star):
+
+            if isinstance(node, exp.Column) and not isinstance(node.this, exp.Star):
                 self._raw_columns.append(node)
             elif isinstance(node, exp.Table) and not isinstance(node.parent, exp.JoinHint):
                 self._tables.append(node)
@@ -477,17 +478,12 @@ def traverse_scope(expression: exp.Expression) -> t.List[Scope]:
         ('SELECT a FROM (SELECT a FROM x) AS y', ['y'])
 
     Args:
-        expression (exp.Expression): expression to traverse
+        expression: Expression to traverse
 
     Returns:
-        list[Scope]: scope instances
+        A list of the created scope instances
     """
-    if isinstance(expression, exp.Query) or (
-        isinstance(expression, exp.DDL) and isinstance(expression.expression, exp.Query)
-    ):
-        return list(_traverse_scope(Scope(expression)))
-
-    return []
+    return list(_traverse_scope(Scope(expression))) if isinstance(expression, exp.Query) else []
 
 
 def build_scope(expression: exp.Expression) -> t.Optional[Scope]:
@@ -495,14 +491,12 @@ def build_scope(expression: exp.Expression) -> t.Optional[Scope]:
     Build a scope tree.
 
     Args:
-        expression (exp.Expression): expression to build the scope tree for
+        expression: Expression to build the scope tree for.
+
     Returns:
-        Scope: root scope
+        The root scope
     """
-    scopes = traverse_scope(expression)
-    if scopes:
-        return scopes[-1]
-    return None
+    return seq_get(traverse_scope(expression), -1)
 
 
 def _traverse_scope(scope):
diff --git a/sqlglot/transforms.py b/sqlglot/transforms.py
index 04c1f7b1e3..392143d886 100644
--- a/sqlglot/transforms.py
+++ b/sqlglot/transforms.py
@@ -561,9 +561,7 @@ def move_partitioned_by_to_schema_columns(expression: exp.Expression) -> exp.Exp
 
 
 def struct_kv_to_alias(expression: exp.Expression) -> exp.Expression:
-    """
-    Convert struct arguments to aliases: STRUCT(1 AS y) .
-    """
+    """Converts struct arguments to aliases, e.g. STRUCT(1 AS y)."""
     if isinstance(expression, exp.Struct):
         expression.set(
             "expressions",
@@ -576,6 +574,30 @@ def struct_kv_to_alias(expression: exp.Expression) -> exp.Expression:
     return expression
 
 
+def extract_ddl_query(expression: exp.Expression) -> exp.Expression:
+    """
+    Extracts the query of a given DDL expression. If the latter has CTEs attached to it,
+    they are removed from it and attached to the query.
+
+    Example:
+        >>> import sqlglot
+        >>> ddl = sqlglot.parse_one("WITH cte AS (SELECT 1 AS c) INSERT INTO t SELECT c FROM cte")
+        >>> ddl.transform(extract_ddl_query).sql()
+        'WITH cte AS (SELECT 1 AS c) SELECT c FROM cte'
+    """
+    query = expression.expression
+    if isinstance(expression, exp.DDL) and isinstance(query, exp.Query):
+        ddl_with = expression.args.get("with")
+        expression = query
+        if ddl_with:
+            for cte in ddl_with.pop().expressions:
+                query.with_(cte.alias, cte.this, copy=False)
+
+            query.args["with"].set("recursive", ddl_with.recursive)
+
+    return expression
+
+
 def preprocess(
     transforms: t.List[t.Callable[[exp.Expression], exp.Expression]],
 ) -> t.Callable[[Generator, exp.Expression], str]:
diff --git a/tests/fixtures/optimizer/optimizer.sql b/tests/fixtures/optimizer/optimizer.sql
index 990453ba32..a62813c73f 100644
--- a/tests/fixtures/optimizer/optimizer.sql
+++ b/tests/fixtures/optimizer/optimizer.sql
@@ -32,17 +32,18 @@ FROM (
               )
          WHERE object_pointstext IS NOT NULL
      );
-CREATE OR REPLACE TEMPORARY VIEW `latest_boo` AS
-SELECT
-  TRIM(SPLIT(`_q_1`.`points`, ':')[0]) AS `points_type`,
-  TRIM(SPLIT(`_q_1`.`points`, ':')[1]) AS `points_value`
-FROM (
+CREATE OR REPLACE TEMPORARY VIEW latest_boo AS
+WITH `_q_1` AS (
   SELECT
     EXPLODE_OUTER(SPLIT(`boo`.`object_pointstext`, ',')) AS `points`
   FROM `boo` AS `boo`
   WHERE
     NOT `boo`.`object_pointstext` IS NULL
-) AS `_q_1`;
+)
+SELECT
+  TRIM(SPLIT(`_q_1`.`points`, ':')[0]) AS `points_type`,
+  TRIM(SPLIT(`_q_1`.`points`, ':')[1]) AS `points_value`
+FROM `_q_1` AS `_q_1`;
 
 # title: Union in CTE
 WITH cte AS (
diff --git a/tests/fixtures/optimizer/qualify_columns_ddl.sql b/tests/fixtures/optimizer/qualify_columns_ddl.sql
index 907780b70c..3a31e44ed4 100644
--- a/tests/fixtures/optimizer/qualify_columns_ddl.sql
+++ b/tests/fixtures/optimizer/qualify_columns_ddl.sql
@@ -1,6 +1,6 @@
 # title: Create with CTE
 WITH cte AS (SELECT b FROM y) CREATE TABLE s AS SELECT * FROM cte;
-WITH cte AS (SELECT y.b AS b FROM y AS y) CREATE TABLE s AS SELECT cte.b AS b FROM cte AS cte;
+CREATE TABLE s AS WITH cte AS (SELECT y.b AS b FROM y AS y) SELECT cte.b AS b FROM cte AS cte;
 
 # title: Create without CTE
 CREATE TABLE foo AS SELECT a FROM tbl;
@@ -8,15 +8,15 @@ CREATE TABLE foo AS SELECT tbl.a AS a FROM tbl AS tbl;
 
 # title: Create with complex CTE with derived table
 WITH cte AS (SELECT a FROM (SELECT a from x)) CREATE TABLE s AS SELECT * FROM cte;
-WITH cte AS (SELECT _q_0.a AS a FROM (SELECT x.a AS a FROM x AS x) AS _q_0) CREATE TABLE s AS SELECT cte.a AS a FROM cte AS cte;
+CREATE TABLE s AS WITH cte AS (SELECT _q_0.a AS a FROM (SELECT x.a AS a FROM x AS x) AS _q_0) SELECT cte.a AS a FROM cte AS cte;
 
 # title: Create wtih multiple CTEs
 WITH cte1 AS (SELECT b FROM y), cte2 AS (SELECT b FROM cte1) CREATE TABLE s AS SELECT * FROM cte2;
-WITH cte1 AS (SELECT y.b AS b FROM y AS y), cte2 AS (SELECT cte1.b AS b FROM cte1 AS cte1) CREATE TABLE s AS SELECT cte2.b AS b FROM cte2 AS cte2;
+CREATE TABLE s AS WITH cte1 AS (SELECT y.b AS b FROM y AS y), cte2 AS (SELECT cte1.b AS b FROM cte1 AS cte1) SELECT cte2.b AS b FROM cte2 AS cte2;
 
 # title: Create with multiple CTEs, selecting only from the first CTE (unnecessary code)
 WITH cte1 AS (SELECT b FROM y), cte2 AS (SELECT b FROM cte1) CREATE TABLE s AS SELECT * FROM cte1;
-WITH cte1 AS (SELECT y.b AS b FROM y AS y), cte2 AS (SELECT cte1.b AS b FROM cte1 AS cte1) CREATE TABLE s AS SELECT cte1.b AS b FROM cte1 AS cte1;
+CREATE TABLE s AS WITH cte1 AS (SELECT y.b AS b FROM y AS y), cte2 AS (SELECT cte1.b AS b FROM cte1 AS cte1) SELECT cte1.b AS b FROM cte1 AS cte1;
 
 # title: Create with multiple derived tables
 CREATE TABLE s AS SELECT * FROM (SELECT b FROM (SELECT b FROM y));
@@ -24,9 +24,10 @@ CREATE TABLE s AS SELECT _q_1.b AS b FROM (SELECT _q_0.b AS b FROM (SELECT y.b A
 
 # title: Create with a CTE and a derived table
 WITH cte AS (SELECT b FROM y) CREATE TABLE s AS SELECT * FROM (SELECT b FROM (SELECT b FROM cte));
-WITH cte AS (SELECT y.b AS b FROM y AS y) CREATE TABLE s AS SELECT _q_1.b AS b FROM (SELECT _q_0.b AS b FROM (SELECT cte.b AS b FROM cte AS cte) AS _q_0) AS _q_1;
+CREATE TABLE s AS WITH cte AS (SELECT y.b AS b FROM y AS y) SELECT _q_1.b AS b FROM (SELECT _q_0.b AS b FROM (SELECT cte.b AS b FROM cte AS cte) AS _q_0) AS _q_1;
 
 # title: Insert with CTE
+# dialect: spark
 WITH cte AS (SELECT b FROM y) INSERT INTO s SELECT * FROM cte;
 WITH cte AS (SELECT y.b AS b FROM y AS y) INSERT INTO s SELECT cte.b AS b FROM cte AS cte;
 
diff --git a/tests/fixtures/optimizer/qualify_tables.sql b/tests/fixtures/optimizer/qualify_tables.sql
index 99b5153eae..f651a871ef 100644
--- a/tests/fixtures/optimizer/qualify_tables.sql
+++ b/tests/fixtures/optimizer/qualify_tables.sql
@@ -159,6 +159,7 @@ CREATE TABLE t1 AS (WITH cte AS (SELECT x FROM t2) SELECT * FROM cte);
 CREATE TABLE c.db.t1 AS (WITH cte AS (SELECT x FROM c.db.t2 AS t2) SELECT * FROM cte AS cte);
 
 # title: insert statement with cte
+# dialect: spark
 WITH cte AS (SELECT b FROM y) INSERT INTO s SELECT * FROM cte;
 WITH cte AS (SELECT b FROM c.db.y AS y) INSERT INTO c.db.s SELECT * FROM cte AS cte;
 
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
index c48ce4564c..6312e86997 100644
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -388,3 +388,28 @@ def test_lineage_normalize(self) -> None:
 
         with self.assertRaises(sqlglot.errors.SqlglotError):
             lineage('"a"', "WITH x AS (SELECT 1 a) SELECT a FROM x", dialect="snowflake")
+
+    def test_ddl_lineage(self) -> None:
+        sql = """
+        INSERT /*+ HINT1 */
+        INTO target (x, y)
+        SELECT subq.x, subq.y
+        FROM (
+          SELECT /*+ HINT2 */
+            t.x AS x,
+            TO_DATE('2023-12-19', 'YYYY-MM-DD') AS y
+          FROM s.t t
+          WHERE 1 = 1 AND y = TO_DATE('2023-12-19', 'YYYY-MM-DD')
+        ) subq
+        """
+
+        node = lineage("y", sql, dialect="oracle")
+
+        self.assertEqual(node.name, "Y")
+        self.assertEqual(node.expression.sql(dialect="oracle"), "SUBQ.Y AS Y")
+
+        downstream = node.downstream[0]
+        self.assertEqual(downstream.name, "SUBQ.Y")
+        self.assertEqual(
+            downstream.expression.sql(dialect="oracle"), "TO_DATE('2023-12-19', 'YYYY-MM-DD') AS Y"
+        )
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 046e5a60c6..ce2e526b59 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -298,7 +298,9 @@ def test_qualify_columns(self, logger):
         self.check_file(
             "qualify_columns", qualify_columns, execute=True, schema=self.schema, set_dialect=True
         )
-        self.check_file("qualify_columns_ddl", qualify_columns, schema=self.schema)
+        self.check_file(
+            "qualify_columns_ddl", qualify_columns, schema=self.schema, set_dialect=True
+        )
 
     def test_qualify_columns__with_invisible(self):
         schema = MappingSchema(self.schema, {"x": {"a"}, "y": {"b"}, "z": {"b"}})


Code review comments:
diff_hunk: @@ -89,4 +94,8 @@ def qualify(
     if validate_qualify_columns:
         validate_qualify_columns_func(expression)
 
+    if isinstance(original, exp.DDL) and isinstance(original.expression, exp.Query):
path: sqlglot/optimizer/qualify.py
text: @user1:
This might be a little confusing since its so far from `extract_ddl_query`.

I think this might be more explicit:
```python
original = expression
expression = extract_ddl_query(expression)
ddl = original is not expression

...

if ddl:
    # etc
```

@author:
Done

diff_hunk: @@ -576,6 +574,30 @@ def struct_kv_to_alias(expression: exp.Expression) -> exp.Expression:
     return expression
 
 
+def extract_ddl_query(expression: exp.Expression) -> exp.Expression:
+    """
+    Extracts the query of a given DDL expression. If the latter has CTEs attached to it,
+    they are removed from it and attached to the query.
+
+    Example:
+        >>> import sqlglot
+        >>> ddl = sqlglot.parse_one("WITH cte AS (SELECT 1 AS c) INSERT INTO t SELECT c FROM cte")
+        >>> ddl.transform(extract_ddl_query).sql()
+        'WITH cte AS (SELECT 1 AS c) SELECT c FROM cte'
+    """
+    query = expression.expression
+    if isinstance(expression, exp.DDL) and isinstance(query, exp.Query):
+        ddl_with = expression.args.get("with")
+        expression = query
+        if ddl_with:
+            for cte in ddl_with.pop().expressions:
+                query.with_(cte.alias, cte.this, copy=False)
path: sqlglot/transforms.py
text: @user1:
We need to prepend, in case the query _also_ has CTEs

@author:
Yeah I kinda did this on purpose because I was only thinking about the name shadowing concerns (which seemed unnecessary to deal with), but thinking about it again, you're right. If the query also has CTEs that reference the DDL ctes then this breaks them. Will fix, thanks for catching.

@author:
Fixed