Repository: reata/sqllineage

Issue:
No target tables in UPDATE statement using CTE
SQL statement (pure UPDATE) (see below). After parsing in parser_result target_tables are empty!
```
WITH data_table AS ( 
	SELECT 
		id, 
		user_id,
		recommended_vacancies,
		push_dtm,
		push_check_datetime
	FROM dap_pushs_multyrec_sent
	WHERE 
		status = 'CHECKED'
		AND push_dtm > NOW() - INTERVAL '2 months'
		AND push_success = 0
		AND (push_success_notimeframe_total IS NULL OR push_success_notimeframe_total = 0)
	),
-----
unnest_vacs AS (
	SELECT user_id, UNNEST (recommended_vacancies) AS vacancy_id FROM data_table
	),
----
activities_count_table AS ( 
	SELECT * FROM unnest_vacs	
	LEFT JOIN LATERAL(
		SELECT 
			count (*) AS num_activities_vacancy,
			min(ewa_created_on) AS reaction_dtm
		FROM dap_ewa 
		WHERE ewa_executor_id = unnest_vacs.user_id
		AND ewa_vacancy_id = unnest_vacs.vacancy_id
	) activities ON TRUE),
-----
data_with_actcount_unnested AS (
	SELECT 
		id, data_table.user_id, recommended_vacancies, vacancy_id, num_activities_vacancy, reaction_dtm,
		data_table.push_dtm, data_table.push_check_datetime
	FROM data_table
	LEFT JOIN activities_count_table ON activities_count_table.vacancy_id = ANY (data_table.recommended_vacancies) AND data_table.user_id = activities_count_table.user_id
	),
----
result_table AS (
	SELECT 
		id, user_id, 	
		recommended_vacancies,
		array_agg((num_activities_vacancy>0)::integer) AS push_success_notimeframe,
		sum((num_activities_vacancy>0)::integer) push_success_notimeframe_total,
		array_agg (reaction_dtm AT time ZONE 'UTC-3' - push_dtm) AS reaction_time_notimeframe		
	FROM data_with_actcount_unnested
	GROUP BY id, user_id, recommended_vacancies)
----
UPDATE dap_pushs_multyrec_sent
SET
	push_success_notimeframe = result_table.push_success_notimeframe,
	push_success_notimeframe_total = result_table.push_success_notimeframe_total,
	reaction_time_notimeframe = result_table.reaction_time_notimeframe
FROM result_table
WHERE result_table.id = dap_pushs_multyrec_sent.id;
```

