Repository: reata/sqllineage
Pull request #395 (2023-06-19 16:08:23)

fix: cte inside bracketed of insert
re-fix issue: #377 

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/395?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#395](https://app.codecov.io/gh/reata/sqllineage/pull/395?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (8795fe1) into [master](https://app.codecov.io/gh/reata/sqllineage/commit/52f3b758b302658b4d539098adec6946b4409b74?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (52f3b75) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master     #395   +/-   ##
=======================================
  Coverage   99.40%   99.40%           
=======================================
  Files          42       42           
  Lines        2008     2015    +7     
=======================================
+ Hits         1996     2003    +7     
  Misses         12       12           
```


| [Impacted Files](https://app.codecov.io/gh/reata/sqllineage/pull/395?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [...parser/sqlfluff/extractors/dml\_insert\_extractor.py](https://app.codecov.io/gh/reata/sqllineage/pull/395?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL3BhcnNlci9zcWxmbHVmZi9leHRyYWN0b3JzL2RtbF9pbnNlcnRfZXh0cmFjdG9yLnB5) | `100.00% <100.00%> (ø)` | |
| [sqllineage/core/parser/sqlfluff/utils.py](https://app.codecov.io/gh/reata/sqllineage/pull/395?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL3BhcnNlci9zcWxmbHVmZi91dGlscy5weQ==) | `98.20% <100.00%> (+0.02%)` | :arrow_up: |


reata: I have removed the `is_cte` function as it doesn't seem that we can reuse it. Also the test case is moved to `test_sqlfluff.py`, which contains all the new features we don't/won't support in sqlparse. Other than that, this PR is perfect fine.

Patch:
diff --git a/sqllineage/core/parser/sqlfluff/extractors/dml_insert_extractor.py b/sqllineage/core/parser/sqlfluff/extractors/dml_insert_extractor.py
index 4e4943c..538b6d9 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/dml_insert_extractor.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/dml_insert_extractor.py
@@ -62,6 +62,19 @@ class DmlInsertExtractor(LineageHolderExtractor):
                     segment,
                     AnalyzerContext(prev_cte=holder.cte, prev_write=holder.write),
                 )
+            elif segment.type == "bracketed" and any(
+                s.type == "with_compound_statement" for s in segment.segments
+            ):
+                for sgmt in segment.segments:
+                    if sgmt.type == "with_compound_statement":
+                        from .cte_extractor import DmlCteExtractor
+
+                        holder |= DmlCteExtractor(self.dialect).extract(
+                            sgmt,
+                            AnalyzerContext(
+                                prev_cte=holder.cte, prev_write=holder.write
+                            ),
+                        )
             elif segment.type == "bracketed" and (
                 self.parse_subquery(segment) or is_union(segment)
             ):


Test patch:
diff --git a/tests/test_others_dialect_specific.py b/tests/test_others_dialect_specific.py
index c7cbc8f..ddbf51e 100644
--- a/tests/test_others_dialect_specific.py
+++ b/tests/test_others_dialect_specific.py
@@ -94,38 +94,6 @@ def test_alter_table_exchange_partition(dialect: str):
     )
 
 
-@pytest.mark.parametrize("dialect", ["snowflake", "bigquery"])
-def test_create_clone(dialect: str):
-    """
-    Language manual:
-        https://cloud.google.com/bigquery/docs/table-clones-create
-        https://docs.snowflake.com/en/sql-reference/sql/create-clone
-    Note clone is not a keyword in sqlparse, we'll skip testing for it.
-    """
-    assert_table_lineage_equal(
-        "create table tab2 CLONE tab1;",
-        {"tab1"},
-        {"tab2"},
-        dialect=dialect,
-        test_sqlparse=False,
-    )
-
-
-@pytest.mark.parametrize("dialect", ["snowflake"])
-def test_alter_table_swap_partition(dialect: str):
-    """
-    See https://docs.snowflake.com/en/sql-reference/sql/alter-table for language manual
-    Note swap is not a keyword in sqlparse, we'll skip testing for it.
-    """
-    assert_table_lineage_equal(
-        "alter table tab1 swap with tab2",
-        {"tab2"},
-        {"tab1"},
-        dialect=dialect,
-        test_sqlparse=False,
-    )
-
-
 @pytest.mark.parametrize("dialect", ["databricks", "sparksql"])
 def test_refresh_table(dialect: str):
     assert_table_lineage_equal("refresh table tab1", None, None, dialect)
diff --git a/tests/test_sqlfluff.py b/tests/test_sqlfluff.py
index 6fde533..975426e 100644
--- a/tests/test_sqlfluff.py
+++ b/tests/test_sqlfluff.py
@@ -1,3 +1,5 @@
+import pytest
+
 from sqllineage.utils.entities import ColumnQualifierTuple
 from .helpers import assert_column_lineage_equal, assert_table_lineage_equal
 
@@ -104,6 +106,38 @@ FROM tab2 a"""
     )
 
 
+@pytest.mark.parametrize("dialect", ["snowflake", "bigquery"])
+def test_create_clone(dialect: str):
+    """
+    Language manual:
+        https://cloud.google.com/bigquery/docs/table-clones-create
+        https://docs.snowflake.com/en/sql-reference/sql/create-clone
+    Note clone is not a keyword in sqlparse, we'll skip testing for it.
+    """
+    assert_table_lineage_equal(
+        "create table tab2 CLONE tab1;",
+        {"tab1"},
+        {"tab2"},
+        dialect=dialect,
+        test_sqlparse=False,
+    )
+
+
+@pytest.mark.parametrize("dialect", ["snowflake"])
+def test_alter_table_swap_partition(dialect: str):
+    """
+    See https://docs.snowflake.com/en/sql-reference/sql/alter-table for language manual
+    Note swap is not a keyword in sqlparse, we'll skip testing for it.
+    """
+    assert_table_lineage_equal(
+        "alter table tab1 swap with tab2",
+        {"tab2"},
+        {"tab1"},
+        dialect=dialect,
+        test_sqlparse=False,
+    )
+
+
 # For top-level query parenthesis in DML, we don't treat it as subquery.
 # sqlparse has some problems identifying these subqueries.
 # note the table-level lineage works, only column-level lineage breaks for sqlparse
@@ -117,6 +151,20 @@ def test_create_as_with_parenthesis_around_both():
     assert_table_lineage_equal(sql, {"tab2"}, {"tab1"}, test_sqlparse=False)
 
 
+def test_cte_inside_bracket_of_insert():
+    sql = """INSERT INTO tab3 (WITH tab1 AS (SELECT * FROM tab2) SELECT * FROM tab1)"""
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("*", "tab2"),
+                ColumnQualifierTuple("*", "tab3"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
+
+
 # specify columns in CREATE statement, sqlparse would parse my_view as function call
 def test_view_with_subquery_custom_columns():
     # select as subquery