Repository: reata/sqllineage

Issue:
Not fully processed top-level subquery in DML
**SQL**
```sql
INSERT INTO tab1 (col1, col2) ((SELECT * FROM tab2 UNION SELECT * FROM tab3))
INSERT INTO tab1 (col1, col2) (((SELECT * FROM tab2 UNION SELECT * FROM tab3)))

CREATE VIEW my_view (random1,random2) AS ((SELECT col1,col2 FROM tbl))
CREATE VIEW my_view (random1,random2) AS (((SELECT col1,col2 FROM tbl)))
```

**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_table_lineage()
```

**Actual behavior**
```
Statements(#): 1
Source Tables:
    
Target Tables:
    <default>.tab1
```

**Expected behavior**
```
Statements(#): 1
Source Tables:
    <default>.tab2
    <default>.tab3
Target Tables:
    <default>.tab1
```

**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/565?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 [(`4c6c65e`)](https://app.codecov.io/gh/reata/sqllineage/commit/4c6c65e7c94b87b50f26ad7f35ac26ed26a54cee?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) 99.50% compared to head [(`9060655`)](https://app.codecov.io/gh/reata/sqllineage/pull/565?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) 99.50%.


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


```diff
@@            Coverage Diff             @@
##           master     #565      +/-   ##
==========================================
- Coverage   99.50%   99.50%   -0.01%     
==========================================
  Files          41       41              
  Lines        2208     2203       -5     
==========================================
- Hits         2197     2192       -5     
  Misses         11       11              
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/reata/sqllineage/pull/565?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).


Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index a0e4a96..aa4edca 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,15 +1,15 @@
 repos:
   - repo: https://github.com/psf/black
-    rev: 23.9.1
+    rev: 24.1.0
     hooks:
       - id: black
         language_version: python3.8
   - repo: https://github.com/pycqa/flake8.git
-    rev: 6.1.0
+    rev: 7.0.0
     hooks:
       - id: flake8
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.5.1
+    rev: v1.8.0
     hooks:
     - id: mypy
       additional_dependencies:
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/cli.py b/sqllineage/cli.py
index 87104bb..9159349 100644
--- a/sqllineage/cli.py
+++ b/sqllineage/cli.py
@@ -106,9 +106,11 @@ def main(args=None) -> None:
         runner = LineageRunner(
             sql,
             dialect=args.dialect,
-            metadata_provider=SQLAlchemyMetaDataProvider(args.sqlalchemy_url)
-            if args.sqlalchemy_url
-            else DummyMetaDataProvider(),
+            metadata_provider=(
+                SQLAlchemyMetaDataProvider(args.sqlalchemy_url)
+                if args.sqlalchemy_url
+                else DummyMetaDataProvider()
+            ),
             verbose=args.verbose,
             draw_options={
                 "host": args.host,
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/create_insert.py b/sqllineage/core/parser/sqlfluff/extractors/create_insert.py
index 526d726..7610bac 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/create_insert.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/create_insert.py
@@ -6,7 +6,6 @@ from sqllineage.core.parser.sqlfluff.extractors.base import BaseExtractor
 from sqllineage.core.parser.sqlfluff.extractors.select import SelectExtractor
 from sqllineage.core.parser.sqlfluff.models import SqlFluffColumn, SqlFluffTable
 from sqllineage.core.parser.sqlfluff.utils import (
-    is_set_expression,
     list_child_segments,
 )
 from sqllineage.utils.entities import AnalyzerContext
@@ -56,14 +55,16 @@ class CreateInsertExtractor(BaseExtractor):
                                         select_statement, holder
                                     )
             elif segment.type == "bracketed" and (
-                self.list_subquery(segment) or is_set_expression(segment)
+                subquery_segments := list(
+                    segment.recursive_crawl(
+                        "select_statement", "set_expression", recurse_into=False
+                    )
+                )
             ):
                 # note regular subquery within SELECT statement is handled by SelectExtractor, this is only to handle
                 # top-level subquery in DML like: 1) create table foo as (subquery); 2) insert into foo (subquery)
                 # subquery here isn't added as read source, and it inherits DML-level write_columns if parsed
-                if subquery_segment := segment.get_child(
-                    "select_statement", "set_expression"
-                ):
+                for subquery_segment in subquery_segments:
                     holder |= self.delegate_to_select(subquery_segment, holder)
 
             elif segment.type == "bracketed":
diff --git a/sqllineage/core/parser/sqlfluff/extractors/cte.py b/sqllineage/core/parser/sqlfluff/extractors/cte.py
index 656d893..0cead85 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/cte.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/cte.py
@@ -46,23 +46,16 @@ class CteExtractor(BaseExtractor):
                     UpdateExtractor, segment, AnalyzerContext(cte=holder.cte)
                 )
             elif segment.type == "common_table_expression":
-                identifier = None
-                segment_has_alias = any(
-                    s for s in segment.get_children("keyword") if s.raw_upper == "AS"
-                )
+                alias = None
                 sub_segments = list_child_segments(segment)
                 for sub_segment in sub_segments:
                     if sub_segment.type == "identifier":
-                        identifier = sub_segment.raw
-                        if not segment_has_alias:
-                            holder.add_cte(SqlFluffSubQuery.of(sub_segment, identifier))
-                    if sub_segment.type == "bracketed":
+                        alias = sub_segment.raw
+                    elif sub_segment.type == "bracketed":
                         for sq in self.list_subquery(sub_segment):
-                            if identifier:
-                                sq.alias = identifier
+                            sq.alias = alias
                             subqueries.append(sq)
-                        if segment_has_alias:
-                            holder.add_cte(SqlFluffSubQuery.of(sub_segment, identifier))
+                        holder.add_cte(SqlFluffSubQuery.of(sub_segment, alias))
 
         self.extract_subquery(subqueries, holder)
 
diff --git a/sqllineage/core/parser/sqlfluff/extractors/merge.py b/sqllineage/core/parser/sqlfluff/extractors/merge.py
index 1758e30..01d54be 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/merge.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/merge.py
@@ -109,9 +109,11 @@ class MergeExtractor(BaseExtractor):
                     next_segment = segments[i + 1]
                     direct_source = SqlFluffSubQuery.of(
                         extract_innermost_bracketed(segment),
-                        extract_identifier(next_segment)
-                        if next_segment.type == "alias_expression"
-                        else None,
+                        (
+                            extract_identifier(next_segment)
+                            if next_segment.type == "alias_expression"
+                            else None
+                        ),
                     )
                     holder.add_read(direct_source)
 
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]:
diff --git a/sqllineage/core/parser/sqlfluff/utils.py b/sqllineage/core/parser/sqlfluff/utils.py
index 2f80106..09a7559 100644
--- a/sqllineage/core/parser/sqlfluff/utils.py
+++ b/sqllineage/core/parser/sqlfluff/utils.py
@@ -8,6 +8,7 @@ naming convention:
     extract_x: other pattern
 first parameter of each function must be sqlfluff BaseSegment
 """
+
 from typing import List, Optional, Tuple
 
 from sqlfluff.core.parser import BaseSegment
@@ -156,9 +157,11 @@ def list_subqueries(segment: BaseSegment) -> List[SubQueryTuple]:
         if is_subquery(target):
             subquery = [
                 SubQueryTuple(
-                    extract_innermost_bracketed(target)
-                    if not is_set_expression(target)
-                    else target,
+                    (
+                        extract_innermost_bracketed(target)
+                        if not is_set_expression(target)
+                        else target
+                    ),
                     extract_identifier(as_segment) if as_segment else None,
                 )
             ]
diff --git a/sqllineage/io.py b/sqllineage/io.py
index 555f093..dd82b16 100644
--- a/sqllineage/io.py
+++ b/sqllineage/io.py
@@ -12,9 +12,11 @@ def to_cytoscape(graph: DiGraph, compound=False) -> List[Dict[str, Dict[str, Any
         parents_dict = {
             node.parent: {
                 "name": str(node.parent) if node.parent is not None else "<unknown>",
-                "type": type(node.parent).__name__
-                if node.parent is not None
-                else "Table or SubQuery",
+                "type": (
+                    type(node.parent).__name__
+                    if node.parent is not None
+                    else "Table or SubQuery"
+                ),
             }
             for node in graph.nodes
         }


Test patch:
diff --git a/tests/sql/column/test_column_select_column_specified_in_dml.py b/tests/sql/column/test_column_select_column_specified_in_dml.py
index a68db1e..0cc8939 100644
--- a/tests/sql/column/test_column_select_column_specified_in_dml.py
+++ b/tests/sql/column/test_column_select_column_specified_in_dml.py
@@ -24,6 +24,36 @@ def test_view_with_subquery_custom_columns():
         ],
         test_sqlparse=False,
     )
+    sql = "CREATE VIEW my_view (random1,random2) AS ((SELECT col1,col2 FROM tbl))"
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("col1", "tbl"),
+                ColumnQualifierTuple("random1", "my_view"),
+            ),
+            (
+                ColumnQualifierTuple("col2", "tbl"),
+                ColumnQualifierTuple("random2", "my_view"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
+    sql = "CREATE VIEW my_view (random1,random2) AS (((SELECT col1,col2 FROM tbl)))"
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("col1", "tbl"),
+                ColumnQualifierTuple("random1", "my_view"),
+            ),
+            (
+                ColumnQualifierTuple("col2", "tbl"),
+                ColumnQualifierTuple("random2", "my_view"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
     sql = "CREATE VIEW my_view (random1,random2) AS SELECT col1,col2 FROM tbl"
     assert_column_lineage_equal(
         sql,
@@ -82,6 +112,18 @@ def test_insert_into_with_columns_and_select_union():
         {"tab1"},
         test_sqlparse=False,
     )
+    assert_table_lineage_equal(
+        "INSERT INTO tab1 (col1, col2) ((SELECT * FROM tab2 UNION SELECT * FROM tab3))",
+        {"tab2", "tab3"},
+        {"tab1"},
+        test_sqlparse=False,
+    )
+    assert_table_lineage_equal(
+        "INSERT INTO tab1 (col1, col2) (((SELECT * FROM tab2 UNION SELECT * FROM tab3)))",
+        {"tab2", "tab3"},
+        {"tab1"},
+        test_sqlparse=False,
+    )
 
 
 def test_insert_with_custom_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
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,
     )