Repository: reata/sqllineage

Issue:
update statement column lineage
```sql
update 
    public.tgt_tbl1
set 
    email = s.email
from
    public.src_tbl1 s
where
    s.id = tgt_tbl1.id
;
```


Discussion:
Column lineage for UPDATE statement is not supported as of v1.4.8.

It should be quite similar to what we did with MERGE statement.



Patch:
diff --git a/sqllineage/core/holders.py b/sqllineage/core/holders.py
index fe663d00..688833fe 100644
--- a/sqllineage/core/holders.py
+++ b/sqllineage/core/holders.py
@@ -1,5 +1,5 @@
 import itertools
-from typing import List, Optional, Set, Tuple, Union
+from typing import Dict, List, Optional, Set, Tuple, Union
 
 import networkx as nx
 from networkx import DiGraph
@@ -172,6 +172,27 @@ def expand_wildcard(self, metadata_provider: MetaDataProvider) -> None:
                                     src_wildcard,
                                 )
 
+    def get_alias_mapping_from_table_group(
+        self, table_group: List[Union[Path, Table, SubQuery]]
+    ) -> Dict[str, Union[Path, Table, SubQuery]]:
+        """
+        A table can be referred to as alias, table name, or database_name.table_name, create the mapping here.
+        For SubQuery, it's only alias then.
+        """
+        return {
+            **{
+                tgt: src
+                for src, tgt, attr in self.graph.edges(data=True)
+                if attr.get("type") == EdgeType.HAS_ALIAS and src in table_group
+            },
+            **{
+                table.raw_name: table
+                for table in table_group
+                if isinstance(table, Table)
+            },
+            **{str(table): table for table in table_group if isinstance(table, Table)},
+        }
+
     def _get_target_table(self) -> Optional[Union[SubQuery, Table]]:
         table = None
         if write_only := self.write.difference(self.read):
diff --git a/sqllineage/core/parser/__init__.py b/sqllineage/core/parser/__init__.py
index f9c7ea1b..ac26d62c 100644
--- a/sqllineage/core/parser/__init__.py
+++ b/sqllineage/core/parser/__init__.py
@@ -1,9 +1,8 @@
-from typing import Dict, List, Tuple, Union
+from typing import List, Tuple, Union
 
 from sqllineage.core.holders import SubQueryLineageHolder
 from sqllineage.core.models import Column, Path, SubQuery, Table
 from sqllineage.exceptions import SQLLineageException
-from sqllineage.utils.constant import EdgeType
 
 
 class SourceHandlerMixin:
@@ -36,7 +35,7 @@ def end_of_query_cleanup(self, holder: SubQueryLineageHolder) -> None:
                             lateral_aliases.add(tgt_col.raw_name)
                             break
                     for src_col in tgt_col.to_source_columns(
-                        self.get_alias_mapping_from_table_group(tbl_grp, holder)
+                        holder.get_alias_mapping_from_table_group(tbl_grp)
                     ):
                         if len(write_columns := holder.write_columns) == len(col_grp):
                             # example query: create view test (col3) select col1 as col2 from tab
@@ -62,27 +61,3 @@ def end_of_query_cleanup(self, holder: SubQueryLineageHolder) -> None:
                         if is_lateral_alias_ref:
                             continue
                         holder.add_column_lineage(src_col, tgt_col)
-
-    @classmethod
-    def get_alias_mapping_from_table_group(
-        cls,
-        table_group: List[Union[Path, Table, SubQuery]],
-        holder: SubQueryLineageHolder,
-    ) -> Dict[str, Union[Path, Table, SubQuery]]:
-        """
-        A table can be referred to as alias, table name, or database_name.table_name, create the mapping here.
-        For SubQuery, it's only alias then.
-        """
-        return {
-            **{
-                tgt: src
-                for src, tgt, attr in holder.graph.edges(data=True)
-                if attr.get("type") == EdgeType.HAS_ALIAS and src in table_group
-            },
-            **{
-                table.raw_name: table
-                for table in table_group
-                if isinstance(table, Table)
-            },
-            **{str(table): table for table in table_group if isinstance(table, Table)},
-        }
diff --git a/sqllineage/core/parser/sqlfluff/extractors/base.py b/sqllineage/core/parser/sqlfluff/extractors/base.py
index ede1cd81..1f1ca364 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/base.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/base.py
@@ -1,18 +1,23 @@
 from functools import reduce
 from operator import add
