fix: column lineage for nested cast

sqllineage/core/models.py

@@ -288,11 +288,14 @@ def _extract_source_columns(token: Token) -> List[ColumnQualifierTuple]:
             ]
         elif isinstance(token, Identifier):
             real_name = token.get_real_name()
-            if real_name is None or (
-                real_name == "decimal" and isinstance(token.tokens[-1], Function)
-            ):
+            if (
                 # real name is None: col1=1 AS int
+                real_name is None
                 # real_name is decimal: case when col1 > 0 then col2 else col3 end as decimal(18, 0)
+                or (real_name == "decimal" and isinstance(token.tokens[-1], Function))
+                # real_name is cast: cast(col1 AS string) AS string
+                or (real_name == "cast" and isinstance(token.tokens[0], Function))
+            ):
                 source_columns = [
                     cqt
                     for tk in token.get_sublists()

tests/test_columns.py

@@ -577,6 +577,26 @@ def test_cast_to_data_type(dtype):
     )
 
 
+@pytest.mark.parametrize(
+    "dtype", ["string", "timestamp", "date", "datetime", "decimal(18, 0)"]
+)
+def test_nested_cast_to_data_type(dtype):
+    sql = f"""INSERT OVERWRITE TABLE tab1
+SELECT cast(cast(col1 AS {dtype}) AS {dtype}) AS col1
+FROM tab2"""
+    assert_column_lineage_equal(
+        sql,
+        [(ColumnQualifierTuple("col1", "tab2"), ColumnQualifierTuple("col1", "tab1"))],
+    )
+    sql = f"""INSERT OVERWRITE TABLE tab1
+SELECT cast(cast(cast(cast(cast(col1 AS {dtype}) AS {dtype}) AS {dtype}) AS {dtype}) AS {dtype}) AS col1
+FROM tab2"""
+    assert_column_lineage_equal(
+        sql,
+        [(ColumnQualifierTuple("col1", "tab2"), ColumnQualifierTuple("col1", "tab1"))],
+    )
+
+
 @pytest.mark.parametrize(
     "dtype", ["string", "timestamp", "date", "datetime", "decimal(18, 0)"]
 )