Repository: reata/sqllineage

Issue:
Add an Argument to Exclude SubQuery Column Node in Column Lineage Path
**Describe the bug**
now get_column_lineage's result include SubQuery

**To Reproduce**
For example:
```python
from sqllineage.runner import LineageRunner
v_sql = "insert into ta select b from (select b from tb union all select c from tc ) sub"
parse = LineageRunner(sql=v_sql)
for col_tuple in parse.get_column_lineage():
    print(col_tuple)
```
```
(Column: <default>.tb.b, Column: sub.b, Column: <default>.ta.b)
(Column: <default>.tc.c, Column: sub.b, Column: <default>.ta.b)
```

**Expected behavior**
```
(Column: <default>.tb.b, Column: <default>.ta.b)
(Column: <default>.tc.c, Column: <default>.ta.b)
```
**Python version (available via `python --version`)**
 - 3.10.13

**SQLLineage version (available via `sqllineage --version`):**
 - 1.4.9

**Additional context**
I will submit a PR later


Discussion:

Patch:
diff --git a/sqllineage/core/holders.py b/sqllineage/core/holders.py
index fe663d00..049b12f1 100644
--- a/sqllineage/core/holders.py
+++ b/sqllineage/core/holders.py
@@ -13,7 +13,16 @@
 
 
 class ColumnLineageMixin:
-    def get_column_lineage(self, exclude_subquery=True) -> Set[Tuple[Column, ...]]:
+    def get_column_lineage(
+        self, exclude_path_ending_in_subquery=True, exclude_subquery_columns=False
+    ) -> Set[Tuple[Column, ...]]:
+        """
+        :param exclude_path_ending_in_subquery:  exclude_subquery rename to exclude_path_ending_in_subquery
+               exclude column from SubQuery in the ending path
+        :param exclude_subquery_columns: exclude column from SubQuery in the path.
+
+        return a list of column tuple :class:`sqllineage.models.Column`
+        """
         self.graph: DiGraph  # For mypy attribute checking
         # filter all the column node in the graph
         column_nodes = [n for n in self.graph.nodes if isinstance(n, Column)]
@@ -25,7 +34,7 @@ def get_column_lineage(self, exclude_subquery=True) -> Set[Tuple[Column, ...]]:
             for node, deg in column_graph.out_degree
             if isinstance(node, Column) and deg == 0
         }
-        if exclude_subquery:
+        if exclude_path_ending_in_subquery:
             target_columns = {
                 node for node in target_columns if isinstance(node.parent, Table)
             }
@@ -33,7 +42,14 @@ def get_column_lineage(self, exclude_subquery=True) -> Set[Tuple[Column, ...]]:
         for source, target in itertools.product(source_columns, target_columns):
             simple_paths = list(nx.all_simple_paths(self.graph, source, target))
             for path in simple_paths:
-                columns.add(tuple(path))
+                if exclude_subquery_columns:
+                    path = [
+                        node for node in path if not isinstance(node.parent, SubQuery)
+                    ]
+                    if len(path) > 1:
+                        columns.add(tuple(path))
+                else:
+                    columns.add(tuple(path))
         return columns
 
 
diff --git a/sqllineage/core/parser/sqlfluff/models.py b/sqllineage/core/parser/sqlfluff/models.py
index 921600ea..529b812b 100644
--- a/sqllineage/core/parser/sqlfluff/models.py
+++ b/sqllineage/core/parser/sqlfluff/models.py
@@ -195,7 +195,7 @@ def _get_column_from_subquery(
             lineage[0]
             for lineage in LineageRunner(
                 sub_segment.raw, dialect=SQLPARSE_DIALECT
-            ).get_column_lineage(exclude_subquery=False)
+            ).get_column_lineage(exclude_path_ending_in_subquery=False)
         ]
         source_columns = [
             ColumnQualifierTuple(
diff --git a/sqllineage/core/parser/sqlparse/models.py b/sqllineage/core/parser/sqlparse/models.py
index 6a499062..82a7c4ac 100644
--- a/sqllineage/core/parser/sqlparse/models.py
+++ b/sqllineage/core/parser/sqlparse/models.py
@@ -110,7 +110,7 @@ def _extract_source_columns(token: Token) -> List[ColumnQualifierTuple]:
                     lineage[0]
                     for lineage in LineageRunner(
                         token.value, dialect=SQLPARSE_DIALECT
-                    ).get_column_lineage(exclude_subquery=False)
+                    ).get_column_lineage(exclude_path_ending_in_subquery=False)
                 ]
                 source_columns = [
                     ColumnQualifierTuple(
diff --git a/sqllineage/runner.py b/sqllineage/runner.py
index 7db37f4c..335a0eec 100644
--- a/sqllineage/runner.py
+++ b/sqllineage/runner.py
@@ -151,13 +151,17 @@ def intermediate_tables(self) -> List[Table]:
         return sorted(self._sql_holder.intermediate_tables, key=lambda x: str(x))
 
     @lazy_method
-    def get_column_lineage(self, exclude_subquery=True) -> List[Tuple[Column, Column]]:
+    def get_column_lineage(
+        self, exclude_path_ending_in_subquery=True, exclude_subquery_columns=False
+    ) -> List[Tuple[Column, Column]]:
         """
         a list of column tuple :class:`sqllineage.models.Column`
         """
         # sort by target column, and then source column
         return sorted(
-            self._sql_holder.get_column_lineage(exclude_subquery),
+            self._sql_holder.get_column_lineage(
+                exclude_path_ending_in_subquery, exclude_subquery_columns
+            ),
             key=lambda x: (str(x[-1]), str(x[0])),
         )
 


Test patch:
diff --git a/tests/core/test_runner.py b/tests/core/test_runner.py
index ead0ef62..96593928 100644
--- a/tests/core/test_runner.py
+++ b/tests/core/test_runner.py
@@ -1,3 +1,4 @@
+from sqllineage.core.models import SubQuery
 from sqllineage.runner import LineageRunner
 from sqllineage.utils.constant import LineageLevel
 
@@ -22,3 +23,11 @@ def test_statements_trim_comment():
 def test_silent_mode():
     sql = "begin; select * from dual;"
     LineageRunner(sql, dialect="greenplum", silent_mode=True)._eval()
+
+
+def test_get_column_lineage_exclude_subquery_inpath():
+    v_sql = "insert into ta select b from (select b from tb union all select c from tc ) sub"
+    parse = LineageRunner(sql=v_sql)
+    for col_tuple in parse.get_column_lineage(exclude_subquery_columns=True):
+        for col in col_tuple:
+            assert not isinstance(col.parent, SubQuery)