From 93c6ba0ce57db151b78acc7f7087d700d9537872 Mon Sep 17 00:00:00 2001
From: kkozhakin <kkozhakin@mail.ru>
Date: Mon, 5 Aug 2024 10:42:52 +0300
Subject: [PATCH] feat: del wildcard only with metadata_provider (#612)

---
 sqllineage/core/holders.py                    |  6 +++--
 tests/sql/column/test_column_select_column.py | 26 +++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/sqllineage/core/holders.py b/sqllineage/core/holders.py
index 8a7086d9..4d6d8c53 100644
--- a/sqllineage/core/holders.py
+++ b/sqllineage/core/holders.py
@@ -185,6 +185,7 @@ def expand_wildcard(self, metadata_provider: MetaDataProvider) -> None:
                                     src_table_columns,
                                     tgt_wildcard,
                                     src_wildcard,
+                                    metadata_provider,
                                 )
 
     def get_alias_mapping_from_table_group(
@@ -227,6 +228,7 @@ def _replace_wildcard(
         src_table_columns: List[Column],
         tgt_wildcard: Column,
         src_wildcard: Column,
+        metadata_provider: MetaDataProvider,
     ) -> None:
         target_columns = self.get_table_columns(tgt_table)
         for src_col in src_table_columns:
@@ -238,9 +240,9 @@ def _replace_wildcard(
             self.graph.add_edge(src_col.parent, src_col, type=EdgeType.HAS_COLUMN)
             self.graph.add_edge(src_col, new_column, type=EdgeType.LINEAGE)
         # remove wildcard
-        if self.graph.has_node(tgt_wildcard):
+        if metadata_provider and self.graph.has_node(tgt_wildcard):
             self.graph.remove_node(tgt_wildcard)
-        if self.graph.has_node(src_wildcard):
+        if metadata_provider and self.graph.has_node(src_wildcard):
             self.graph.remove_node(src_wildcard)
 
 
diff --git a/tests/sql/column/test_column_select_column.py b/tests/sql/column/test_column_select_column.py
index c946aa4b..5d1bcff5 100644
--- a/tests/sql/column/test_column_select_column.py
+++ b/tests/sql/column/test_column_select_column.py
@@ -48,6 +48,32 @@ def test_select_column_wildcard():
     )
 
 
+# TODO
+# Expected Lineage: {(Column: <default>.t2.f1, Column: <default>.result_table.f1), ...}
+# Actual Lineage: {(Column: t2.f1, Column: <default>.result_table.f1), ...}
+
+# def test_select_column_wildcard_with_fields():
+#     sql = """
+#         WITH t2 AS (SELECT * FROM t1),
+#              t3 AS (SELECT f1, * FROM t2)
+#         INSERT INTO result_table
+#         SELECT *
+#         FROM t3;"""
+#     assert_column_lineage_equal(
+#         sql,
+#         [
+#             (
+#                 ColumnQualifierTuple("f1", "t2"),
+#                 ColumnQualifierTuple("f1", "result_table"),
+#             ),
+#             (
+#                 ColumnQualifierTuple("*", "t1"),
+#                 ColumnQualifierTuple("*", "result_table"),
+#             ),
+#         ],
+#     )
+
+
 def test_select_column_wildcard_with_qualifier():
     sql = """INSERT INTO tab1
 SELECT tab2.*