```
parser_result = LineageRunner(sql_query, dialect = 'postgres')
```

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/495?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 [(`b66274e`)](https://app.codecov.io/gh/reata/sqllineage/commit/b66274eb702d81afb0b7f2a6478e349e66979adc?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) 99.41% compared to head [(`5fb704d`)](https://app.codecov.io/gh/reata/sqllineage/pull/495?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) 99.41%.


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


```diff
@@           Coverage Diff           @@
##           master     #495   +/-   ##
=======================================
  Coverage   99.41%   99.41%           
=======================================
  Files          40       40           
  Lines        2034     2045   +11     
=======================================
+ Hits         2022     2033   +11     
  Misses         12       12           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/reata/sqllineage/pull/495?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/cte.py b/sqllineage/core/parser/sqlfluff/extractors/cte.py
index 4356dcd..656d893 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/cte.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/cte.py
@@ -6,6 +6,7 @@ from sqllineage.core.parser.sqlfluff.extractors.create_insert import (
     CreateInsertExtractor,
 )
 from sqllineage.core.parser.sqlfluff.extractors.select import SelectExtractor
+from sqllineage.core.parser.sqlfluff.extractors.update import UpdateExtractor
 from sqllineage.core.parser.sqlfluff.models import SqlFluffSubQuery
 from sqllineage.core.parser.sqlfluff.utils import list_child_segments
 from sqllineage.utils.entities import AnalyzerContext
@@ -40,6 +41,10 @@ class CteExtractor(BaseExtractor):
                 holder |= self.delegate_to(
                     CreateInsertExtractor, segment, AnalyzerContext(cte=holder.cte)
                 )
+            elif segment.type == "update_statement":
+                holder |= self.delegate_to(
+                    UpdateExtractor, segment, AnalyzerContext(cte=holder.cte)
+                )
             elif segment.type == "common_table_expression":
                 identifier = None
                 segment_has_alias = any(
diff --git a/sqllineage/core/parser/sqlfluff/extractors/update.py b/sqllineage/core/parser/sqlfluff/extractors/update.py
index 64a8eae..a5c78ad 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/update.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/update.py
@@ -1,10 +1,11 @@
-from typing import Optional
+from typing import Optional, Union
 
 from sqlfluff.core.parser import BaseSegment
 
-from sqllineage.core.holders import StatementLineageHolder
-from sqllineage.core.models import Table
+from sqllineage.core.holders import SubQueryLineageHolder
+from sqllineage.core.models import SubQuery, Table
 from sqllineage.core.parser.sqlfluff.extractors.base import BaseExtractor
+from sqllineage.core.parser.sqlfluff.models import SqlFluffSubQuery
 from sqllineage.core.parser.sqlfluff.utils import (
     find_from_expression_element,
     find_table_identifier,
@@ -23,19 +24,19 @@ class UpdateExtractor(BaseExtractor):
 
     def extract(
         self, statement: BaseSegment, context: AnalyzerContext
-    ) -> StatementLineageHolder:
-        holder = StatementLineageHolder()
+    ) -> SubQueryLineageHolder:
+        holder = self._init_holder(context)
         tgt_flag = False
         for segment in list_child_segments(statement):
             if segment.type == "from_expression":
-                # UPDATE with JOIN
+                # UPDATE with JOIN, mysql only syntax
                 if table := self.find_table_from_from_expression_or_join_clause(
-                    segment
+                    segment, holder
                 ):
                     holder.add_write(table)
                 for join_clause in list_join_clause(statement):
                     if table := self.find_table_from_from_expression_or_join_clause(
-                        join_clause
+                        join_clause, holder
                     ):
                         holder.add_read(table)
 
@@ -48,13 +49,25 @@ class UpdateExtractor(BaseExtractor):
                     holder.add_write(table)
                 tgt_flag = False
 
+            if segment.type == "from_clause":
+                # UPDATE FROM, ansi syntax
+                if from_expression := segment.get_child("from_expression"):
+                    if table := self.find_table_from_from_expression_or_join_clause(
+                        from_expression, holder
+                    ):
+                        holder.add_read(table)
+
         return holder
 
     def find_table_from_from_expression_or_join_clause(
-        self, segment
-    ) -> Optional[Table]:
-        table = None
+        self, segment, holder
+    ) -> Optional[Union[Table, SubQuery]]:
+        table: Optional[Union[Table, SubQuery]] = None
         if from_expression_element := find_from_expression_element(segment):
             if table_identifier := find_table_identifier(from_expression_element):
-                table = self.find_table(table_identifier)
+                cte_dict = {c.alias: c for c in holder.cte}
+                if cte := cte_dict.get(table_identifier.raw):
+                    table = SqlFluffSubQuery.of(cte.query, table_identifier.raw)
+                else:
+                    table = self.find_table(table_identifier)
         return table


Test patch:
diff --git a/tests/sql/table/test_cte.py b/tests/sql/table/test_cte.py
index c9ab271..b9ddede 100644
--- a/tests/sql/table/test_cte.py
+++ b/tests/sql/table/test_cte.py
@@ -48,7 +48,7 @@ def test_with_select_join_table_with_same_name():
     )
 
 
-def test_with_insert():
+def test_cte_insert():
     assert_table_lineage_equal(
         "WITH tab1 AS (SELECT * FROM tab2) INSERT INTO tab3 SELECT * FROM tab1",
         {"tab2"},
@@ -56,7 +56,7 @@ def test_with_insert():
     )
 
 
-def test_with_insert_in_query():
+def test_cte_insert_cte_in_query():
     assert_table_lineage_equal(
         "INSERT INTO tab3 WITH tab1 AS (SELECT * FROM tab2) SELECT * FROM tab1",
         {"tab2"},
@@ -64,6 +64,17 @@ def test_with_insert_in_query():
     )
 
 
+def test_cte_update_from():
+    assert_table_lineage_equal(
+        """WITH cte1 AS (SELECT col1, col2 FROM tab1)
+UPDATE tab2
+SET tab2.col2 = cte1.col2 FROM cte1
+WHERE tab2.col1 = cte1.col1;""",
+        {"tab1"},
+        {"tab2"},
+    )
+
+
 def test_cte_and_union():
     sql = """WITH cte_1 AS (select col1 from tab1)
 SELECT col2 from tab2
diff --git a/tests/sql/table/test_update.py b/tests/sql/table/test_update.py
index ef5887a..7736e7e 100644
--- a/tests/sql/table/test_update.py
+++ b/tests/sql/table/test_update.py
@@ -5,3 +5,13 @@ def test_update():
     assert_table_lineage_equal(
         "UPDATE tab1 SET col1='val1' WHERE col2='val2'", None, {"tab1"}
     )
+
+
+def test_update_from():
+    assert_table_lineage_equal(
+        """UPDATE tab2
+SET tab2.col2 = tab1.col2 FROM tab1
+WHERE tab2.col1 = tab1.col1""",
+        {"tab1"},
+        {"tab2"},
+    )