-from typing import List, Optional, Type
+from typing import List, Optional, Type, Union
 
 from sqlfluff.core.parser import BaseSegment
 
 from sqllineage.core.holders import SubQueryLineageHolder
 from sqllineage.core.metadata_provider import MetaDataProvider
-from sqllineage.core.models import SubQuery, Table
+from sqllineage.core.models import Path, SubQuery, Table
 from sqllineage.core.parser.sqlfluff.models import SqlFluffSubQuery, SqlFluffTable
 from sqllineage.core.parser.sqlfluff.utils import (
+    find_from_expression_element,
+    find_table_identifier,
     is_subquery,
+    list_child_segments,
+    list_join_clause,
     list_subqueries,
 )
 from sqllineage.utils.entities import AnalyzerContext, SubQueryTuple
+from sqllineage.utils.helpers import escape_identifier_name
 
 
 class BaseExtractor:
@@ -80,6 +85,101 @@ def list_subquery(cls, segment: BaseSegment) -> List[SubQuery]:
             result = [SqlFluffSubQuery.of(segment, None)]
         return result
 
+    def _list_table_from_from_clause_or_join_clause(
+        self, segment: BaseSegment, holder: SubQueryLineageHolder
+    ) -> List[Union[Table, SubQuery, Path]]:
+        """
+        Extract table from from_clause or join_clause, join_clause is a child node of from_clause.
+        """
+        tables = []
+        if segment.type in ["from_clause", "join_clause", "update_statement"]:
+            if len(from_expressions := segment.get_children("from_expression")) > 1:
+                # SQL89 style of join
+                for from_expression in from_expressions:
+                    if from_expression_element := find_from_expression_element(
+                        from_expression
+                    ):
+                        tables += self._add_dataset_from_expression_element(
+                            from_expression_element, holder
+                        )
+            else:
+                if from_expression_element := find_from_expression_element(segment):
+                    tables += self._add_dataset_from_expression_element(
+                        from_expression_element, holder
+                    )
+                for join_clause in list_join_clause(segment):
+                    tables += self._list_table_from_from_clause_or_join_clause(
+                        join_clause, holder
+                    )
+
+        return tables
+
+    @staticmethod
+    def _add_dataset_from_expression_element(
+        segment: BaseSegment, holder: SubQueryLineageHolder
+    ) -> List[Union[Table, SubQuery, Path]]:
+        """
+        Append tables and subqueries identified in the 'from_expression_element' type segment to the table and
+        holder extra subqueries sets
+        """
+        tables: List[Union[Table, SubQuery, Path]] = []
+        all_segments = [
+            seg for seg in list_child_segments(segment) if seg.type != "keyword"
+        ]
+        if table_expression := segment.get_child("table_expression"):
+            if table_expression.get_child("function"):
+                # for UNNEST or generator function, no dataset involved
+                return tables
+        first_segment = all_segments[0]
+        if first_segment.type == "bracketed":
+            if table_expression := first_segment.get_child("table_expression"):
+                if table_expression.get_child("values_clause"):
+                    # (VALUES ...) AS alias, no dataset involved
+                    return tables
+        subqueries = list_subqueries(segment)
+        if subqueries:
+            for sq in subqueries:
+                bracketed, alias = sq
+                read_sq = SqlFluffSubQuery.of(bracketed, alias)
+                tables.append(read_sq)
+        else:
+            table_identifier = find_table_identifier(segment)
+            if table_identifier:
+                subquery_flag = False
+                alias = None
+                if len(all_segments) > 1 and all_segments[1].type == "alias_expression":
+                    all_segments = list_child_segments(all_segments[1])
+                    alias = str(
+                        all_segments[1].raw
+                        if len(all_segments) > 1
+                        else all_segments[0].raw
+                    )
+                if "." not in table_identifier.raw:
+                    cte_dict = {s.alias: s for s in holder.cte}
+                    cte = cte_dict.get(table_identifier.raw)
+                    if cte is not None:
+                        # could reference CTE with or without alias
+                        tables.append(
+                            SqlFluffSubQuery.of(
+                                cte.query,
+                                alias or table_identifier.raw,
+                            )
+                        )
+                        subquery_flag = True
+                if subquery_flag is False:
+                    if table_identifier.type == "file_reference":
+                        tables.append(
+                            Path(
+                                escape_identifier_name(
+                                    table_identifier.segments[-1].raw
+                                )
+                            )
+                        )
+                    else:
+                        tables.append(SqlFluffTable.of(table_identifier, alias=alias))
+
+        return tables
+
     @classmethod
     def _parse_subquery(cls, subqueries: List[SubQueryTuple]) -> List[SubQuery]:
         """
diff --git a/sqllineage/core/parser/sqlfluff/extractors/select.py b/sqllineage/core/parser/sqlfluff/extractors/select.py
index 005a5cda..5864d3d3 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/select.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/select.py
@@ -2,21 +2,16 @@
 
 from sqllineage.core.holders import SubQueryLineageHolder
 from sqllineage.core.metadata_provider import MetaDataProvider
-from sqllineage.core.models import Path
 from sqllineage.core.parser import SourceHandlerMixin
 from sqllineage.core.parser.sqlfluff.extractors.base import BaseExtractor
 from sqllineage.core.parser.sqlfluff.models import (
     SqlFluffColumn,
-    SqlFluffSubQuery,
     SqlFluffTable,
 )
 from sqllineage.core.parser.sqlfluff.utils import (
-    find_from_expression_element,
     find_table_identifier,
     is_set_expression,
     list_child_segments,
-    list_join_clause,
-    list_subqueries,
 )
 from sqllineage.utils.entities import AnalyzerContext
 from sqllineage.utils.helpers import escape_identifier_name
@@ -55,7 +50,9 @@ def extract(
 
             self._handle_swap_partition(segment, holder)
             self._handle_select_into(segment, holder)
-            self._handle_table(segment, holder)
+            self.tables.extend(
+                self._list_table_from_from_clause_or_join_clause(segment, holder)
+            )
             self._handle_column(segment)
 
             if is_set_expression(segment):
@@ -69,7 +66,11 @@ def extract(
                     for seg in list_child_segments(sub_segment):
                         for sq in self.list_subquery(seg):
                             subqueries.append(sq)
-                        self._handle_table(seg, holder)
+                        self.tables.extend(
+                            self._list_table_from_from_clause_or_join_clause(
+                                seg, holder
+                            )
+                        )
                         self._handle_column(seg)
         self.end_of_query_cleanup(holder)
 
@@ -113,31 +114,6 @@ def _handle_select_into(self, segment: BaseSegment, holder: SubQueryLineageHolde
                 if table := self.find_table(identifier):
                     holder.add_write(table)
 
-    def _handle_table(
-        self, segment: BaseSegment, holder: SubQueryLineageHolder
-    ) -> None:
-        """
-        handle from_clause or join_clause, join_clause is a child node of from_clause.
-        """
-        if segment.type in ["from_clause", "join_clause"]:
-            from_expressions = segment.get_children("from_expression")
-            if len(from_expressions) > 1:
-                # SQL89 style of join
-                for from_expression in from_expressions:
-                    if from_expression_element := find_from_expression_element(
-                        from_expression
-                    ):
-                        self._add_dataset_from_expression_element(
-                            from_expression_element, holder
-                        )
-            else:
-                if from_expression_element := find_from_expression_element(segment):
-                    self._add_dataset_from_expression_element(
-                        from_expression_element, holder
-                    )
-                for join_clause in list_join_clause(segment):
-                    self._handle_table(join_clause, holder)
-
     def _handle_column(self, segment: BaseSegment) -> None:
         """
         Column handler method
