fix: case when subqueries in the select_clause are not correctly reco…

…gnized (#560)
reata · Jan 20, 2024 · 8a02894 · 8a02894
1 parent 06a474b
commit 8a02894
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git 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"):

diff --git a/tests/sql/table/test_select.py b/tests/sql/table/test_select.py
@@ -140,6 +140,18 @@ def test_select_subquery_in_case():
 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():