Repository: reata/sqllineage

Issue:
Not Using Column Name Specified in Query For CTE within Query
**Describe the bug**
* When the column names in the target table are different from those in the source table, the column lineage is incorrect.

**SQL**
Paste the SQL text here. For example:
```sql
insert into
    public.tgt_tbl1
(
    id
)
with
    cte1 as (
        select name from public.src_tbl1
    )
select
    name
from
    cte1
;
```

**To Reproduce**
*Note here we refer to SQL provided in prior step as stored in a file named `test.sql`*

```python
with open(sql_file, 'r') as f:
    sql = f.read()

lr = LineageRunner(sql, dialect='redshift')

lr.print_column_lineage()
```
```
public.tgt_tbl1.name <- cte1.name <- public.src_tbl1.name
```

**Expected behavior**
```
public.tgt_tbl1.id <- cte1.name <- public.src_tbl1.name
```

**Python version (available via `python --version`)**
- 3.11.5

**SQLLineage version (available via `sqllineage --version`):**
 - 1.4.8


Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/492?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 [(`2d91264`)](https://app.codecov.io/gh/reata/sqllineage/commit/2d912648c14f7010b34026d5a0b1ae53dfe49d9c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) 99.40% compared to head [(`9f75139`)](https://app.codecov.io/gh/reata/sqllineage/pull/492?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) 99.40%.


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master     #492   +/-   ##
=======================================
  Coverage   99.40%   99.40%           
=======================================
  Files          40       40           
  Lines        2031     2031           
=======================================
  Hits         2019     2019           
  Misses         12       12           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/reata/sqllineage/pull/492?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/parser/sqlfluff/extractors/create_insert.py b/sqllineage/core/parser/sqlfluff/extractors/create_insert.py
index 563851a..526d726 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/create_insert.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/create_insert.py
@@ -122,7 +122,11 @@ class CreateInsertExtractor(BaseExtractor):
         from .cte import CteExtractor
 
         return self.delegate_to(
-            CteExtractor, segment, AnalyzerContext(cte=holder.cte, write=holder.write)
+            CteExtractor,
+            segment,
+            AnalyzerContext(
+                cte=holder.cte, write=holder.write, write_columns=holder.write_columns
+            ),
         )
 
     def delegate_to_select(
diff --git a/sqllineage/core/parser/sqlfluff/extractors/cte.py b/sqllineage/core/parser/sqlfluff/extractors/cte.py
index f05732d..0427a3b 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/cte.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/cte.py
@@ -30,7 +30,11 @@ class CteExtractor(BaseExtractor):
                 holder |= self.delegate_to(
                     SelectExtractor,
                     segment,
-                    AnalyzerContext(cte=holder.cte, write=holder.write),
+                    AnalyzerContext(
+                        cte=holder.cte,
+                        write=holder.write,
+                        write_columns=holder.write_columns,
+                    ),
                 )
             elif segment.type == "insert_statement":
                 holder |= self.delegate_to(


Test patch:
diff --git a/tests/sql/sqlfluff_only/test_dml_specify_column.py b/tests/sql/sqlfluff_only/test_dml_specify_column.py
index 5d57e66..6ae4b67 100644
--- a/tests/sql/sqlfluff_only/test_dml_specify_column.py
+++ b/tests/sql/sqlfluff_only/test_dml_specify_column.py
@@ -117,3 +117,19 @@ def test_insert_with_custom_columns():
         ],
         test_sqlparse=False,
     )
+
+
+def test_insert_with_custom_columns_and_cte_within_query():
+    sql = """INSERT INTO tab2 (col1)
+WITH cte1 AS (SELECT col2 FROM tab1)
+SELECT col2 FROM cte1"""
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("col2", "tab1"),
+                ColumnQualifierTuple("col1", "tab2"),
+            )
+        ],
+        test_sqlparse=False,
+    )