Repository: reata/sqllineage

Issue:
Handling CTE at the start of query in DML
Given that query, I expect that `model` should be in `cte` section. 
Is it some kind of reserved word?

```sql
INSERT INTO db.table1 ("date")
WITH raw AS (
  SELECT
    date
  FROM db.table2
), model AS (
  SELECT
    raw.date
  FROM raw
  LEFT JOIN db.table3 USING (date)
)
SELECT * FROM model

table read: [Table: <default>.model, Table: db.table2, Table: db.table3]
table write: [Table: db.table1]
table cte: []
```

When there is only one CTE it works correctly
```sql
INSERT INTO db.table1 ("date")
WITH model AS (
  SELECT
    db.date
  FROM db.table2
)
SELECT * FROM model

table read: [Table: db.table2]
table write: [Table: db.table1]
table cte: [SubQuery: model]
```

If I wrap `raw` in backticks it also works

```sql
INSERT INTO db.table1 ("date")
WITH `raw` AS (
  SELECT
    date
  FROM db.table2
), model AS (
  SELECT
    raw.date
  FROM raw
  LEFT JOIN db.table3
    USING (date)
)
SELECT * FROM model

table read: [Table: db.table2, Table: db.table3]
table write: [Table: db.table1]
table cte: [SubQuery: `raw`, SubQuery: model]
```

versions that I use:
```
sqllineage               1.3.7
sqlparse                 0.4.1
```

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/369?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#369](https://app.codecov.io/gh/reata/sqllineage/pull/369?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (c2f5b62) into [master](https://app.codecov.io/gh/reata/sqllineage/commit/c6a470135a4a2633b0f217ab147efebd727e9f75?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (c6a4701) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master     #369   +/-   ##
=======================================
  Coverage   99.39%   99.39%           
=======================================
  Files          42       42           
  Lines        1972     1975    +3     
=======================================
+ Hits         1960     1963    +3     
  Misses         12       12           
```


| [Impacted Files](https://app.codecov.io/gh/reata/sqllineage/pull/369?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [...parser/sqlfluff/extractors/dml\_insert\_extractor.py](https://app.codecov.io/gh/reata/sqllineage/pull/369?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL3BhcnNlci9zcWxmbHVmZi9leHRyYWN0b3JzL2RtbF9pbnNlcnRfZXh0cmFjdG9yLnB5) | `100.00% <100.00%> (ø)` | |



Patch:
diff --git a/sqllineage/core/parser/sqlfluff/extractors/dml_insert_extractor.py b/sqllineage/core/parser/sqlfluff/extractors/dml_insert_extractor.py
index 3d3a475..dd2202a 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/dml_insert_extractor.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/dml_insert_extractor.py
@@ -59,6 +59,14 @@ class DmlInsertExtractor(LineageHolderExtractor):
             for current_handler in handlers:
                 current_handler.handle(segment, holder)
 
+            if segment.type == "with_compound_statement":
+                from .cte_extractor import DmlCteExtractor
+
+                holder |= DmlCteExtractor(self.dialect).extract(
+                    segment,
+                    AnalyzerContext(prev_cte=holder.cte, prev_write=holder.write),
+                )
+
             if segment.type == "select_statement":
                 holder |= DmlSelectExtractor(self.dialect).extract(
                     segment,


Test patch:
diff --git a/tests/test_cte.py b/tests/test_cte.py
index 1fd979f..1934505 100644
--- a/tests/test_cte.py
+++ b/tests/test_cte.py
@@ -54,3 +54,11 @@ def test_with_insert():
         {"tab2"},
         {"tab3"},
     )
+
+
+def test_with_insert_in_query():
+    assert_table_lineage_equal(
+        "INSERT INTO tab3 WITH tab1 AS (SELECT * FROM tab2) SELECT * FROM tab1",
+        {"tab2"},
+        {"tab3"},
+    )