Repository: reata/sqllineage

Issue:
qualified wildcard recognized as wrong column name
The following SQL query uses "sales" as an alias for "T_SALES". How could I ask the `LineageRunner()` to return the table/view name rather than the alias?

```sql
-- sample_query_with_semicolons.sql

SELECT CUST_ID, CUST_NAME, CUST_CITY, CUST_COUNTRY
INTO #TEMP_CUSTOMERS
FROM T_CUSTOMERS
WHERE CUST_COUNTRY = 'USA';

SELECT CUST.*
    , SALES.CUM_SALES
INTO #TEMP_FINAL_RESULTS
FROM #TEMP_CUSTOMERS AS CUST
LEFT JOIN T_SALES AS SALES
ON CUST.CUST_ID = SALES.CUST_ID
WHERE SALES.CUST_ID IS NULL;

SELECT *
FROM #TEMP_FINAL_RESULTS;
```

```python
from sqllineage.runner import LineageRunner
import csv
import pandas as pd

# read text from sample_query.sql
sql_script = open('sample_query_with_semicolons.sql', 'r').read()

# parse sql_script with LineageRunner
parsed_results = LineageRunner(sql_script, dialect="tsql")

# write parsed_results to data frame
df = pd.DataFrame(parsed_results.get_column_lineage())

df.columns = ['Source', 'Target']

print(df)
```

Expected output:

Source                |                   Target
---|---
`<default>.t_customers.cust_city`    |  `<default>.#temp_customers.cust_city`
`<default>.t_customers.cust_country` |  `<default>.#temp_customers.cust_country`
`<default>.t_customers.cust_id`      |  `<default>.#temp_customers.cust_id`
`<default>.t_customers.cust_name`   |   `<default>.#temp_customers.cust_name`
`<default>.`**t_sales.cum_sales**  | `<default>.#temp_final_results.cum_sales`
`<default>.#temp_customers.cust.*`   |  `<default>.#temp_final_results.cust.*`


Actual output:

 Source                |                   Target
---|---
`<default>.t_customers.cust_city`    |  `<default>.#temp_customers.cust_city`
`<default>.t_customers.cust_country` |  `<default>.#temp_customers.cust_country`
`<default>.t_customers.cust_id`      |  `<default>.#temp_customers.cust_id`
`<default>.t_customers.cust_name`   |   `<default>.#temp_customers.cust_name`
`<default>.sales.cum_sales`  | `<default>.#temp_final_results.cum_sales`
`<default>.#temp_customers.cust.*`   |  `<default>.#temp_final_results.cust.*`

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

```diff
@@            Coverage Diff             @@
##           master     #430      +/-   ##
==========================================
- Coverage   99.54%   99.39%   -0.15%     
==========================================
  Files          41       41              
  Lines        1972     1984      +12     
==========================================
+ Hits         1963     1972       +9     
- Misses          9       12       +3     
```