@@ -145,67 +121,3 @@ def _handle_column(self, segment: BaseSegment) -> None:
         if segment.type == "select_clause":
             for sub_segment in segment.get_children("select_clause_element"):
                 self.columns.append(SqlFluffColumn.of(sub_segment))
-
-    def _add_dataset_from_expression_element(
-        self, segment: BaseSegment, holder: SubQueryLineageHolder
-    ) -> None:
-        """
-        Append tables and subqueries identified in the 'from_expression_element' type segment to the table and
-        holder extra subqueries sets
-        """
-        all_segments = [
-            seg for seg in list_child_segments(segment) if seg.type != "keyword"
-        ]
-        if table_expression := segment.get_child("table_expression"):
-            if table_expression.get_child("function"):
-                # for UNNEST or generator function, no dataset involved
-                return
-        first_segment = all_segments[0]
-        if first_segment.type == "bracketed":
-            if table_expression := first_segment.get_child("table_expression"):
-                if table_expression.get_child("values_clause"):
-                    # (VALUES ...) AS alias, no dataset involved
-                    return
-        subqueries = list_subqueries(segment)
-        if subqueries:
-            for sq in subqueries:
-                bracketed, alias = sq
-                read_sq = SqlFluffSubQuery.of(bracketed, alias)
-                self.tables.append(read_sq)
-        else:
-            table_identifier = find_table_identifier(segment)
-            if table_identifier:
-                subquery_flag = False
-                alias = None
-                if len(all_segments) > 1 and all_segments[1].type == "alias_expression":
-                    all_segments = list_child_segments(all_segments[1])
-                    alias = str(
-                        all_segments[1].raw
-                        if len(all_segments) > 1
-                        else all_segments[0].raw
-                    )
-                if "." not in table_identifier.raw:
-                    cte_dict = {s.alias: s for s in holder.cte}
-                    cte = cte_dict.get(table_identifier.raw)
-                    if cte is not None:
-                        # could reference CTE with or without alias
-                        self.tables.append(
-                            SqlFluffSubQuery.of(
-                                cte.query,
-                                alias or table_identifier.raw,
-                            )
-                        )
-                        subquery_flag = True
-                if subquery_flag is False:
-                    if table_identifier.type == "file_reference":
-                        self.tables.append(
-                            Path(
-                                escape_identifier_name(
-                                    table_identifier.segments[-1].raw
-                                )
-                            )
-                        )
-                    else:
-                        self.tables.append(
-                            SqlFluffTable.of(table_identifier, alias=alias)
-                        )
diff --git a/sqllineage/core/parser/sqlfluff/extractors/update.py b/sqllineage/core/parser/sqlfluff/extractors/update.py
index a5c78ad8..32fd3465 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/update.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/update.py
@@ -1,16 +1,11 @@
-from typing import Optional, Union
-
 from sqlfluff.core.parser import BaseSegment
 
 from sqllineage.core.holders import SubQueryLineageHolder
