Repository: reata/sqllineage

Issue:
Make Lateral Column Alias Reference Configurable
**SQL**
```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
```

**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()
```

**Actual behavior**
```
public.tgt_tbl1.email <- public.src_tbl1.email
```

**Expected behavior**
```
public.tgt_tbl1.name <- public.src_tbl1.name
public.tgt_tbl1.name <- public.src_tbl1.email
public.tgt_tbl1.email <- public.src_tbl1.email
```

**Python version (available via `python --version`)**
 - 3.10

**SQLLineage version (available via `sqllineage --version`):**
 - 1.5.0

**Additional context**


Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/540?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 [(`0326b59`)](https://app.codecov.io/gh/reata/sqllineage/commit/0326b59c18e6ad926e183973083de9b94a52193c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) 99.49% compared to head [(`d3e987d`)](https://app.codecov.io/gh/reata/sqllineage/pull/540?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) 99.49%.

> :exclamation: Current head d3e987d differs from pull request most recent head 047fc3d. Consider uploading reports for the commit 047fc3d to get more accurate results


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


```diff
@@           Coverage Diff           @@
##           master     #540   +/-   ##
=======================================
  Coverage   99.49%   99.49%           
=======================================
  Files          41       41           
  Lines        2196     2199    +3     
=======================================
+ Hits         2185     2188    +3     
  Misses         11       11           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/reata/sqllineage/pull/540?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: ```sql
    insert into public.tgt_tbl1
    (
        id,
        id_original
    )
    select
        a || b || c || id as id,
        id as id_original        -- # noqa: E501 TODO: I need the metadata information for the table public.src_tbl1 to identify whether the column reference 'id' in this context is from the table public.src_tbl1 or from an alias reference, currently being used as an alias reference. Note: This decision may significantly deviate from the actual scenario.
    from
        public.src_tbl1
```

**I need the help of a pro.😜😜😜**
reata: See my comment on #539 or email. Maybe we should take a step back and rethink on the approach.
maoxingda: ```sql
select
    a || b || c as c,
    c as d            -- Metadata for a subquery is needed in this context to confirm whether the reference to 'c' is from the subquery or an alias reference.
from
    (
        select
            1 as a,
            2 as b,
            3 as c
    )
```


maoxingda: https://github.com/maoxingda/sqllineage/blob/ba0f63097ce06d20bb31e564c25af5e0e8ebfb99/sqllineage/core/parser/sqlfluff/extractors/cte.py#L67

https://github.com/maoxingda/sqllineage/blob/ba0f63097ce06d20bb31e564c25af5e0e8ebfb99/sqllineage/core/parser/sqlfluff/extractors/select.py#L77

https://github.com/maoxingda/sqllineage/blob/ba0f63097ce06d20bb31e564c25af5e0e8ebfb99/sqllineage/core/parser/sqlfluff/extractors/update.py#L76

Move the parsing of the subquery ahead of the for loop, so that the metadata information of the subquery is available. This way, can we determine whether the column references in the SELECT clause come from the subquery or lateral alias references?
maoxingda: After #552 merge, I will refactor the configuration of LATERAL_COLUMN_ALIAS_REFERENCE
reata: Not sure if this is already complete. Please re-request review once you're done. 
maoxingda: > Not sure if this is already complete. Please re-request review once you're done.

already done.
reata: Rebase master and make LATERAL_COLUMN_ALIAS_REFERENCE bool type. 

This PR indeed makes huge changes to current design. I need to take more time to review and refactor. The part I dislike most is that we need to do different things by telling if `SQLLineageConfig.LATERAL_COLUMN_ALIAS_REFERENCE and bool(self.metadata_provider) is True` everywhere. It will be better we can do this in a unified way as use this condition as few as possible.

The first option I was trying is to move `self.extract_subquery(subqueries, holder)` up in `select.py` regardless. But that doesn't work yet. I left a open question, not sure if that will work. If you can help me investigate, that will be helpful.
maoxingda: OK, let me try.
maoxingda: already refactor done, please re-review, thanks. ✅
maoxingda: https://www.databricks.com/blog/introducing-support-lateral-column-alias

I don’t know much about databricks but it seems that it also supports lateral column alias (**LCA**) reference.
maoxingda: https://sqlkover.com/cool-stuff-in-snowflake-part-4-aliasing-all-the-things/

