fix: missing target table for select into with union (#572)

reata · Jan 31, 2024 · cefe7bb · cefe7bb
1 parent e0fbf32
commit cefe7bb
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 12 deletions.
diff --git a/sqllineage/core/parser/sqlfluff/extractors/select.py b/sqllineage/core/parser/sqlfluff/extractors/select.py
@@ -59,12 +59,7 @@ def extract(
         self.extract_subquery(subqueries, holder)
 
         for segment in segments:
-            self._handle_swap_partition(segment, holder)
-            self._handle_select_into(segment, holder)
-            self.tables.extend(
-                self._list_table_from_from_clause_or_join_clause(segment, holder)
-            )
-            self._handle_column(segment)
+            self._handle_select_statement_child_segments(segment, holder)
 
             if is_set_expression(segment):
                 for idx, sub_segment in enumerate(
@@ -75,19 +70,24 @@ def extract(
                             (len(self.columns), len(self.tables))
                         )
                     for seg in list_child_segments(sub_segment):
-                        self.tables.extend(
-                            self._list_table_from_from_clause_or_join_clause(
-                                seg, holder
-                            )
-                        )
-                        self._handle_column(seg)
+                        self._handle_select_statement_child_segments(seg, holder)
 
         self.end_of_query_cleanup(holder)
 
         holder.expand_wildcard(self.metadata_provider)
 
         return holder
 
+    def _handle_select_statement_child_segments(
+        self, segment: BaseSegment, holder: SubQueryLineageHolder
+    ):
+        self._handle_swap_partition(segment, holder)
+        self._handle_select_into(segment, holder)
+        self.tables.extend(
+            self._list_table_from_from_clause_or_join_clause(segment, holder)
+        )
+        self._handle_column(segment)
+
     @staticmethod
     def _handle_swap_partition(segment: BaseSegment, holder: SubQueryLineageHolder):
         """

diff --git a/tests/sql/table/test_select_dialect_specific.py b/tests/sql/table/test_select_dialect_specific.py
@@ -60,6 +60,14 @@ def test_select_into(dialect: str):
     assert_table_lineage_equal(sql, {"films"}, {"films_recent"}, dialect=dialect)
 
 
+@pytest.mark.parametrize("dialect", ["postgres", "tsql"])
+def test_select_into_with_union(dialect: str):
+    sql = "SELECT * INTO films_all FROM films UNION ALL SELECT * FROM films_backup"
+    assert_table_lineage_equal(
+        sql, {"films", "films_backup"}, {"films_all"}, dialect=dialect
+    )
+
+
 @pytest.mark.parametrize("dialect", ["athena"])
 def test_select_from_unnest_with_ordinality(dialect: str):
     """