-from sqllineage.core.models import SubQuery, Table
+from sqllineage.core.models import Column
 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,
+    extract_column_qualifier,
     list_child_segments,
-    list_join_clause,
 )
 from sqllineage.utils.entities import AnalyzerContext
 
@@ -27,47 +22,57 @@ def extract(
     ) -> SubQueryLineageHolder:
         holder = self._init_holder(context)
         tgt_flag = False
+        columns = []
+        subqueries = []
         for segment in list_child_segments(statement):
             if segment.type == "from_expression":
                 # UPDATE with JOIN, mysql only syntax
-                if table := self.find_table_from_from_expression_or_join_clause(
-                    segment, holder
+                # from_expression is wrapped directly under update_statement and there's no from_clause in this case
+                # we need to pass update_statement to the function
+                if from_join_tables := self._list_table_from_from_clause_or_join_clause(
+                    statement, 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, holder
-                    ):
-                        holder.add_read(table)
+                    holder.add_write(from_join_tables[0])
+                    for join_table in from_join_tables[1:]:
+                        holder.add_read(join_table)
 
             if segment.type == "keyword" and segment.raw_upper == "UPDATE":
                 tgt_flag = True
                 continue
 
             if tgt_flag:
-                if table := self.find_table(segment):
-                    holder.add_write(table)
+                if write_table := self.find_table(segment):
+                    holder.add_write(write_table)
                 tgt_flag = False
 
+            if segment.type == "set_clause_list":
+                for set_clause in segment.get_children("set_clause"):
+                    column_references = set_clause.get_children("column_reference")
+                    if len(column_references) == 2:
+                        tgt_cqt = extract_column_qualifier(column_references[0])
+                        src_cqt = extract_column_qualifier(column_references[1])
+                        if tgt_cqt is not None and src_cqt is not None:
+                            columns.append(
+                                Column(tgt_cqt.column, source_columns=[src_cqt])
+                            )
+
             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)
+                # there can be multiple from items, each may be a table or a subquery
+                for sq in self.list_subquery(segment):
+                    subqueries.append(sq)
+                for read_table in self._list_table_from_from_clause_or_join_clause(
+                    segment, holder
+                ):
+                    holder.add_read(read_table)
 
-        return holder
+        for tgt_col in columns:
+            tgt_col.parent = list(holder.write)[0]
+            for src_col in tgt_col.to_source_columns(
+                holder.get_alias_mapping_from_table_group(list(holder.read))
+            ):
+                holder.add_column_lineage(src_col, tgt_col)
 
