Repository: reata/sqllineage
Pull request #391 (2023-06-15 06:43:20)

parse union inside cte
fix issues #376 
commit [bd16d45](https://github.com/reata/sqllineage/pull/391/commits/bd16d45b885511e7b3d3f6b4585f333d3d69fa6e): add check union direct inside cte, we current just have check union in raw segment
commit [94672b3](https://github.com/reata/sqllineage/pull/391/commits/94672b33bbd93aaf253d195f479be322a6b1852e): use union direct method to check union inside cte, then add it to queue for parsing
commit [d8eb39d](https://github.com/reata/sqllineage/pull/391/commits/d8eb39df4f3d094df7ee6dd29f9d1f81cf7aab02): add test for this issue


reata: Would you fix the CI? You can run it in your local env via `tox -e py38`. Or change to your local Python version like py39 or py310. You can rely on GitHub Actions for multiple Python version test though if you don't have to set it up locally.
duongphannamhung: I just fix the CI, thank you
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/391?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#391](https://app.codecov.io/gh/reata/sqllineage/pull/391?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (56b3d87) into [master](https://app.codecov.io/gh/reata/sqllineage/commit/341435ac413475c9e600d47762a1e7e441f4dd43?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (341435a) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master     #391   +/-   ##
=======================================
  Coverage   99.40%   99.40%           
=======================================
  Files          42       42           
  Lines        2006     2013    +7     
=======================================
+ Hits         1994     2001    +7     
  Misses         12       12           
```


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


duongphannamhung: Really appreciate your comments, to be `is_isomorphic` like sqlparse makes me rethink the way my commit work so I have a little few changes, hope it's working out

Patch:
diff --git a/sqllineage/core/parser/sqlfluff/extractors/lineage_holder_extractor.py b/sqllineage/core/parser/sqlfluff/extractors/lineage_holder_extractor.py
index 337a56b..c99d8ca 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/lineage_holder_extractor.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/lineage_holder_extractor.py
@@ -15,7 +15,6 @@ from sqllineage.core.parser.sqlfluff.utils import (
     get_multiple_identifiers,
     get_subqueries,
     is_subquery,
-    is_union,
 )
 from sqllineage.utils.entities import SubQueryTuple
 
@@ -74,7 +73,7 @@ class LineageHolderExtractor:
             )
         if segment.type in ["select_clause", "from_clause", "where_clause"]:
             result = cls._parse_subquery(get_subqueries(segment))
-        elif is_subquery(segment) and not is_union(segment):
+        elif is_subquery(segment):
             # Parenthesis for SubQuery without alias, this is valid syntax for certain SQL dialect
             result = [SqlFluffSubQuery.of(segment, None)]
         return result
diff --git a/sqllineage/core/parser/sqlfluff/utils.py b/sqllineage/core/parser/sqlfluff/utils.py
index 2ca660e..10268e3 100644
--- a/sqllineage/core/parser/sqlfluff/utils.py
+++ b/sqllineage/core/parser/sqlfluff/utils.py
@@ -87,8 +87,6 @@ def get_subqueries(segment: BaseSegment) -> List[SubQueryTuple]:
             return [SubQueryTuple(get_innermost_bracketed(bracketed_segments[0]), None)]
         return []
     elif is_union(segment):
-        if segment.type != "set_expression":
-            segment = segment.get_child("set_expression")
         for s in retrieve_segments(segment, check_bracketed=True):
             if s.type == "bracketed" or s.type == "select_statement":
                 subquery.append(SubQueryTuple(s, None))
@@ -262,7 +260,11 @@ def retrieve_segments(
     :return: a list of segments
     """
     if segment.type == "bracketed" and is_union(segment):
-        return [segment]
+        result = []
+        for sgmt in segment.segments:
+            if sgmt.type == "set_expression":
+                result = [sgmt]
+        return result
     elif segment.type == "bracketed" and check_bracketed:
         segments = [
             sg


Test patch:
diff --git a/tests/test_columns.py b/tests/test_columns.py
index d4bba13..ebe7e21 100644
--- a/tests/test_columns.py
+++ b/tests/test_columns.py
@@ -1082,3 +1082,21 @@ WHEN NOT MATCHED THEN INSERT (k, v) VALUES (b.k, b.v_max)"""
             (ColumnQualifierTuple("k", "src"), ColumnQualifierTuple("k", "target")),
         ],
     )
+
+
+def test_union_inside_cte():
+    sql = """INSERT INTO dataset.target WITH temp_cte AS (SELECT col1 FROM dataset.tab1 UNION ALL
+SELECT col1 FROM dataset.tab2) SELECT col1 FROM temp_cte"""
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("col1", "dataset.tab1"),
+                ColumnQualifierTuple("col1", "dataset.target"),
+            ),
+            (
+                ColumnQualifierTuple("col1", "dataset.tab2"),
+                ColumnQualifierTuple("col1", "dataset.target"),
+            ),
+        ],
+    )