It seems that snowflake also supports it, but unfortunately I don’t have the environment to verify it.
maoxingda: https://forums.oracle.com/ords/apexds/post/how-to-reference-column-alias-names-in-other-columns-of-a-q-7737

It seems that snowflake also supports it, but unfortunately I don’t find the official document.
reata: Popularity is one of the considerations. The major reason I'd like to move this feature as configurable is that it won't function without metadata. And the assumption is that by default sqllineage only does static code analysis and metadata is not present.
maoxingda: 👌
reata: Postpone merging to next week. 

Right now `to_source_columns` does part of the column to table/subquery resolution work. It would be better that we can handle this universally in  `_build_digraph` method of class SQLLineageHolder.

I'm investigating whether we can move all logic to `end_of_query_cleanup` and not modify `to_source_columns`.

Patch:
diff --git a/docs/gear_up/configuration.rst b/docs/gear_up/configuration.rst
index 76301cb..8653358 100644
--- a/docs/gear_up/configuration.rst
+++ b/docs/gear_up/configuration.rst
@@ -15,6 +15,8 @@ default schema name.
 
 Default: ``""``
 
+Since: 1.5.0
+
 DIRECTORY
 =========
 Frontend app SQL directory. By default the frontend app is showing the data directory packaged with sqllineage,
@@ -22,18 +24,56 @@ which includes tpcds queries for demo purposes. User can customize with this key
 
 Default: ``sqllineage/data``
 
+Since: 1.2.1
+
+LATERAL_COLUMN_ALIAS_REFERENCE
+==============================
+Enable lateral column alias reference. This is a syntax feature supported by some SQL dialects. See:
+
+- Amazon Redshift: `Amazon Redshift announces support for lateral column alias reference`_
+- Spark (since 3.4): `Support "lateral column alias references" to allow column aliases to be used within SELECT clauses`_
+- Databricks: `Introducing the Support of Lateral Column Alias`_
+
+.. note::
+     Lateral column alias reference is a feature that must be used together with metadata for each column to be
+     correctly resolved. Take below example:
+
+     .. code-block:: sql
+
+        SELECT clicks / impressions as probability,
+               round(100 * probability, 1) as percentage
+        FROM raw_data
+
+     If table raw_data has a column named **probability**, **probability** in the second selected column is from table
+     raw_data. Otherwise, it's referencing alias **clicks / impressions as probability**.
+
+     That means with SQLLineage, besides making LATERAL_COLUMN_ALIAS_REFERENCE=TRUE, MetaDataProvider must also be
+     provided so we can query raw_data table to see if it has a column named **probability** and then check alias reference.
+     If not provided, we will fallback to default behavior to simply assume column **probability** is from table raw_data
+     even if LATERAL_COLUMN_ALIAS_REFERENCE is set to TRUE.
+
+Default: ``False``
+
+Since: 1.5.1
+
 TSQL_NO_SEMICOLON
 =================
-Enable tsql no semicolon splitter mode.
+Enable tsql no semicolon splitter mode. Transact-SQL offers this feature that even when SQL statements are not delimited
+by semicolon, it can still be parsed and executed.
 
 .. warning::
-     Transact-SQL offers this feature that even when SQL statements are not delimited by semicolon, it can still be
-     parsed and executed. But quoting `tsql_syntax_convention`_, "although the semicolon isn't required for most
-     statements in this version (v16) of SQL Server, it will be required in a future version". So this config key is
-     kept mostly for backward-compatible purposes. We may drop the support any time without warning. Bear this in mind
-     when using this feature with sqllineage.
+     Quoting `Transact-SQL syntax conventions (Transact-SQL)`_, "although the semicolon isn't required for most
+     statements in this version (v16) of SQL Server, it will be required in a future version".
+
+     So with SQLLineage, this config key is kept mostly for backward-compatible purposes. We may drop the support any
+     time without warning. Bear this in mind when using this feature.
 
 Default: ``False``
 
+Since: 1.4.8
+
 