-    def find_table_from_from_expression_or_join_clause(
-        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):
-                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
+        self.extract_subquery(subqueries, holder)
+
+        return holder
diff --git a/sqllineage/core/parser/sqlfluff/models.py b/sqllineage/core/parser/sqlfluff/models.py
index 921600ea..6565f510 100644
--- a/sqllineage/core/parser/sqlfluff/models.py
+++ b/sqllineage/core/parser/sqlfluff/models.py
@@ -98,9 +98,9 @@ class SqlFluffColumn(Column):
     @staticmethod
     def of(column: BaseSegment, **kwargs) -> Column:
         """
-        Build a 'SqlFluffSubQuery' object
+        Build a 'Column' object
         :param column: column segment
-        :return:
+        :return: 'Column' object
         """
         if column.type == "select_clause_element":
             source_columns, alias = SqlFluffColumn._get_column_and_alias(column)
diff --git a/sqllineage/core/parser/sqlfluff/utils.py b/sqllineage/core/parser/sqlfluff/utils.py
index c9b732cd..60926737 100644
--- a/sqllineage/core/parser/sqlfluff/utils.py
+++ b/sqllineage/core/parser/sqlfluff/utils.py
@@ -59,7 +59,7 @@ def find_from_expression_element(segment: BaseSegment) -> Optional[BaseSegment]:
         from_expression/join_clause as parent
     """
     from_expression_element = None
-    if segment.type == "from_clause":
+    if segment.type in ["from_clause", "update_statement"]:
         if from_expression := segment.get_child("from_expression"):
             non_bracket = from_expression
             while bracketed := non_bracket.get_child("bracketed"):


Test patch:
diff --git a/tests/sql/column/test_column_update.py b/tests/sql/column/test_column_update.py
new file mode 100644
index 00000000..ab78317f
--- /dev/null
+++ b/tests/sql/column/test_column_update.py
@@ -0,0 +1,160 @@
+from sqllineage.utils.entities import ColumnQualifierTuple
+from ...helpers import assert_column_lineage_equal
+
+
+def test_column_update_with_single_table():
+    sql = """
+    update
+        public.tgt_tbl1 as t
+    set
+        name  = s.name,
+        email = s.address
+    from
+        public.src_tbl1 as s
+    where
+        s.id = t.id
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("name", "public.src_tbl1"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("address", "public.src_tbl1"),
+                ColumnQualifierTuple("email", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
+
+
+def test_column_update_with_cte():
+    sql = """
+    with
+        s as (
+            select name, address from public.src_tbl1
+        )
+    update
+        public.tgt_tbl1 as t
+    set
+        name  = s.name,
+        email = s.address
+    from
+        s
+    where
+        s.id = t.id
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("name", "public.src_tbl1"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("address", "public.src_tbl1"),
+                ColumnQualifierTuple("email", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
+
+
+def test_column_update_with_multi_tables():
+    sql = """
+    update
+        public.tgt_tbl1 as t
+    set
+        name  = s1.name,
+        email = s2.email
+    from
+        public.src_tbl1 as s1
+        join
+            public.src_tbl2 as s2
+        on
+            s1.id = s2.id
+    where
+        s1.id = t.id
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("name", "public.src_tbl1"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("email", "public.src_tbl2"),
+                ColumnQualifierTuple("email", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
+    sql = """
+    update
+        public.tgt_tbl1 as t
+    set
+        name = s1.name,
+        email = s2.email
+    from
+        public.src_tbl1 as s1,
+        public.src_tbl2 as s2
+    where
+        s1.id = s2.id and
+        s1.id = t.id
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("name", "public.src_tbl1"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("email", "public.src_tbl2"),
+                ColumnQualifierTuple("email", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
+
+
+def test_column_update_with_subquery():
+    sql = """
+    update
+        public.tgt_tbl1 as t
+    set
+        name  = s.name,
+        email = s.email
+    from
+        (
+            select
+                s1.id,
+                s1.name,
+                s2.email
+            from
+                public.src_tbl1 as s1
+                join
+                    public.src_tbl2 as s2
+                on
+                    s1.id = s2.id
+        ) as s
+    where
+        s.id = t.id
+    """
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("name", "public.src_tbl1"),
+                ColumnQualifierTuple("name", "public.tgt_tbl1"),
+            ),
+            (
+                ColumnQualifierTuple("email", "public.src_tbl2"),
+                ColumnQualifierTuple("email", "public.tgt_tbl1"),
+            ),
+        ],
+        test_sqlparse=False,
+    )