fix: misidentify-column-name-as-alias (#539)

Author: maoxingda

sqllineage/core/models.py

@@ -223,6 +223,7 @@ def _to_src_col(
                 else:
                     # select unqualified column
                     source = _to_src_col(src_col, None)
+                    setattr(source, "has_qualifier", False)
                     for table in set(alias_mapping.values()):
                         # in case of only one table, we get the right answer
                         # in case of multiple tables, a bunch of possible tables are set

sqllineage/core/parser/__init__.py

@@ -27,13 +27,15 @@ def end_of_query_cleanup(self, holder: SubQueryLineageHolder) -> None:
                 lateral_aliases = set()
                 for idx, tgt_col in enumerate(col_grp):
                     tgt_col.parent = tgt_tbl
-                    for lateral_alias_ref in col_grp[idx + 1 :]:  # noqa: E203
-                        if any(
-                            src_col[0] == tgt_col.raw_name
-                            for src_col in lateral_alias_ref.source_columns
-                        ):
-                            lateral_aliases.add(tgt_col.raw_name)
-                            break
+                    if hasattr(tgt_col, "has_alias") and tgt_col.has_alias is True:
+                        for lateral_alias_ref in col_grp[idx + 1 :]:  # noqa: E203
+                            if any(
+                                src_col[0] == tgt_col.raw_name
+                                for src_col in lateral_alias_ref.source_columns
+                                if src_col[1] is None
+                            ):
+                                lateral_aliases.add(tgt_col.raw_name)
+                                break
                     for src_col in tgt_col.to_source_columns(
                         holder.get_alias_mapping_from_table_group(tbl_grp)
                     ):
@@ -49,7 +51,9 @@ def end_of_query_cleanup(self, holder: SubQueryLineageHolder) -> None:
                             if wc.raw_name == "*":
                                 continue
                             if (
-                                src_col.raw_name == wc.raw_name
+                                hasattr(src_col, "has_qualifier")
+                                and src_col.has_qualifier is False
+                                and src_col.raw_name == wc.raw_name
                                 and src_col.raw_name in lateral_aliases
                             ):
                                 is_lateral_alias_ref = True

sqllineage/core/parser/sqlfluff/models.py

@@ -105,10 +105,12 @@ def of(column: BaseSegment, **kwargs) -> Column:
         if column.type == "select_clause_element":
             source_columns, alias = SqlFluffColumn._get_column_and_alias(column)
             if alias:
-                return Column(
+                alias_column = Column(
                     alias,
                     source_columns=source_columns,
                 )
+                setattr(alias_column, "has_alias", True)
+                return alias_column
             if source_columns:
                 column_name = None
                 for sub_segment in list_child_segments(column):

tests/sql/column/test_column_select_lateral_alias_ref.py

@@ -27,6 +27,37 @@ def test_column_top_level_lateral_ref():
                 ColumnQualifierTuple("id", "public.tgt_tbl1"),
             ),
         ],
+        test_sqlparse=False,
+    )
+    sql = """
+    insert into public.tgt_tbl1
+    (
+        name,
+        email
+    )
+    select
+        st1.name,
+        st1.name || st1.email || '@gmail.com' as email
+    from
+        public.src_tbl1 as st1
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("name", "public.src_tbl1"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("name", "public.src_tbl1"),
+                ColumnQualifierTuple("email", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("email", "public.src_tbl1"),
+                ColumnQualifierTuple("email", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
     )
 
 
@@ -60,6 +91,7 @@ def test_column_lateral_ref_within_subquery():
                 ColumnQualifierTuple("name", "public.tgt_tbl1"),
             ),
         ],
+        test_sqlparse=False,
     )
 
     sql = """
@@ -95,4 +127,5 @@ def test_column_lateral_ref_within_subquery():
                 ColumnQualifierTuple("name", "public.tgt_tbl1"),
             ),
         ],
+        test_sqlparse=False,
     )