Repository: reata/sqllineage

Issue:
Support Lateral Column Alias Reference Analyzing
**SQL**
```sql
insert into public.tgt_tbl1
(
    id
)
select
    sq.id
from
    (
        select
            name      as user_name,
            user_name as id -- backward reference
        from
            public.src_tbl1
    ) sq
;
```

**To Reproduce**
*Note here we refer to SQL provided in prior step as stored in a file named `test.sql`*

```python
from sqllineage.runner import LineageRunner
with open("test.sql") as f:
    sql = f.read()
lr = LineageRunner(sql, dialect='redshift')
lr.print_column_lineage()
```
```
public.tgt_tbl1.id <- sq.id <- public.src_tbl1.user_name
```

**Expected behavior**
```
public.tgt_tbl1.id <- sq.id <- public.src_tbl1.name
```
**Python version (available via `python --version`)**
 - 3.11.5

**SQLLineage version (available via `sqllineage --version`):**
 - 1.4.9


Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/521?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`b620d78`)](https://app.codecov.io/gh/reata/sqllineage/commit/b620d78f34224737071510f0368b9eee588ad4ce?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) 99.49% compared to head [(`f0efc8a`)](https://app.codecov.io/gh/reata/sqllineage/pull/521?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) 99.49%.


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master     #521   +/-   ##
=======================================
  Coverage   99.49%   99.49%           
=======================================
  Files          41       41           
  Lines        2170     2188   +18     
=======================================
+ Hits         2159     2177   +18     
  Misses         11       11           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/reata/sqllineage/pull/521?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None).

maoxingda: Thank you for your response. I'm not sure if I can handle it, but I'll give it a try.
reata: We can also merge this and not closing #507, if top-level lateral column alias takes too much effort.

Just let me know your conclusion. Thanks for taking time investigating.
maoxingda: If it can be like this, that would be the best.
maoxingda: I will try to expand at top level.
maoxingda: > We can also merge this and not closing #507, if top-level lateral column alias takes too much effort.
> 
> Just let me know your conclusion. Thanks for taking time investigating.

It seems that it doesn't require too much effort.

Patch:
diff --git a/sqllineage/core/holders.py b/sqllineage/core/holders.py
index ad74141..fe663d0 100644
--- a/sqllineage/core/holders.py
+++ b/sqllineage/core/holders.py
@@ -153,7 +153,7 @@ class SubQueryLineageHolder(ColumnLineageMixin):
             for column in self.write_columns:
                 if column.raw_name == "*":
                     tgt_wildcard = column
-                    for src_wildcard in self._get_source_columns(tgt_wildcard):
+                    for src_wildcard in self.get_source_columns(tgt_wildcard):
                         if source_table := src_wildcard.parent:
                             src_table_columns = []
                             if isinstance(source_table, SubQuery):
@@ -178,7 +178,7 @@ class SubQueryLineageHolder(ColumnLineageMixin):
             table = next(iter(write_only))
         return table
 
-    def _get_source_columns(self, node: Column) -> List[Column]:
+    def get_source_columns(self, node: Column) -> List[Column]:
         return [
             src
             for (src, tgt, edge_type) in self.graph.in_edges(nbunch=node, data="type")
diff --git a/sqllineage/core/models.py b/sqllineage/core/models.py
index 610571b..85457d4 100644
--- a/sqllineage/core/models.py
+++ b/sqllineage/core/models.py
@@ -147,7 +147,7 @@ class Column:
         """
         self._parent: Set[Union[Path, Table, SubQuery]] = set()
         self.raw_name = escape_identifier_name(name)
-        self.source_columns = (
+        self.source_columns = [
             (
                 escape_identifier_name(raw_name),
                 escape_identifier_name(qualifier) if qualifier is not None else None,
@@ -155,7 +155,7 @@ class Column:
             for raw_name, qualifier in kwargs.pop(
                 "source_columns", ((self.raw_name, None),)
             )
-        )
+        ]
 
     def __str__(self):
         return (
diff --git a/sqllineage/core/parser/__init__.py b/sqllineage/core/parser/__init__.py
index 73d24af..f9c7ea1 100644
--- a/sqllineage/core/parser/__init__.py
+++ b/sqllineage/core/parser/__init__.py
@@ -25,8 +25,16 @@ class SourceHandlerMixin:
                 if len(holder.write) > 1:
                     raise SQLLineageException
                 tgt_tbl = list(holder.write)[0]
+                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
                     for src_col in tgt_col.to_source_columns(
                         self.get_alias_mapping_from_table_group(tbl_grp, holder)
                     ):
@@ -37,6 +45,22 @@ class SourceHandlerMixin:
                             # for invalid query: create view test (col3, col4) select col1 as col2 from tab,
                             # when the length doesn't match, we fall back to default behavior
                             tgt_col = write_columns[idx]
+                        is_lateral_alias_ref = False
+                        for wc in holder.write_columns:
+                            if wc.raw_name == "*":
+                                continue
+                            if (
+                                src_col.raw_name == wc.raw_name
+                                and src_col.raw_name in lateral_aliases
+                            ):
+                                is_lateral_alias_ref = True
+                                for lateral_alias_col in holder.get_source_columns(wc):
+                                    holder.add_column_lineage(
+                                        lateral_alias_col, tgt_col
+                                    )
+                                break
+                        if is_lateral_alias_ref:
+                            continue
                         holder.add_column_lineage(src_col, tgt_col)
 
     @classmethod


Test patch:
diff --git a/tests/sql/column/test_column_select_lateral_alias_ref.py b/tests/sql/column/test_column_select_lateral_alias_ref.py
new file mode 100644
index 0000000..28b7dc8
--- /dev/null
+++ b/tests/sql/column/test_column_select_lateral_alias_ref.py
@@ -0,0 +1,98 @@
+from sqllineage.utils.entities import ColumnQualifierTuple
+from ...helpers import assert_column_lineage_equal
+
+
+def test_column_top_level_lateral_ref():
+    sql = """
+    insert into public.tgt_tbl1
+    select
+        name               as user_name,
+        user_name || email as id         -- lateral ref
+    from
+        public.src_tbl1
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("name", "public.src_tbl1"),
+                ColumnQualifierTuple("user_name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("name", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("email", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+        ],
+    )
+
+
+def test_column_lateral_ref_within_subquery():
+    sql = """
+    insert into public.tgt_tbl1
+    select
+        sq.name
+    from
+        (
+            select
+                id || name as alias1,
+                alias1 || email as name
+            from
+                public.src_tbl1
+        ) as sq
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("name", "public.src_tbl1"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("email", "public.src_tbl1"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+        ],
+    )
+
+    sql = """
+    insert into public.tgt_tbl1
+    select
+        sq.name
+    from
+        (
+            select
+                st1.id || st1.name as alias1,
+                alias1 || st2.email as name
+            from
+                public.src_tbl1 as st1
+                join
+                    public.src_tbl2 as st2
+                on
+                    st1.id = st2.id
+        ) as sq
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("name", "public.src_tbl1"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("email", "public.src_tbl2"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+        ],
+    )