Repository: reata/sqllineage
Pull request #490 (2023-12-09 13:02:42)

fix: union and subquery case
making progress in #481 

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/490?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`5d5ebec`)](https://app.codecov.io/gh/reata/sqllineage/commit/5d5ebec78669a565a99bec91a54b80aa70c11a6f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) 99.40% compared to head [(`00c8905`)](https://app.codecov.io/gh/reata/sqllineage/pull/490?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) 99.40%.


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master     #490   +/-   ##
=======================================
  Coverage   99.40%   99.40%           
=======================================
  Files          40       40           
  Lines        2031     2031           
=======================================
  Hits         2019     2019           
  Misses         12       12           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/reata/sqllineage/pull/490?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None).


Patch:
diff --git a/sqllineage/core/parser/sqlfluff/extractors/select.py b/sqllineage/core/parser/sqlfluff/extractors/select.py
index c431a56..72c1eda 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/select.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/select.py
@@ -56,7 +56,20 @@ class SelectExtractor(BaseExtractor, SourceHandlerMixin):
             self._handle_select_into(segment, holder)
             self._handle_table(segment, holder)
             self._handle_column(segment)
-            self._handle_set(segment, holder)
+
+            if is_set_expression(segment):
+                for idx, sub_segment in enumerate(
+                    segment.get_children("select_statement", "bracketed")
+                ):
+                    if idx != 0:
+                        self.union_barriers.append(
+                            (len(self.columns), len(self.tables))
+                        )
+                    for seg in list_child_segments(sub_segment):
+                        for sq in self.list_subquery(seg):
+                            subqueries.append(sq)
+                        self._handle_table(seg, holder)
+                        self._handle_column(seg)
 
         self.end_of_query_cleanup(holder)
 
@@ -135,20 +148,6 @@ class SelectExtractor(BaseExtractor, SourceHandlerMixin):
             for sub_segment in segment.get_children("select_clause_element"):
                 self.columns.append(SqlFluffColumn.of(sub_segment))
 
-    def _handle_set(self, segment: BaseSegment, holder: SubQueryLineageHolder) -> None:
-        """
-        set handler method
-        """
-        if is_set_expression(segment):
-            for idx, sub_segment in enumerate(
-                segment.get_children("select_statement", "bracketed")
-            ):
-                if idx != 0:
-                    self.union_barriers.append((len(self.columns), len(self.tables)))
-                for seg in list_child_segments(sub_segment):
-                    self._handle_table(seg, holder)
-                    self._handle_column(seg)
-
     def _add_dataset_from_expression_element(
         self, segment: BaseSegment, holder: SubQueryLineageHolder
     ) -> None:


Test patch:
diff --git a/tests/sql/column/test_column_select_union.py b/tests/sql/column/test_column_select_union.py
index fb640b3..681041d 100644
--- a/tests/sql/column/test_column_select_union.py
+++ b/tests/sql/column/test_column_select_union.py
@@ -59,3 +59,27 @@ SELECT col1 FROM dataset.tab2) SELECT col1 FROM temp_cte"""
             ),
         ],
     )
+
+
+def test_union_with_subquery():
+    sql = """INSERT INTO tab3
+SELECT sq1.id
+FROM (SELECT id
+      FROM tab1) sq1
+UNION ALL
+SELECT sq2.id
+FROM (SELECT id
+      FROM tab2) sq2"""
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("id", "tab1"),
+                ColumnQualifierTuple("id", "tab3"),
+            ),
+            (
+                ColumnQualifierTuple("id", "tab2"),
+                ColumnQualifierTuple("id", "tab3"),
+            ),
+        ],
+    )
diff --git a/tests/sql/table/test_select.py b/tests/sql/table/test_select.py
index 2a0b90c..8cd7730 100644
--- a/tests/sql/table/test_select.py
+++ b/tests/sql/table/test_select.py
@@ -271,3 +271,14 @@ ORDER BY col1"""
         sql,
         {"tab1", "tab2", "tab3"},
     )
+
+
+def test_select_union_all_with_nested_subquery():
+    sql = """SELECT sq1.id
+FROM (SELECT id
+      FROM tab1) sq1
+UNION ALL
+SELECT sq2.id
+FROM (SELECT id
+      FROM tab2) sq2"""
+    assert_table_lineage_equal(sql, {"tab1", "tab2"})