Repository: reata/sqllineage (python)

Issue:
Problem with SELECT CAST(CASE WHEN ...END AS DECIMAL(M,N)) AS col_name
Sorry for this problem, but if statement is like:
SELECT CAST(CASE WHEN ..., 
                               WHEN ...,
                               END AS DECIMAL(M, N)) AS col_name
FROM ...
the result is:
<default>table_name.col_name -> decimal
I don't know how to deal with this statement, or could I cut 'CAST’ function for analysis?

Pull request description:
Fix Two Bugs When Cast Used Together With Case When
Closes #215 

Patch:
diff --git a/sqllineage/__init__.py b/sqllineage/__init__.py
index 694c13e..6fd30ce 100644
--- a/sqllineage/__init__.py
+++ b/sqllineage/__init__.py
@@ -2,12 +2,25 @@ import logging.config
 import os
 
 
+def _patch_adding_window_function_token() -> None:
+    from sqlparse.engine import grouping
+    from sqllineage.utils.sqlparse import group_function_with_window
+
+    grouping.group_functions = group_function_with_window
+
+
+def _patch_adding_builtin_type() -> None:
+    from sqlparse import tokens
+    from sqlparse.keywords import KEYWORDS
+
+    KEYWORDS["STRING"] = tokens.Name.Builtin
+    KEYWORDS["DATETIME"] = tokens.Name.Builtin
+
+
 def _monkey_patch() -> None:
     try:
-        from sqlparse.engine import grouping
-        from sqllineage.utils.sqlparse import group_function_with_window
-
-        grouping.group_functions = group_function_with_window
+        _patch_adding_window_function_token()
+        _patch_adding_builtin_type()
     except ImportError:
         # when imported by setup.py for constant variables, dependency is not ready yet
         pass
diff --git a/sqllineage/core/models.py b/sqllineage/core/models.py
index a8be030..93a0af0 100644
--- a/sqllineage/core/models.py
+++ b/sqllineage/core/models.py
@@ -287,18 +287,22 @@ class Column:
                 for cqt in Column._extract_source_columns(tk)
             ]
         elif isinstance(token, Identifier):
-            if token.get_real_name():
-                # col1 AS col2
-                source_columns = [
-                    ColumnQualifierTuple(token.get_real_name(), token.get_parent_name())
-                ]
-            else:
-                # col1=1 AS int
+            real_name = token.get_real_name()
+            if real_name is None or (
+                real_name == "decimal" and isinstance(token.tokens[-1], Function)
+            ):
+                # real name is None: col1=1 AS int
+                # real_name is decimal: case when col1 > 0 then col2 else col3 end as decimal(18, 0)
                 source_columns = [
                     cqt
                     for tk in token.get_sublists()
                     for cqt in Column._extract_source_columns(tk)
                 ]
+            else:
+                # col1 AS col2
+                source_columns = [
+                    ColumnQualifierTuple(token.get_real_name(), token.get_parent_name())
+                ]
         else:
             # Handle literals other than *
             if (


Test patch:
diff --git a/tests/test_columns.py b/tests/test_columns.py
index aa5ff3f..99ec782 100644
--- a/tests/test_columns.py
+++ b/tests/test_columns.py
@@ -535,7 +535,9 @@ FROM tab2"""
     )
 
 
-@pytest.mark.parametrize("dtype", ["string", "timestamp", "date", "datetime"])
+@pytest.mark.parametrize(
+    "dtype", ["string", "timestamp", "date", "datetime", "decimal(18, 0)"]
+)
 def test_cast_to_data_type(dtype):
     sql = f"""INSERT OVERWRITE TABLE tab1
 SELECT cast(col1 as {dtype}) AS col1
@@ -546,6 +548,32 @@ FROM tab2"""
     )
 
 
+@pytest.mark.parametrize(
+    "dtype", ["string", "timestamp", "date", "datetime", "decimal(18, 0)"]
+)
+def test_cast_to_data_type_with_case_when(dtype):
+    sql = f"""INSERT OVERWRITE TABLE tab1
+SELECT cast(case when col1 > 0 then col2 else col3 end as {dtype}) AS col1
+FROM tab2"""
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("col1", "tab2"),
+                ColumnQualifierTuple("col1", "tab1"),
+            ),
+            (
+                ColumnQualifierTuple("col2", "tab2"),
+                ColumnQualifierTuple("col1", "tab1"),
+            ),
+            (
+                ColumnQualifierTuple("col3", "tab2"),
+                ColumnQualifierTuple("col1", "tab1"),
+            ),
+        ],
+    )
+
+
 def test_cast_using_constant():
     sql = """INSERT OVERWRITE TABLE tab1
 SELECT cast('2012-12-21' as date) AS col2"""