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

reata · Jan 9, 2024 · 37ea08a · 37ea08a
1 parent 0326b59
commit 37ea08a
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 7 deletions.
diff --git a/sqllineage/core/parser/__init__.py b/sqllineage/core/parser/__init__.py
@@ -27,13 +27,14 @@ 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:
+                        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
                     for src_col in tgt_col.to_source_columns(
                         holder.get_alias_mapping_from_table_group(tbl_grp)
                     ):

diff --git a/sqllineage/core/parser/sqlfluff/models.py b/sqllineage/core/parser/sqlfluff/models.py
@@ -107,6 +107,7 @@ def of(column: BaseSegment, **kwargs) -> Column:
             if alias:
                 return Column(
                     alias,
+                    has_alias=True,
                     source_columns=source_columns,
                 )
             if source_columns:

diff --git a/tests/sql/column/test_column_select_lateral_alias_ref.py b/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,
     )