fix: misidentify binary operator * as wildcard (#510)

reata · Dec 24, 2023 · 6c646a5 · 6c646a5
1 parent c0cf6aa
commit 6c646a5
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 8 deletions.
diff --git a/sqllineage/core/parser/sqlfluff/models.py b/sqllineage/core/parser/sqlfluff/models.py
@@ -209,20 +209,20 @@ def _get_column_from_subquery(
     def _get_column_from_parenthesis(
         sub_segment: BaseSegment,
     ) -> List[ColumnQualifierTuple]:
-        """
-        :param sub_segment: segment to be processed
-        :return: list of columns and alias from the segment
-        """
-        col, _ = SqlFluffColumn._get_column_and_alias(sub_segment)
-        if col:
-            return col
+        # windows function has an extra layer, get rid of it so that it can be handled as regular functions
+        if window_specification := sub_segment.get_child("window_specification"):
+            sub_segment = window_specification
         col, _ = SqlFluffColumn._get_column_and_alias(sub_segment, False)
         return col if col else []
 
     @staticmethod
     def _get_column_and_alias(
         segment: BaseSegment, check_bracketed: bool = True
     ) -> Tuple[List[ColumnQualifierTuple], Optional[str]]:
+        """
+        check_bracketed is True for top-level column definition, like (col1 + col2) as col3
+        set to False for bracket in function call, like coalesce(col1, col2) as col3
+        """
         alias = None
         columns = []
         sub_segments = list_child_segments(segment, check_bracketed)

diff --git a/sqllineage/core/parser/sqlfluff/utils.py b/sqllineage/core/parser/sqlfluff/utils.py
@@ -48,7 +48,7 @@ def is_subquery(segment: BaseSegment) -> bool:
 
 def is_wildcard(segment: BaseSegment) -> bool:
     return segment.type == "wildcard_expression" or (
-        segment.type == "symbol" and segment.raw == "*"
+        segment.type == "symbol" and segment.raw == "*" and segment.get_type() == "star"
     )
 
 

diff --git a/tests/sql/column/test_column_select_expression.py b/tests/sql/column/test_column_select_expression.py
@@ -37,6 +37,21 @@ def test_select_column_using_expression():
     )
 
 
+def test_select_column_multiply_expression():
+    sql = """INSERT INTO tab1
+SELECT col1 * 1 AS col1
+FROM tab2"""
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("col1", "tab2"),
+                ColumnQualifierTuple("col1", "tab1"),
+            ),
+        ],
+    )
+
+
 def test_select_column_using_expression_in_parenthesis():
     sql = """INSERT INTO tab1
 SELECT (col1 + col2)