-.. _tsql_syntax_convention: https://learn.microsoft.com/en-us/sql/t-sql/language-elements/transact-sql-syntax-conventions-transact-sql?view=sql-server-ver16
+.. _Amazon Redshift announces support for lateral column alias reference: https://aws.amazon.com/about-aws/whats-new/2018/08/amazon-redshift-announces-support-for-lateral-column-alias-reference/
+.. _Support "lateral column alias references" to allow column aliases to be used within SELECT clauses: https://issues.apache.org/jira/browse/SPARK-27561
+.. _Introducing the Support of Lateral Column Alias: https://www.databricks.com/blog/introducing-support-lateral-column-alias
+.. _Transact-SQL syntax conventions (Transact-SQL): https://learn.microsoft.com/en-us/sql/t-sql/language-elements/transact-sql-syntax-conventions-transact-sql?view=sql-server-ver16
diff --git a/sqllineage/config.py b/sqllineage/config.py
index 4cc7dc1..6e75ac6 100644
--- a/sqllineage/config.py
+++ b/sqllineage/config.py
@@ -14,6 +14,8 @@ class _SQLLineageConfigLoader:
         "DEFAULT_SCHEMA": (str, ""),
         # to enable tsql no semicolon splitter mode
         "TSQL_NO_SEMICOLON": (bool, False),
+        # lateral column alias reference supported by some dialect (redshift, spark 3.4+, etc)
+        "LATERAL_COLUMN_ALIAS_REFERENCE": (bool, False),
     }
     BOOLEAN_TRUE_STRINGS = ("true", "on", "ok", "y", "yes", "1")
 
diff --git a/sqllineage/core/holders.py b/sqllineage/core/holders.py
index 42823fc..8a7086d 100644
--- a/sqllineage/core/holders.py
+++ b/sqllineage/core/holders.py
@@ -420,7 +420,7 @@ class SQLLineageHolder(ColumnLineageMixin):
         for unresolved_col, tgt_col in unresolved_cols:
             # check if there's only one parent candidate contains the column with same name
             src_cols = []
-            # check if source column exists in graph
+            # check if source column exists in graph (either from subquery or from table created in prev statement)
             for parent in unresolved_col.parent_candidates:
                 src_col = Column(unresolved_col.raw_name)
                 src_col.parent = parent
@@ -442,7 +442,9 @@ class SQLLineageHolder(ColumnLineageMixin):
             # It incorrect for JOIN with ON, but sql without specifying an alias in this case will be invalid
             for src_col in src_cols:
                 g.add_edge(src_col, tgt_col, type=EdgeType.LINEAGE)
-            g.remove_edge(unresolved_col, tgt_col)
+            if len(src_cols) > 0:
+                # only delete unresolved column when it's resolved
+                g.remove_edge(unresolved_col, tgt_col)
 
         # when unresolved column got resolved, it will be orphan node, and we can remove it
         for node in [n for n, deg in g.degree if deg == 0]:
diff --git a/sqllineage/core/models.py b/sqllineage/core/models.py
index 8eb012f..01e6985 100644
--- a/sqllineage/core/models.py
+++ b/sqllineage/core/models.py
@@ -158,6 +158,7 @@ class Column:
                 "source_columns", ((self.raw_name, None),)
             )
         ]