| [Files Changed](https://app.codecov.io/gh/reata/sqllineage/pull/430?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [.../parser/sqlfluff/extractors/dml\_merge\_extractor.py](https://app.codecov.io/gh/reata/sqllineage/pull/430?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL3BhcnNlci9zcWxmbHVmZi9leHRyYWN0b3JzL2RtbF9tZXJnZV9leHRyYWN0b3IucHk=) | `100.00% <100.00%> (ø)` | |
| [sqllineage/core/parser/sqlfluff/models.py](https://app.codecov.io/gh/reata/sqllineage/pull/430?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL3BhcnNlci9zcWxmbHVmZi9tb2RlbHMucHk=) | `100.00% <100.00%> (ø)` | |
| [sqllineage/core/parser/sqlfluff/utils.py](https://app.codecov.io/gh/reata/sqllineage/pull/430?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL3BhcnNlci9zcWxmbHVmZi91dGlscy5weQ==) | `99.31% <100.00%> (-0.69%)` | :arrow_down: |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/reata/sqllineage/pull/430/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/extractors/dml_merge_extractor.py b/sqllineage/core/parser/sqlfluff/extractors/dml_merge_extractor.py
index 2495c82..89bbba8 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/dml_merge_extractor.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/dml_merge_extractor.py
@@ -13,6 +13,7 @@ from sqllineage.core.parser.sqlfluff.extractors.lineage_holder_extractor import
 )
 from sqllineage.core.parser.sqlfluff.models import SqlFluffSubQuery, SqlFluffTable
 from sqllineage.core.parser.sqlfluff.utils import (
+    extract_column_qualifier,
     extract_identifier,
     extract_innermost_bracketed,
     get_child,
@@ -93,11 +94,15 @@ class DmlMergeExtractor(LineageHolderExtractor):
                         for set_clause in get_children(set_clause_list, "set_clause"):
                             columns = get_children(set_clause, "column_reference")
                             if len(columns) == 2:
-                                src_col = Column(extract_identifier(columns[1]))
-                                src_col.parent = direct_source
-                                tgt_col = Column(extract_identifier(columns[0]))
-                                tgt_col.parent = list(holder.write)[0]
-                                holder.add_column_lineage(src_col, tgt_col)
+                                src_col = tgt_col = None
+                                if src_cqt := extract_column_qualifier(columns[1]):
+                                    src_col = Column(src_cqt.column)
+                                    src_col.parent = direct_source
+                                if tgt_cqt := extract_column_qualifier(columns[0]):
+                                    tgt_col = Column(tgt_cqt.column)
+                                    tgt_col.parent = list(holder.write)[0]
+                                if src_col is not None and tgt_col is not None:
+                                    holder.add_column_lineage(src_col, tgt_col)
             for merge_when_not_matched_clause in get_children(
                 merge_match, "merge_when_not_matched_clause"
             ):
@@ -107,20 +112,22 @@ class DmlMergeExtractor(LineageHolderExtractor):
                 insert_columns = []
                 if bracketed := get_child(merge_insert, "bracketed"):
                     for column_reference in get_children(bracketed, "column_reference"):
-                        tgt_col = Column(extract_identifier(column_reference))
-                        tgt_col.parent = list(holder.write)[0]
-                        insert_columns.append(tgt_col)
+                        if cqt := extract_column_qualifier(column_reference):
+                            tgt_col = Column(cqt.column)
+                            tgt_col.parent = list(holder.write)[0]
+                            insert_columns.append(tgt_col)
                     if values_clause := get_child(merge_insert, "values_clause"):
                         if bracketed := get_child(values_clause, "bracketed"):
                             for j, e in enumerate(
                                 get_children(bracketed, "literal", "expression")
                             ):
                                 if column_reference := get_child(e, "column_reference"):
-                                    src_col = Column(
-                                        extract_identifier(column_reference)
-                                    )
-                                    src_col.parent = direct_source
-                                    holder.add_column_lineage(
-                                        src_col, insert_columns[j]
-                                    )
+                                    if cqt := extract_column_qualifier(
+                                        column_reference
+                                    ):
+                                        src_col = Column(cqt.column)
+                                        src_col.parent = direct_source
+                                        holder.add_column_lineage(
+                                            src_col, insert_columns[j]
+                                        )
         return holder
diff --git a/sqllineage/core/parser/sqlfluff/models.py b/sqllineage/core/parser/sqlfluff/models.py
index 10fb704..6c60309 100644
--- a/sqllineage/core/parser/sqlfluff/models.py
+++ b/sqllineage/core/parser/sqlfluff/models.py
@@ -6,6 +6,7 @@ from sqlfluff.core.parser import BaseSegment
 from sqllineage import SQLPARSE_DIALECT
 from sqllineage.core.models import Column, Schema, SubQuery, Table
 from sqllineage.core.parser.sqlfluff.utils import (
+    extract_column_qualifier,
     extract_identifier,
     is_subquery,
     is_wildcard,
@@ -111,8 +112,11 @@ class SqlFluffColumn(Column):
             if source_columns:
                 column_name = None
                 for sub_segment in list_child_segments(column):
-                    if sub_segment.type == "column_reference":
-                        column_name = extract_identifier(sub_segment)
+                    if sub_segment.type == "column_reference" or is_wildcard(
+                        sub_segment
+                    ):
+                        if cqt := extract_column_qualifier(sub_segment):
+                            column_name = cqt.column
                     elif sub_segment.type == "expression":
                         # special handling for postgres style type cast, col as target column name instead of col::type
                         if len(sub2_segments := list_child_segments(sub_segment)) == 1:
@@ -130,7 +134,10 @@ class SqlFluffColumn(Column):
                                     if (
                                         sub3_segment := sub3_segments[0]
                                     ).type == "column_reference":
-                                        column_name = extract_identifier(sub3_segment)
+                                        if cqt := extract_column_qualifier(
+                                            sub3_segment
+                                        ):
+                                            column_name = cqt.column
                 return Column(
                     column.raw if column_name is None else column_name,
                     source_columns=source_columns,
@@ -149,12 +156,11 @@ class SqlFluffColumn(Column):
         :param segment: segment to be processed
         :return: list of extracted source columns
         """
-        if segment.type == "identifier" or is_wildcard(segment):
-            return [ColumnQualifierTuple(segment.raw, None)]
-        if segment.type == "column_reference":
-            parent, column = SqlFluffColumn._get_column_and_parent(segment)
-            return [ColumnQualifierTuple(column, parent)]
-        if segment.type in NON_IDENTIFIER_OR_COLUMN_SEGMENT_TYPE:
+        col_list = []
+        if segment.type in ("identifier", "column_reference") or is_wildcard(segment):
+            if cqt := extract_column_qualifier(segment):
+                col_list = [cqt]
+        elif segment.type in NON_IDENTIFIER_OR_COLUMN_SEGMENT_TYPE:
             sub_segments = list_child_segments(segment)
             col_list = []
             for sub_segment in sub_segments:
@@ -172,8 +178,7 @@ class SqlFluffColumn(Column):
                 ):
                     res = SqlFluffColumn._extract_source_columns(sub_segment)
                     col_list.extend(res)
-            return col_list
-        return []
+        return col_list
 
     @staticmethod
     def _get_column_from_subquery(
@@ -229,12 +234,4 @@ class SqlFluffColumn(Column):
             ):
                 res = SqlFluffColumn._extract_source_columns(sub_segment)
                 columns += res if res else []
-
         return columns, alias
-
-    @staticmethod
-    def _get_column_and_parent(col_segment: BaseSegment) -> Tuple[Optional[str], str]:
-        identifiers = list_child_segments(col_segment)
-        if len(identifiers) > 1:
-            return identifiers[-2].raw, identifiers[-1].raw
-        return None, identifiers[-1].raw
diff --git a/sqllineage/core/parser/sqlfluff/utils.py b/sqllineage/core/parser/sqlfluff/utils.py
index 0697bc1..8957494 100644
--- a/sqllineage/core/parser/sqlfluff/utils.py
+++ b/sqllineage/core/parser/sqlfluff/utils.py
@@ -12,7 +12,7 @@ from typing import List, Optional, Tuple
 
 from sqlfluff.core.parser import BaseSegment
 
-from sqllineage.utils.entities import SubQueryTuple
+from sqllineage.utils.entities import ColumnQualifierTuple, SubQueryTuple
 
 
 def is_negligible(segment: BaseSegment) -> bool:
@@ -228,6 +228,23 @@ def extract_as_and_target_segment(
     return as_segment, target
 
 
+def extract_column_qualifier(segment: BaseSegment) -> Optional[ColumnQualifierTuple]:
+    cqt = None
+    if is_wildcard(segment):
+        identifiers = segment.raw.split(".")
+        column = identifiers[-1]
+        parent = identifiers[-2] if len(identifiers) > 1 else None
+        cqt = ColumnQualifierTuple(column, parent)
+    elif segment.type == "column_reference":
+        sub_segments = list_child_segments(segment)
+        column = sub_segments[-1].raw
+        parent = sub_segments[-2].raw if len(sub_segments) > 1 else None
+        cqt = ColumnQualifierTuple(column, parent)
+    elif segment.type == "identifier":
+        cqt = ColumnQualifierTuple(segment.raw, None)
+    return cqt
+
+
 def extract_innermost_bracketed(bracketed_segment: BaseSegment) -> BaseSegment:
     # in case of subquery in nested parenthesis like: `SELECT * FROM ((table))`, find the innermost one first
     while True:


Test patch:
diff --git a/tests/sql/column/test_column_select_column.py b/tests/sql/column/test_column_select_column.py
index bced8c7..f836a02 100644
--- a/tests/sql/column/test_column_select_column.py
+++ b/tests/sql/column/test_column_select_column.py
@@ -48,6 +48,27 @@ FROM tab2 a
     )
 
 
+def test_select_column_wildcard_with_qualifier():
+    sql = """INSERT INTO tab1
+SELECT tab2.*
+FROM tab2 a
+         INNER JOIN tab3 b
+                    ON a.id = b.id"""
+    assert_column_lineage_equal(
+        sql,
+        [(ColumnQualifierTuple("*", "tab2"), ColumnQualifierTuple("*", "tab1"))],
+    )
+    sql = """INSERT INTO tab1
+SELECT a.*
+FROM tab2 a
+         INNER JOIN tab3 b
+                    ON a.id = b.id"""
+    assert_column_lineage_equal(
+        sql,
+        [(ColumnQualifierTuple("*", "tab2"), ColumnQualifierTuple("*", "tab1"))],
+    )
+
+
 def test_select_distinct_column():
     sql = """INSERT INTO tab1
 SELECT DISTINCT col1