Repository: reata/sqllineage

Issue:
Missing Source Table for MERGE statement when UNION involved in source subquery
version: sqllineage 1.4.2 

when i use LineageRunner to parse sql below, it turned out error like:"return name.strip("`").strip('"').strip("'") AttributeError: 'NoneType' object has no attribute 'strip'":

create or replace view scn.vwn as
SELECT 1 as file_type
FROM scn.tbn

LineageRunner parse error

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/417?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#417](https://app.codecov.io/gh/reata/sqllineage/pull/417?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (0f9e70e) into [master](https://app.codecov.io/gh/reata/sqllineage/commit/60219d9bfdb2d31b84e1f6747db1446df834c42f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (60219d9) will **decrease** coverage by `0.16%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master     #417      +/-   ##
==========================================
- Coverage   99.40%   99.24%   -0.16%     
==========================================
  Files          41       41              
  Lines        2003     1998       -5     
==========================================
- Hits         1991     1983       -8     
- Misses         12       15       +3     
```


| [Files Changed](https://app.codecov.io/gh/reata/sqllineage/pull/417?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [sqllineage/core/parser/sqlfluff/utils.py](https://app.codecov.io/gh/reata/sqllineage/pull/417?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL3BhcnNlci9zcWxmbHVmZi91dGlscy5weQ==) | `98.17% <100.00%> (-0.06%)` | :arrow_down: |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/reata/sqllineage/pull/417/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)



Patch:
diff --git a/sqllineage/core/parser/sqlfluff/handlers/source.py b/sqllineage/core/parser/sqlfluff/handlers/source.py
index ed24748..7294735 100644
--- a/sqllineage/core/parser/sqlfluff/handlers/source.py
+++ b/sqllineage/core/parser/sqlfluff/handlers/source.py
@@ -129,10 +129,6 @@ class SourceHandler(SourceHandlerMixin, ConditionalSegmentBaseHandler):
         elif first_segment.type == "bracketed" and is_values_clause(first_segment):
             # (VALUES ...) AS alias, no dataset involved
             return
-        elif is_union(segment):
-            subqueries = get_subqueries(segment)
-            subquery, alias = subqueries[0]
-            self.tables.append(SqlFluffSubQuery.of(subquery, alias))
         else:
             subqueries = get_subqueries(segment)
             if subqueries:
diff --git a/sqllineage/core/parser/sqlfluff/utils.py b/sqllineage/core/parser/sqlfluff/utils.py
index c11e863..d068350 100644
--- a/sqllineage/core/parser/sqlfluff/utils.py
+++ b/sqllineage/core/parser/sqlfluff/utils.py
@@ -259,33 +259,26 @@ def retrieve_segments(
     segment: BaseSegment, check_bracketed: bool = True
 ) -> List[BaseSegment]:
     """
-    Filter segments for a given segment's segments if they are not negligible
-    :param segment: segment to be processed
-    :param check_bracketed: process segment if it is of type "bracketed"
-    :return: a list of segments
-    """
-    if segment.type == "bracketed" and is_union(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
-            for sg in segment.iter_segments(expanding=["expression"], pass_through=True)
-        ]
-        result = []
-        for sgmnt in segments:
-            if sgmnt.type == "column_reference":
-                result.append(sgmnt)
-            else:
-                for sg in sgmnt.segments:
-                    if not is_segment_negligible(sg):
-                        result.append(sg)
-        return result
+    Filter segments for a given segment's children
+    Recursive goes into bracket by default
+    """
+    if segment.type == "bracketed" and check_bracketed:
+        if is_union(segment):
+            return [seg for seg in segment.segments if seg.type == "set_expression"]
+        else:
+            result = []
+            for seg in segment.iter_segments(
+                expanding=["expression"], pass_through=True
+            ):
+                if seg.type == "column_reference":
+                    result.append(seg)
+                else:
+                    for s in seg.segments:
+                        if not is_segment_negligible(s):
+                            result.append(s)
+            return result
     else:
-        return [sgmnt for sgmnt in segment.segments if not is_segment_negligible(sgmnt)]
+        return [seg for seg in segment.segments if not is_segment_negligible(seg)]
 
 
 def get_identifier(col_segment: BaseSegment) -> str:
@@ -404,19 +397,8 @@ def get_statement_segment(parsed_string: ParsedString) -> BaseSegment:
 
 
 def is_union(segment: BaseSegment) -> bool:
-    """
-    :param segment: segment to be processed
-    :return: True if the segment contains 'UNION' or 'UNION ALL' keyword
-    """
-    return (
-        len(
-            [
-                s
-                for s in segment.raw_segments
-                if (s.raw_upper == "UNION" or s.raw_upper == "UNION ALL")
-            ]
-        )
-        > 0
+    return segment.type == "set_expression" or any(
+        seg.type == "set_expression" for seg in segment.segments
     )
 
 


Test patch:
diff --git a/tests/sql/table/test_merge.py b/tests/sql/table/test_merge.py
index 2d7a87a..329c755 100644
--- a/tests/sql/table/test_merge.py
+++ b/tests/sql/table/test_merge.py
@@ -43,3 +43,18 @@ def test_merge_into_insert_one_column():
 USING src ON target.k = src.k
 WHEN NOT MATCHED THEN INSERT VALUES (src.k)"""
     assert_table_lineage_equal(sql, {"src"}, {"target"})
+
+
+def test_merge_with_union_in_subquery_and_join():
+    sql = """MERGE INTO tgt t
+USING (SELECT s1.id, baz.value
+       FROM (SELECT id
+             FROM foo
+             UNION ALL
+             SELECT id
+             FROM bar) s1
+                CROSS JOIN baz) s
+ON t.id = s.id
+WHEN MATCHED THEN UPDATE SET t.value = s.value
+"""
+    assert_table_lineage_equal(sql, {"foo", "bar", "baz"}, {"tgt"})
diff --git a/tests/sql/table/test_select.py b/tests/sql/table/test_select.py
index a17c860..33bdeed 100644
--- a/tests/sql/table/test_select.py
+++ b/tests/sql/table/test_select.py
@@ -221,9 +221,9 @@ def test_select_from_unnest_parsed_as_keyword():
 
 def test_table_name_case_insensitive():
     assert_table_lineage_equal(
-        """select * from tab_b
-union all
-select * from TAB_B""",
+        """SELECT * FROM tab_b
+UNION ALL
+SELECT * FROM TAB_B""",
         {"tab_b"},
     )