+        self.from_alias = kwargs.pop("from_alias", False)
 
     def __str__(self):
         return (
diff --git a/sqllineage/core/parser/__init__.py b/sqllineage/core/parser/__init__.py
index ac26d62..ddbd237 100644
--- a/sqllineage/core/parser/__init__.py
+++ b/sqllineage/core/parser/__init__.py
@@ -1,5 +1,6 @@
-from typing import List, Tuple, Union
+from typing import Dict, List, Tuple, Union
 
+from sqllineage.config import SQLLineageConfig
 from sqllineage.core.holders import SubQueryLineageHolder
 from sqllineage.core.models import Column, Path, SubQuery, Table
 from sqllineage.exceptions import SQLLineageException
@@ -23,19 +24,14 @@ class SourceHandlerMixin:
             if holder.write:
                 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(
-                        holder.get_alias_mapping_from_table_group(tbl_grp)
+                tgt_tbl = next(iter(holder.write))
+                lateral_column_aliases: Dict[str, List[Column]] = {}
+                for idx, tgt_col_from_query in enumerate(col_grp):
+                    tgt_col_from_query.parent = tgt_tbl
+                    tgt_col_resolved = tgt_col_from_query
+                    src_cols_resolved = []
+                    for src_col in tgt_col_from_query.to_source_columns(
+                        holder.get_alias_mapping_from_table_group(tbl_grp),
                     ):
                         if len(write_columns := holder.write_columns) == len(col_grp):
                             # example query: create view test (col3) select col1 as col2 from tab
@@ -43,21 +39,43 @@ class SourceHandlerMixin:
                             # when write_columns exist and length matches, we want tgt_col = col3 instead of col2
                             # 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
+                            tgt_col_resolved = write_columns[idx]
+                        # lateral column alias handling
+                        lca_flag = False
+                        if SQLLineageConfig.LATERAL_COLUMN_ALIAS_REFERENCE:
+                            if metadata_provider := getattr(
+                                self, "metadata_provider", None
                             ):
-                                is_lateral_alias_ref = True
-                                for lateral_alias_col in holder.get_source_columns(wc):
-                                    holder.add_column_lineage(
-                                        lateral_alias_col, tgt_col
+                                from_dataset = False
+                                for parent_candidate in src_col.parent_candidates:
+                                    if isinstance(
+                                        parent_candidate, Table
+                                    ) and src_col in metadata_provider.get_table_columns(
+                                        parent_candidate
+                                    ):
+                                        from_dataset = True
+                                    elif isinstance(
+                                        parent_candidate, SubQuery
+                                    ) and src_col in holder.get_table_columns(
+                                        parent_candidate
+                                    ):
+                                        from_dataset = True
+                                if not from_dataset and (
+                                    lca_cols_resolved := lateral_column_aliases.get(
+                                        src_col.raw_name, []
                                     )
-                                break
-                        if is_lateral_alias_ref:
-                            continue
-                        holder.add_column_lineage(src_col, tgt_col)
+                                ):
+                                    src_cols_resolved.extend(lca_cols_resolved)
+                                    lca_flag = True
+                        if not lca_flag:
+                            src_cols_resolved.append(src_col)
+
+                    for src_col_resolved in src_cols_resolved:
+                        holder.add_column_lineage(src_col_resolved, tgt_col_resolved)
+                        if (
+                            SQLLineageConfig.LATERAL_COLUMN_ALIAS_REFERENCE
+                            and tgt_col_from_query.from_alias
+                        ):
+                            lateral_column_aliases[tgt_col_from_query.raw_name] = (
+                                src_cols_resolved
+                            )
diff --git a/sqllineage/core/parser/sqlfluff/extractors/base.py b/sqllineage/core/parser/sqlfluff/extractors/base.py
index 275e3f1..640ac35 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/base.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/base.py
@@ -2,6 +2,7 @@ from functools import reduce
 from operator import add
 from typing import List, Optional, Type, Union
 
+import networkx as nx
 from sqlfluff.core.parser import BaseSegment
 
 from sqllineage.core.holders import SubQueryLineageHolder
@@ -16,6 +17,7 @@ from sqllineage.core.parser.sqlfluff.utils import (
     list_join_clause,
     list_subqueries,
 )
+from sqllineage.utils.constant import NodeTag
 from sqllineage.utils.entities import AnalyzerContext, SubQueryTuple
 from sqllineage.utils.helpers import escape_identifier_name
 
@@ -220,9 +222,12 @@ class BaseExtractor:
                 if sq.query.get_child("with_compound_statement")
                 else SelectExtractor
             )
-            holder |= extractor_cls(self.dialect, self.metadata_provider).extract(
-                sq.query, AnalyzerContext(cte=holder.cte, write={sq})
-            )
+            subquery_holder = extractor_cls(
+                self.dialect, self.metadata_provider
+            ).extract(sq.query, AnalyzerContext(cte=holder.cte, write={sq}))
+            # remove WRITE tag from subquery so that the combined holder won't have multiple WRITE dataset
+            nx.set_node_attributes(subquery_holder.graph, {sq, False}, NodeTag.WRITE)
+            holder |= subquery_holder
 
     @staticmethod
     def _init_holder(context: AnalyzerContext) -> SubQueryLineageHolder:
diff --git a/sqllineage/core/parser/sqlfluff/extractors/select.py b/sqllineage/core/parser/sqlfluff/extractors/select.py
index 5864d3d..f19a3b0 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/select.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/select.py
@@ -48,6 +48,17 @@ class SelectExtractor(BaseExtractor, SourceHandlerMixin):
                 # so that each handler don't have to worry about what's inside subquery
                 subqueries.append(sq)
 
+            if is_set_expression(segment):
+                for _, sub_segment in enumerate(
+                    segment.get_children("select_statement", "bracketed")
+                ):
+                    for seg in list_child_segments(sub_segment):
+                        for sq in self.list_subquery(seg):
+                            subqueries.append(sq)
+
+        self.extract_subquery(subqueries, holder)
+
+        for segment in segments:
             self._handle_swap_partition(segment, holder)
             self._handle_select_into(segment, holder)
             self.tables.extend(
@@ -64,17 +75,14 @@ class SelectExtractor(BaseExtractor, SourceHandlerMixin):
                             (len(self.columns), len(self.tables))
                         )
                     for seg in list_child_segments(sub_segment):
-                        for sq in self.list_subquery(seg):
-                            subqueries.append(sq)
                         self.tables.extend(
                             self._list_table_from_from_clause_or_join_clause(
                                 seg, holder
                             )
                         )
                         self._handle_column(seg)
-        self.end_of_query_cleanup(holder)
 
-        self.extract_subquery(subqueries, holder)
+        self.end_of_query_cleanup(holder)
 
         holder.expand_wildcard(self.metadata_provider)
 
diff --git a/sqllineage/core/parser/sqlfluff/models.py b/sqllineage/core/parser/sqlfluff/models.py
index 57ab368..da498d0 100644
--- a/sqllineage/core/parser/sqlfluff/models.py
+++ b/sqllineage/core/parser/sqlfluff/models.py
@@ -105,10 +105,7 @@ class SqlFluffColumn(Column):
         if column.type == "select_clause_element":
             source_columns, alias = SqlFluffColumn._get_column_and_alias(column)
             if alias:
-                return Column(
-                    alias,
-                    source_columns=source_columns,
-                )
+                return Column(alias, source_columns=source_columns, from_alias=True)
             if source_columns:
                 column_name = None
                 for sub_segment in list_child_segments(column):
@@ -145,10 +142,7 @@ class SqlFluffColumn(Column):
 
         # Wildcard, Case, Function without alias (thus not recognized as an Identifier)
         source_columns = SqlFluffColumn._extract_source_columns(column)
-        return Column(
-            column.raw,
-            source_columns=source_columns,
-        )
+        return Column(column.raw, source_columns=source_columns)
 
     @staticmethod
     def _extract_source_columns(segment: BaseSegment) -> List[ColumnQualifierTuple]:


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
index 28b7dc8..9a6f0ed 100644
--- a/tests/sql/column/test_column_select_lateral_alias_ref.py
+++ b/tests/sql/column/test_column_select_lateral_alias_ref.py
@@ -1,13 +1,28 @@
+from unittest.mock import patch
+
+import pytest
+
+from sqllineage.core.metadata_provider import MetaDataProvider
 from sqllineage.utils.entities import ColumnQualifierTuple
-from ...helpers import assert_column_lineage_equal
+from ...helpers import assert_column_lineage_equal, generate_metadata_providers
+
+providers = generate_metadata_providers(
+    {
+        "public.src_tbl1": ["id", "a", "b", "c"],
+    }
+)
 
 
-def test_column_top_level_lateral_ref():
+@pytest.mark.parametrize("provider", providers)
+@patch("os.environ", {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
+def test_column_top_level_enable_lateral_ref(
+    provider: MetaDataProvider,
+):
     sql = """
     insert into public.tgt_tbl1
     select
         name               as user_name,
-        user_name || email as id         -- lateral ref
+        user_name || email as id
     from
         public.src_tbl1
     """
@@ -27,10 +42,215 @@ def test_column_top_level_lateral_ref():
                 ColumnQualifierTuple("id", "public.tgt_tbl1"),
             ),
         ],
+        test_sqlparse=False,
+        metadata_provider=provider,
+    )
+    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,
+        metadata_provider=provider,
+    )
+    sql = """
+    insert into public.tgt_tbl1
+    (
+        id,
+        id_original
+    )
+    select
+        'a || b || c' || id as id,
+        id                  as id_original
+    from
+        public.src_tbl1
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("id_original", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
+        metadata_provider=provider,
+    )
+
+
+@pytest.mark.parametrize("provider", providers)
+@patch("os.environ", {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
+def test_column_top_level_enable_lateral_ref_with_metadata_from_table(
+    provider: MetaDataProvider,
+):
+    sql = """
+    insert into public.tgt_tbl1
+    select
+        a || b || c || id as id,
+        id                as id_original
+    from
+        public.src_tbl1
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("a", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("b", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("c", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("id_original", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
+        metadata_provider=provider,
     )
 
 
-def test_column_lateral_ref_within_subquery():
+@pytest.mark.parametrize("provider", providers)
+@patch("os.environ", {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
+def test_column_top_level_enable_lateral_ref_with_metadata_from_subquery(
+    provider: MetaDataProvider,
+):
+    sql = """
+    insert into public.tgt_tbl1
+    select
+        a || b || c || id as id,
+        id                as id_original
+    from
+        (
+            select
+               *
+            from
+                public.src_tbl1
+        ) as sq
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("a", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("b", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("c", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("id_original", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
+        metadata_provider=provider,
+    )
+
+
+@pytest.mark.parametrize("provider", providers)
+@patch("os.environ", {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
+def test_column_top_level_enable_lateral_ref_with_metadata_from_nested_subquery(
+    provider: MetaDataProvider,
+):
+    sql = """
+    insert into public.tgt_tbl1
+    select
+        a || b || c || id as id,
+        id                as id_original
+    from
+        (
+            select
+                *
+            from
+                (
+                    select
+                       *
+                    from
+                        public.src_tbl1
+                ) as inner_sq
+        ) as outer_sq
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("a", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("b", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("c", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("id_original", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
+        metadata_provider=provider,
+    )
+
+
+@pytest.mark.parametrize("provider", providers)
+@patch("os.environ", {"SQLLINEAGE_LATERAL_COLUMN_ALIAS_REFERENCE": "1"})
+def test_column_enable_lateral_ref_within_subquery(
+    provider: MetaDataProvider,
+):
     sql = """
     insert into public.tgt_tbl1
     select
@@ -38,8 +258,8 @@ def test_column_lateral_ref_within_subquery():
     from
         (
             select
-                id || name as alias1,
-                alias1 || email as name
+                id     || a as alias1,
+                alias1 || b as name
             from
                 public.src_tbl1
         ) as sq
@@ -52,14 +272,16 @@ def test_column_lateral_ref_within_subquery():
                 ColumnQualifierTuple("name", "public.tgt_tbl1"),
             ),
             (
-                ColumnQualifierTuple("name", "public.src_tbl1"),
+                ColumnQualifierTuple("a", "public.src_tbl1"),
                 ColumnQualifierTuple("name", "public.tgt_tbl1"),
             ),
             (
-                ColumnQualifierTuple("email", "public.src_tbl1"),
+                ColumnQualifierTuple("b", "public.src_tbl1"),
                 ColumnQualifierTuple("name", "public.tgt_tbl1"),
             ),
         ],
+        test_sqlparse=False,
+        metadata_provider=provider,
     )
 
     sql = """
@@ -95,4 +317,190 @@ def test_column_lateral_ref_within_subquery():
                 ColumnQualifierTuple("name", "public.tgt_tbl1"),
             ),
         ],
+        test_sqlparse=False,
+        metadata_provider=provider,
+    )
+
+
+def test_column_top_level_disable_lateral_ref():
+    sql = """
+    insert into public.tgt_tbl1
+    select
+        name               as user_name,
+        user_name || email as id
+    from
+        public.src_tbl1
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("name", "public.src_tbl1"),
+                ColumnQualifierTuple("user_name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("user_name", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("email", "public.src_tbl1"),
+                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,
+    )
+    sql = """
+    insert into public.tgt_tbl1
+    (
+        id,
+        id_original
+    )
+    select
+        'a || b || c' || id as id,
+        id                  as id_original
+    from
+        public.src_tbl1
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("id_original", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
+
+
+def test_column_top_level_disable_lateral_ref_with_metadata_from_table():
+    sql = """
+    insert into public.tgt_tbl1
+    select
+        a || b || c || id as id,
+        id                as id_original
+    from
+        public.src_tbl1
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("a", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("b", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("c", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("id", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("id", "public.src_tbl1"),
+                ColumnQualifierTuple("id_original", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
+
+
+def test_column_disable_lateral_ref_within_subquery():
+    sql = """
+    insert into public.tgt_tbl1
+    select
+        sq.name
+    from
+        (
+            select
+                id     || a as alias1,
+                alias1 || b as name
+            from
+                public.src_tbl1
+        ) as sq
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("alias1", "public.src_tbl1"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("b", "public.src_tbl1"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
+
+    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("alias1", None),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("email", "public.src_tbl2"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
     )