Repository: reata/sqllineage
Pull request #561 (2024-01-19 06:54:46)

Allow multiple sources for column lineage
fixes: [558](https://github.com/reata/sqllineage/issues/558)

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/561?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 [(`06a474b`)](https://app.codecov.io/gh/reata/sqllineage/commit/06a474b5e7d6a26994db18f97cf9058246cd24bd?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) 99.50% compared to head [(`345af90`)](https://app.codecov.io/gh/reata/sqllineage/pull/561?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     #561      +/-   ##
==========================================
- Coverage   99.50%   99.50%   -0.01%     
==========================================
  Files          41       41              
  Lines        2211     2208       -3     
==========================================
- Hits         2200     2197       -3     
  Misses         11       11              
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/reata/sqllineage/pull/561?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/sqllineage/core/holders.py b/sqllineage/core/holders.py
index eb5085d..42823fc 100644
--- a/sqllineage/core/holders.py
+++ b/sqllineage/core/holders.py
@@ -6,7 +6,6 @@ from networkx import DiGraph
 
 from sqllineage.core.metadata_provider import MetaDataProvider
 from sqllineage.core.models import Column, Path, Schema, SubQuery, Table
-from sqllineage.exceptions import InvalidSyntaxException
 from sqllineage.utils.constant import EdgeTag, EdgeType, NodeTag
 
 DATASET_CLASSES = (Path, Table)
@@ -439,14 +438,12 @@ class SQLLineageHolder(ColumnLineageMixin):
                             if unresolved_col.raw_name == src_col.raw_name:
                                 src_cols.append(src_col)
 
-            if len(src_cols) > 1:
-                raise InvalidSyntaxException(
-                    f"{unresolved_col.raw_name} is not allowed from more than one table or subquery"
-                )
+            # Multiple sources is a correct case for JOIN with USING
+            # 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) == 1:
-                g.add_edge(src_cols[0], tgt_col, type=EdgeType.LINEAGE)
-                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]:
             if isinstance(node, Column) and len(node.parent_candidates) > 1:
diff --git a/sqllineage/core/parser/sqlfluff/utils.py b/sqllineage/core/parser/sqlfluff/utils.py
index 6092673..2f80106 100644
--- a/sqllineage/core/parser/sqlfluff/utils.py
+++ b/sqllineage/core/parser/sqlfluff/utils.py
@@ -131,7 +131,7 @@ def list_expression_from_when_clause(
 def list_subqueries(segment: BaseSegment) -> List[SubQueryTuple]:
     subquery = []
     if segment.type == "select_clause":
-        if select_clause_element := segment.get_child("select_clause_element"):
+        for select_clause_element in segment.get_children("select_clause_element"):
             if expression := select_clause_element.get_child("expression"):
                 if case_expression := expression.get_child("case_expression"):
                     for when_clause in case_expression.get_children("when_clause"):


Test patch:
diff --git a/tests/sql/column/test_column_select_from_join.py b/tests/sql/column/test_column_select_from_join.py
index f82a44d..cd164bf 100644
--- a/tests/sql/column/test_column_select_from_join.py
+++ b/tests/sql/column/test_column_select_from_join.py
@@ -63,14 +63,22 @@ FROM tab2 a
 
 
 def test_select_column_without_table_qualifier_from_table_join():
-    sql = """INSERT INTO tab1
-SELECT col1
-FROM tab2 a
-         INNER JOIN tab3 b
-                    ON a.id = b.id"""
+    sql = """INSERT INTO tab3
+SELECT f1
+FROM ( SELECT f1 FROM tab1)
+LEFT JOIN ( SELECT f1 FROM tab2) USING (f1)"""
     assert_column_lineage_equal(
         sql,
-        [(ColumnQualifierTuple("col1", None), ColumnQualifierTuple("col1", "tab1"))],
+        [
+            (
+                ColumnQualifierTuple("f1", "tab1"),
+                ColumnQualifierTuple("f1", "tab3"),
+            ),
+            (
+                ColumnQualifierTuple("f1", "tab2"),
+                ColumnQualifierTuple("f1", "tab3"),
+            ),
+        ],
     )
 
 
diff --git a/tests/sql/column/test_metadata_unqualified_column.py b/tests/sql/column/test_metadata_unqualified_column.py
index 8fd9d58..a1ec7ae 100644
--- a/tests/sql/column/test_metadata_unqualified_column.py
+++ b/tests/sql/column/test_metadata_unqualified_column.py
@@ -1,8 +1,6 @@
 import pytest
 
 from sqllineage.core.metadata_provider import MetaDataProvider
-from sqllineage.exceptions import InvalidSyntaxException
-from sqllineage.runner import LineageRunner
 from sqllineage.utils.entities import ColumnQualifierTuple
 from ...helpers import assert_column_lineage_equal, generate_metadata_providers
 
@@ -274,21 +272,3 @@ left join (select pk, p, q as z from db3.table3) t3 on t2.id = t3.pk
         ],
         metadata_provider=provider,
     )
-
-
-@pytest.mark.parametrize("provider", providers)
-def test_sqlparse_exception(provider: MetaDataProvider):
-    sql = """insert into db.tbl
-select id
-from db1.table1 t1
-join db2.table2 t2 on t1.id = t2.id
-"""
-
-    with pytest.raises(
-        InvalidSyntaxException,
-        match="id is not allowed from more than one table or subquery",
-    ):
-        lr = LineageRunner(sql, metadata_provider=provider)
-        col_lineage = lr.get_column_lineage()
-        for e in col_lineage:
-            print(e)
diff --git a/tests/sql/table/test_select.py b/tests/sql/table/test_select.py
index d47ee62..ae70a27 100644
--- a/tests/sql/table/test_select.py
+++ b/tests/sql/table/test_select.py
@@ -140,6 +140,18 @@ CASE WHEN (SELECT count(*) FROM tab1 WHERE col1 = 'tab2') = 1 THEN (SELECT count
 CASE WHEN 1 = (SELECT count(*) FROM tab1 WHERE col1 = 'tab2') THEN (SELECT count(*) FROM tab2) ELSE 0 END AS cnt""",
         {"tab1", "tab2"},
     )
+    assert_table_lineage_equal(
+        """
+        select
+            1,
+            case
+                when 1 = ( select count(*) from tab1 where col1 = 'tab2' )
+                    then ( select count(*) from tab2 )
+                else 0
+            end as cnt
+        """,
+        {"tab1", "tab2"},
+    )
 
 
 def test_select_subquery_without_alias():