Repository: reata/sqllineage

Issue:
Specify Column Names in Insert Statement
Hi ~  
"insert into aaa(de00001) select name from a;"      for this sql string expect to be <default>.aaa.de00001<- <default>.a.name
but it turns out to be   <default>.aaa.name <- <default>.a.name, is this something wrong ?

Discussion:
ulixius9: @reata adding the query to the graph makes sense to me, but I'm not sure if it should be put as a write node, to me it makes sense to put the query as a read node.

also regarding the `HAS_COLUMN` edge type would that logic also work when the column names are not same?

for example `insert into target_table(col1, col2) as (select c1, c2 from src_table)` this is a valid query and as I've added this in tests as well that `c1` will form lineage with `col1` and `c2` with `col2`. 

reata: > @reata adding the query to the graph makes sense to me, but I'm not sure if it should be put as a write node, to me it makes sense to put the query as a read node in this particular scenario.
> 
> also regarding the `HAS_COLUMN` edge type would that logic also work when the column names are not same?
> 
> for example `insert into target_table(col1, col2) as (select c1, c2 from src_table)` this is a valid query and as I've added this in tests as well that `c1` will form lineage with `col1` and `c2` with `col2`.

Made some edit to my previous comment, SubQueryLineageHolder.add_column_lineage might not be the right place for this.

Here's what my instinct tells me now. Again, this time it could still be wrong. Just for your reference:

1. while handling (col1, col2), we should be able to get target_table from holder.write. We can build the HAS_COLUMN relation of target_table -> col1 and  target_table -> col2. Put the index information as edge attributes so we know col1 is the first column and col2 is the second.
2. wrap the corresponding graph query into a method in class SubQueryLineageHolder, you can even call it target_columns and put a @property decorator, so it still looks like an attribute.
3. everything else is basically just what you do right now. 
ulixius9: @reata make sense let me try it out and get back to you on that! 

Patch:
diff --git a/sqllineage/core/holders.py b/sqllineage/core/holders.py
index 7434522..24837cf 100644
--- a/sqllineage/core/holders.py
+++ b/sqllineage/core/holders.py
@@ -1,11 +1,11 @@
 import itertools
-from typing import Set, Tuple, Union
+from typing import List, Set, Tuple, Union
 
 import networkx as nx
 from networkx import DiGraph
 
 from sqllineage.core.models import Column, Path, SubQuery, Table
-from sqllineage.utils.constant import EdgeType, NodeTag
+from sqllineage.utils.constant import EdgeTag, EdgeType, NodeTag
 
 DATASET_CLASSES = (Path, Table)
 
@@ -84,6 +84,35 @@ class SubQueryLineageHolder(ColumnLineageMixin):
     def add_cte(self, value) -> None:
         self._property_setter(value, NodeTag.CTE)
 
+    @property
+    def target_columns(self) -> List[Column]:
+        """
+        in case of DML with column specified, like INSERT INTO tab1 (col1, col2) SELECT ...
+        target_columns tell us that tab1 has column col1 and col2 in that order.
+        """
+        tgt_cols = []
+        if self.write:
+            tgt_tbl = list(self.write)[0]
+            tgt_col_with_idx: List[Tuple[Column, int]] = sorted(
+                [
+                    (col, attr.get(EdgeTag.INDEX, 0))
+                    for tbl, col, attr in self.graph.out_edges(tgt_tbl, data=True)
+                    if attr["type"] == EdgeType.HAS_COLUMN
+                ],
+                key=lambda x: x[1],
+            )
+            tgt_cols = [x[0] for x in tgt_col_with_idx]
+        return tgt_cols
+
+    def add_target_column(self, *tgt_cols: Column) -> None:
+        if self.write:
+            tgt_tbl = list(self.write)[0]
+            for idx, tgt_col in enumerate(tgt_cols):
+                tgt_col.parent = tgt_tbl
+                self.graph.add_edge(
+                    tgt_tbl, tgt_col, type=EdgeType.HAS_COLUMN, **{EdgeTag.INDEX: idx}
+                )
+
     def add_column_lineage(self, src: Column, tgt: Column) -> None:
         self.graph.add_edge(src, tgt, type=EdgeType.LINEAGE)
         self.graph.add_edge(tgt.parent, tgt, type=EdgeType.HAS_COLUMN)
diff --git a/sqllineage/core/models.py b/sqllineage/core/models.py
index 4df575c..c5eb25f 100644
--- a/sqllineage/core/models.py
+++ b/sqllineage/core/models.py
@@ -225,21 +225,22 @@ class AnalyzerContext:
     Data class to hold the analyzer context
     """
 
-    subquery: Optional[SubQuery]
-    prev_cte: Optional[Set[SubQuery]]
-    prev_write: Optional[Set[Union[SubQuery, Table]]]
-
     def __init__(
         self,
         subquery: Optional[SubQuery] = None,
         prev_cte: Optional[Set[SubQuery]] = None,
         prev_write: Optional[Set[Union[SubQuery, Table]]] = None,
+        target_columns=None,
     ):
         """
         :param subquery: subquery
         :param prev_cte: previous CTE queries
         :param prev_write: previous written tables
+        :param target_columns: previous target columns
         """
+        if target_columns is None:
+            target_columns = []
         self.subquery = subquery
         self.prev_cte = prev_cte
         self.prev_write = prev_write
+        self.target_columns = target_columns
diff --git a/sqllineage/core/parser/__init__.py b/sqllineage/core/parser/__init__.py
index 3b2a142..2719084 100644
--- a/sqllineage/core/parser/__init__.py
+++ b/sqllineage/core/parser/__init__.py
@@ -21,17 +21,23 @@ class SourceHandlerMixin:
             )
             col_grp = self.columns[prev_col_barrier:col_barrier]
             tbl_grp = self.tables[prev_tbl_barrier:tbl_barrier]
-            tgt_tbl = None
             if holder.write:
                 if len(holder.write) > 1:
                     raise SQLLineageException
                 tgt_tbl = list(holder.write)[0]
-            if tgt_tbl:
-                for tgt_col in col_grp:
+                for idx in range(len(col_grp)):
+                    tgt_col = col_grp[idx]
                     tgt_col.parent = tgt_tbl
                     for src_col in tgt_col.to_source_columns(
                         self.get_alias_mapping_from_table_group(tbl_grp, holder)
                     ):
+                        if len(target_columns := holder.target_columns) == len(col_grp):
+                            # example query: create view test (col3) select col1 as col2 from tab
+                            # without target_columns = [col3] information, by default src_col = col1 and tgt_col = col2
+                            # when target_columns exist and length matches, we want tgt_col = col3 instead of col2
+                            # for invalid query: create view test (col3, col4) select col1 as col2 from tab,
+                            # when the length doesn't match, we fall back to default behavior
+                            tgt_col = target_columns[idx]
                         holder.add_column_lineage(src_col, tgt_col)
 
     @classmethod
diff --git a/sqllineage/core/parser/sqlfluff/extractors/ddl_alter_extractor.py b/sqllineage/core/parser/sqlfluff/extractors/ddl_alter_extractor.py
index 83a8a21..30d6ccd 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/ddl_alter_extractor.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/ddl_alter_extractor.py
@@ -13,6 +13,8 @@ class DdlAlterExtractor(LineageHolderExtractor):
     DDL Alter queries lineage extractor
     """
 
+    SOURCE_KEYWORDS = {"EXCHANGE", "SWAP"}
+
     SUPPORTED_STMT_TYPES = [
         "alter_table_statement",
         "rename_statement",
@@ -44,7 +46,10 @@ class DdlAlterExtractor(LineageHolderExtractor):
         if any(k.raw_upper == "RENAME" for k in keywords):
             if statement.type == "alter_table_statement" and len(tables) == 2:
                 holder.add_rename(tables[0], tables[1])
-        if any(k.raw_upper == "EXCHANGE" for k in keywords) and len(tables) == 2:
+        if (
+            any(k.raw_upper in self.SOURCE_KEYWORDS for k in keywords)
+            and len(tables) == 2
+        ):
             holder.add_write(tables[0])
             holder.add_read(tables[1])
         return holder
diff --git a/sqllineage/core/parser/sqlfluff/extractors/dml_insert_extractor.py b/sqllineage/core/parser/sqlfluff/extractors/dml_insert_extractor.py
index dd2202a..4e4943c 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/dml_insert_extractor.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/dml_insert_extractor.py
@@ -1,3 +1,5 @@
+from typing import Optional
+
 from sqlfluff.core.parser import BaseSegment
 
 from sqllineage.core.holders import SubQueryLineageHolder
@@ -9,7 +11,7 @@ from sqllineage.core.parser.sqlfluff.extractors.lineage_holder_extractor import
     LineageHolderExtractor,
 )
 from sqllineage.core.parser.sqlfluff.models import SqlFluffSubQuery
-from sqllineage.core.parser.sqlfluff.utils import retrieve_segments
+from sqllineage.core.parser.sqlfluff.utils import get_child, is_union, retrieve_segments
 
 
 class DmlInsertExtractor(LineageHolderExtractor):
@@ -48,14 +50,8 @@ class DmlInsertExtractor(LineageHolderExtractor):
 
         holder = self._init_holder(context)
 
-        subqueries = []
         segments = retrieve_segments(statement)
         for segment in segments:
-            for sq in self.parse_subquery(segment):
-                # Collecting subquery on the way, hold on parsing until last
-                # so that each handler don't have to worry about what's inside subquery
-                subqueries.append(sq)
-
             for current_handler in handlers:
                 current_handler.handle(segment, holder)
 
@@ -66,40 +62,46 @@ class DmlInsertExtractor(LineageHolderExtractor):
                     segment,
                     AnalyzerContext(prev_cte=holder.cte, prev_write=holder.write),
                 )
+            elif segment.type == "bracketed" and (
+                self.parse_subquery(segment) or is_union(segment)
+            ):
+                # note regular subquery within SELECT statement is handled by DmlSelectExtractor, this is only to handle
+                # top-level subquery in DML like: 1) create table foo as (subquery); 2) insert into foo (subquery)
+                # subquery here isn't added as read source, and it inherits DML-level target_columns if parsed
+                if subquery_segment := get_child(segment, "select_statement"):
+                    self._extract_select(holder, subquery_segment)
+                elif subquery_segment := get_child(segment, "set_expression"):
+                    self._extract_set(holder, subquery_segment)
+            elif segment.type == "select_statement":
+                self._extract_select(holder, segment)
+            elif segment.type == "set_expression":
+                self._extract_set(holder, segment)
+            else:
+                for conditional_handler in conditional_handlers:
+                    if conditional_handler.indicate(segment):
+                        conditional_handler.handle(segment, holder)
 
-            if segment.type == "select_statement":
-                holder |= DmlSelectExtractor(self.dialect).extract(
-                    segment,
-                    AnalyzerContext(
-                        SqlFluffSubQuery.of(segment, None),
-                        prev_cte=holder.cte,
-                        prev_write=holder.write,
-                    ),
-                )
-                continue
-
-            if segment.type == "set_expression":
-                sub_segments = retrieve_segments(segment)
-                for sub_segment in sub_segments:
-                    if sub_segment.type == "select_statement":
-                        holder |= DmlSelectExtractor(self.dialect).extract(
-                            sub_segment,
-                            AnalyzerContext(
-                                SqlFluffSubQuery.of(segment, None),
-                                prev_cte=holder.cte,
-                                prev_write=holder.write,
-                            ),
-                        )
-                continue
-
-            for conditional_handler in conditional_handlers:
-                if conditional_handler.indicate(segment):
-                    conditional_handler.handle(segment, holder)
+        return holder
 
-        # By recursively extracting each subquery of the parent and merge, we're doing Depth-first search
-        for sq in subqueries:
-            holder |= DmlSelectExtractor(self.dialect).extract(
-                sq.query, AnalyzerContext(sq, holder.cte)
-            )
+    def _extract_set(self, holder: SubQueryLineageHolder, set_segment: BaseSegment):
+        for sub_segment in retrieve_segments(set_segment):
+            if sub_segment.type == "select_statement":
+                self._extract_select(holder, sub_segment, set_segment)
 
-        return holder
+    def _extract_select(
+        self,
+        holder: SubQueryLineageHolder,
+        select_segment: BaseSegment,
+        set_segment: Optional[BaseSegment] = None,
+    ):
+        holder |= DmlSelectExtractor(self.dialect).extract(
+            select_segment,
+            AnalyzerContext(
+                SqlFluffSubQuery.of(
+                    set_segment if set_segment else select_segment, None
+                ),
+                prev_cte=holder.cte,
+                prev_write=holder.write,
+                target_columns=holder.target_columns,
+            ),
+        )
diff --git a/sqllineage/core/parser/sqlfluff/extractors/dml_merge_extractor.py b/sqllineage/core/parser/sqlfluff/extractors/dml_merge_extractor.py
index 682fea8..c265187 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/dml_merge_extractor.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/dml_merge_extractor.py
@@ -4,6 +4,7 @@ from sqlfluff.core.parser import BaseSegment
 
 from sqllineage.core.holders import StatementLineageHolder, SubQueryLineageHolder
 from sqllineage.core.models import AnalyzerContext, Column, SubQuery, Table
+from sqllineage.core.parser.sqlfluff.extractors.cte_extractor import DmlCteExtractor
 from sqllineage.core.parser.sqlfluff.extractors.dml_select_extractor import (
     DmlSelectExtractor,
 )
@@ -40,7 +41,7 @@ class DmlMergeExtractor(LineageHolderExtractor):
         direct_source: Optional[Union[Table, SubQuery]] = None
         segments = retrieve_segments(statement)
         for i, segment in enumerate(segments):
-            if segment.type == "keyword" and segment.raw_upper == "INTO":
+            if segment.type == "keyword" and segment.raw_upper in {"INTO", "MERGE"}:
                 tgt_flag = True
                 continue
             if segment.type == "keyword" and segment.raw_upper == "USING":
@@ -64,9 +65,18 @@ class DmlMergeExtractor(LineageHolderExtractor):
                         else None,
                     )
                     holder.add_read(direct_source)
-                    holder |= DmlSelectExtractor(self.dialect).extract(
-                        direct_source.query, AnalyzerContext(direct_source, holder.cte)
-                    )
+                    if direct_source.query.get_child("with_compound_statement"):
+                        # in case the subquery is a CTE query
+                        holder |= DmlCteExtractor(self.dialect).extract(
+                            direct_source.query,
+                            AnalyzerContext(direct_source, prev_cte=holder.cte),
+                        )
+                    else:
+                        # in case the subquery is a select query
+                        holder |= DmlSelectExtractor(self.dialect).extract(
+                            direct_source.query,
+                            AnalyzerContext(direct_source, holder.cte),
+                        )
                 src_flag = False
 
         for match in get_grandchildren(
@@ -90,19 +100,23 @@ class DmlMergeExtractor(LineageHolderExtractor):
         ):
             merge_insert = not_match.get_child("merge_insert_clause")
             insert_columns = []
-            for c in merge_insert.get_child("bracketed").get_children(
+            merge_insert_bracketed = merge_insert.get_child("bracketed")
+            if merge_insert_bracketed and merge_insert_bracketed.get_children(
                 "column_reference"
             ):
-                tgt_col = Column(get_identifier(c))
-                tgt_col.parent = list(holder.write)[0]
-                insert_columns.append(tgt_col)
-            for j, e in enumerate(
-                merge_insert.get_child("values_clause")
-                .get_child("bracketed")
-                .get_children("expression")
-            ):
-                if col_ref := e.get_child("column_reference"):
-                    src_col = Column(get_identifier(col_ref))
-                    src_col.parent = direct_source
-                    holder.add_column_lineage(src_col, insert_columns[j])
+                for c in merge_insert.get_child("bracketed").get_children(
+                    "column_reference"
+                ):
+                    tgt_col = Column(get_identifier(c))
+                    tgt_col.parent = list(holder.write)[0]
+                    insert_columns.append(tgt_col)
+                for j, e in enumerate(
+                    merge_insert.get_child("values_clause")
+                    .get_child("bracketed")
+                    .get_children("expression")
+                ):
+                    if col_ref := e.get_child("column_reference"):
+                        src_col = Column(get_identifier(col_ref))
+                        src_col.parent = direct_source
+                        holder.add_column_lineage(src_col, insert_columns[j])
         return holder
diff --git a/sqllineage/core/parser/sqlfluff/extractors/lineage_holder_extractor.py b/sqllineage/core/parser/sqlfluff/extractors/lineage_holder_extractor.py
index b6ec5ea..337a56b 100644
--- a/sqllineage/core/parser/sqlfluff/extractors/lineage_holder_extractor.py
+++ b/sqllineage/core/parser/sqlfluff/extractors/lineage_holder_extractor.py
@@ -130,4 +130,8 @@ class LineageHolderExtractor:
             # If within subquery, then manually add subquery as target table
             holder.add_write(context.subquery)
 
+        if context.target_columns:
+            # target columns can be referred while creating column level lineage
+            holder.add_target_column(*context.target_columns)
+
         return holder
diff --git a/sqllineage/core/parser/sqlfluff/handlers/target.py b/sqllineage/core/parser/sqlfluff/handlers/target.py
index e4d3aa4..3968412 100644
--- a/sqllineage/core/parser/sqlfluff/handlers/target.py
+++ b/sqllineage/core/parser/sqlfluff/handlers/target.py
@@ -5,6 +5,7 @@ from sqllineage.core.models import Path
 from sqllineage.core.parser.sqlfluff.handlers.base import ConditionalSegmentBaseHandler
 from sqllineage.core.parser.sqlfluff.holder_utils import retrieve_holder_data_from
 from sqllineage.core.parser.sqlfluff.models import (
+    SqlFluffColumn,
     SqlFluffTable,
 )
 from sqllineage.core.parser.sqlfluff.utils import (
@@ -22,7 +23,7 @@ class TargetHandler(ConditionalSegmentBaseHandler):
 
     def __init__(self) -> None:
         self.indicator = False
-        self.prev_token_like = False
+        self.prev_token_read = False
         self.prev_token_from = False
 
     TARGET_KEYWORDS = (
@@ -35,7 +36,7 @@ class TargetHandler(ConditionalSegmentBaseHandler):
         "DIRECTORY",
     )
 
-    LIKE_KEYWORD = "LIKE"
+    READ_KEYWORDS = {"LIKE", "CLONE"}
 
     FROM_KEYWORD = "FROM"
 
@@ -44,17 +45,17 @@ class TargetHandler(ConditionalSegmentBaseHandler):
         Check if the segment is a 'LIKE' or 'FROM' keyword
         :param segment: segment to be use for checking
         """
-        if segment.raw_upper == self.LIKE_KEYWORD:
-            self.prev_token_like = True
+        if segment.raw_upper in self.READ_KEYWORDS:
+            self.prev_token_read = True
 
         if segment.raw_upper == self.FROM_KEYWORD:
             self.prev_token_from = True
 
     def _reset_tokens(self) -> None:
         """
-        Set 'prev_token_like' and 'prev_token_like' variable to False
+        Set 'prev_token_from' and 'prev_token_read' variable to False
         """
-        self.prev_token_like = False
+        self.prev_token_read = False
         self.prev_token_from = False
 
     def indicate(self, segment: BaseSegment) -> bool:
@@ -81,7 +82,7 @@ class TargetHandler(ConditionalSegmentBaseHandler):
         """
         if segment.type == "table_reference":
             write_obj = SqlFluffTable.of(segment)
-            if self.prev_token_like:
+            if self.prev_token_read:
                 holder.add_read(write_obj)
             else:
                 holder.add_write(write_obj)
@@ -135,3 +136,18 @@ class TargetHandler(ConditionalSegmentBaseHandler):
                     )
                     if read:
                         holder.add_read(read)
+
+        elif segment.type == "bracketed":
+            """
+            In case of bracketed column reference, add these target columns to holder
+            so that when we compute the column level lineage
+            we keep these columns into consideration
+            """
+            sub_segments = retrieve_segments(segment)
+            if all(
+                sub_segment.type == "column_reference" for sub_segment in sub_segments
+            ):
+                # target columns only apply to bracketed column references
+                holder.add_target_column(
+                    *[SqlFluffColumn.of(sub_segment) for sub_segment in sub_segments]
+                )
diff --git a/sqllineage/core/parser/sqlparse/analyzer.py b/sqllineage/core/parser/sqlparse/analyzer.py
index fcbbd4d..2311077 100644
--- a/sqllineage/core/parser/sqlparse/analyzer.py
+++ b/sqllineage/core/parser/sqlparse/analyzer.py
@@ -107,7 +107,7 @@ class SqlParseLineageAnalyzer(LineageAnalyzer):
             if is_token_negligible(token):
                 continue
             if token.is_keyword:
-                if token.normalized == "INTO":
+                if token.normalized in {"INTO", "MERGE"}:
                     tgt_flag = True
                 elif token.normalized == "USING":
                     src_flag = True
@@ -160,7 +160,7 @@ class SqlParseLineageAnalyzer(LineageAnalyzer):
                         tgt_col = Column(identifier.get_real_name())
                         tgt_col.parent = list(holder.write)[0]
                         insert_columns.append(tgt_col)
-                elif isinstance(token, Values):
+                elif insert_columns and isinstance(token, Values):
                     for sub_token in token.tokens:
                         if isinstance(sub_token, Parenthesis):
                             t = sub_token.tokens[1]
diff --git a/sqllineage/utils/constant.py b/sqllineage/utils/constant.py
index 1665962..3fdc6f4 100644
--- a/sqllineage/utils/constant.py
+++ b/sqllineage/utils/constant.py
@@ -11,6 +11,10 @@ class NodeTag:
     SELFLOOP = "selfloop"
 
 
+class EdgeTag:
+    INDEX = "index"
+
+
 @unique
 class EdgeType(Enum):
     LINEAGE = 1


Test patch:
diff --git a/tests/test_insert.py b/tests/test_insert.py
index d94864c..595fb1a 100644
--- a/tests/test_insert.py
+++ b/tests/test_insert.py
@@ -5,21 +5,9 @@ def test_insert_into():
     assert_table_lineage_equal("INSERT INTO tab1 VALUES (1, 2)", set(), {"tab1"})
 
 
-def test_insert_into_with_columns():
+def test_insert_into_select():
     assert_table_lineage_equal(
-        "INSERT INTO tab1 (col1, col2) SELECT * FROM tab2;", {"tab2"}, {"tab1"}
-    )
-
-
-def test_insert_into_with_columns_and_select():
-    assert_table_lineage_equal(
-        "INSERT INTO tab1 (col1, col2) SELECT * FROM tab2", {"tab2"}, {"tab1"}
-    )
-
-
-def test_insert_into_with_columns_and_select_union():
-    assert_table_lineage_equal(
-        "INSERT INTO tab1 (col1, col2) SELECT * FROM tab2 UNION SELECT * FROM tab3",
-        {"tab2", "tab3"},
+        "INSERT INTO tab1 SELECT * FROM tab2;",
+        {"tab2"},
         {"tab1"},
     )
diff --git a/tests/test_others.py b/tests/test_others.py
index a5e7349..0222e81 100644
--- a/tests/test_others.py
+++ b/tests/test_others.py
@@ -38,22 +38,12 @@ def test_create_as():
     )
 
 
-def test_create_as_with_parenthesis_around_select_statement():
-    sql = "CREATE TABLE tab1 AS (SELECT * FROM tab2)"
-    assert_table_lineage_equal(sql, {"tab2"}, {"tab1"})
-
-
 def test_create_as_with_parenthesis_around_table_name():
     assert_table_lineage_equal(
         "CREATE TABLE tab1 AS SELECT * FROM (tab2)", {"tab2"}, {"tab1"}
     )
 
 
-def test_create_as_with_parenthesis_around_both():
-    sql = "CREATE TABLE tab1 AS (SELECT * FROM (tab2))"
-    assert_table_lineage_equal(sql, {"tab2"}, {"tab1"})
-
-
 def test_create_like():
     assert_table_lineage_equal("CREATE TABLE tab1 LIKE tab2", {"tab2"}, {"tab1"})
 
@@ -211,3 +201,33 @@ def test_split_statements_with_desc():
 
 DESC tab1;"""
     assert len(split(sql)) == 2
+
+
+def test_merge_using_cte_subquery():
+    sql = """MERGE INTO tgt t
+USING (
+    WITH base AS (
+        SELECT
+            id, max(value) AS value
+        FROM src
+        GROUP BY id
+    )
+    SELECT
+        id, value
+    FROM base
+) s
+ON t.id = s.id
+WHEN MATCHED THEN
+UPDATE SET t.value = s.value"""
+    assert_table_lineage_equal(
+        sql,
+        {"src"},
+        {"tgt"},
+    )
+
+
+def test_merge_into_insert_one_column():
+    sql = """MERGE INTO target
+USING src ON target.k = src.k
+WHEN NOT MATCHED THEN INSERT VALUES (src.k)"""
+    assert_table_lineage_equal(sql, {"src"}, {"target"})
diff --git a/tests/test_others_dialect_specific.py b/tests/test_others_dialect_specific.py
index 3ca5c9b..c7cbc8f 100644
--- a/tests/test_others_dialect_specific.py
+++ b/tests/test_others_dialect_specific.py
@@ -94,6 +94,38 @@ def test_alter_table_exchange_partition(dialect: str):
     )
 
 
+@pytest.mark.parametrize("dialect", ["snowflake", "bigquery"])
+def test_create_clone(dialect: str):
+    """
+    Language manual:
+        https://cloud.google.com/bigquery/docs/table-clones-create
+        https://docs.snowflake.com/en/sql-reference/sql/create-clone
+    Note clone is not a keyword in sqlparse, we'll skip testing for it.
+    """
+    assert_table_lineage_equal(
+        "create table tab2 CLONE tab1;",
+        {"tab1"},
+        {"tab2"},
+        dialect=dialect,
+        test_sqlparse=False,
+    )
+
+
+@pytest.mark.parametrize("dialect", ["snowflake"])
+def test_alter_table_swap_partition(dialect: str):
+    """
+    See https://docs.snowflake.com/en/sql-reference/sql/alter-table for language manual
+    Note swap is not a keyword in sqlparse, we'll skip testing for it.
+    """
+    assert_table_lineage_equal(
+        "alter table tab1 swap with tab2",
+        {"tab2"},
+        {"tab1"},
+        dialect=dialect,
+        test_sqlparse=False,
+    )
+
+
 @pytest.mark.parametrize("dialect", ["databricks", "sparksql"])
 def test_refresh_table(dialect: str):
     assert_table_lineage_equal("refresh table tab1", None, None, dialect)
@@ -137,3 +169,36 @@ LATERAL VIEW OUTER explode(sc.json_array) q AS col1"""
 @pytest.mark.parametrize("dialect", ["databricks", "sparksql"])
 def test_show_create_table(dialect: str):
     assert_table_lineage_equal("show create table tab1", None, None, dialect)
+
+
+@pytest.mark.parametrize("dialect", ["bigquery"])
+def test_merge_without_into(dialect: str):
+    """
+    INTO is optional in BigQuery MERGE statement:
+    https://cloud.google.com/bigquery/docs/reference/standard-sql/dml-syntax#merge_statement
+    """
+    sql = """MERGE target USING src AS b ON target.k = b.k
+WHEN MATCHED THEN UPDATE SET target.v = b.v
+WHEN NOT MATCHED THEN INSERT (k, v) VALUES (b.k, b.v)"""
+    assert_table_lineage_equal(sql, {"src"}, {"target"}, dialect=dialect)
+
+
+@pytest.mark.parametrize("dialect", ["bigquery"])
+def test_merge_insert_row(dialect: str):
+    """
+    MERGE INSERT CLAUSE in BigQuery can be INSERT ROW without specifying columns via INSERT VALUES (col, ...)
+    https://cloud.google.com/bigquery/docs/reference/standard-sql/dml-syntax#merge_statement
+    """
+    sql = """MERGE INTO tgt t
+USING src s
+ON t.date = s.date and t.channel = s.channel
+WHEN NOT MATCHED THEN
+INSERT ROW
+WHEN MATCHED THEN
+UPDATE SET t.col = s.col"""
+    assert_table_lineage_equal(
+        sql,
+        {"src"},
+        {"tgt"},
+        dialect=dialect,
+    )
diff --git a/tests/test_sqlfluff.py b/tests/test_sqlfluff.py
index 386c071..6fde533 100644
--- a/tests/test_sqlfluff.py
+++ b/tests/test_sqlfluff.py
@@ -102,3 +102,131 @@ FROM tab2 a"""
         ],
         test_sqlparse=False,
     )
+
+
+# For top-level query parenthesis in DML, we don't treat it as subquery.
+# sqlparse has some problems identifying these subqueries.
+# note the table-level lineage works, only column-level lineage breaks for sqlparse
+def test_create_as_with_parenthesis_around_select_statement():
+    sql = "CREATE TABLE tab1 AS (SELECT * FROM tab2)"
+    assert_table_lineage_equal(sql, {"tab2"}, {"tab1"}, test_sqlparse=False)
+
+
+def test_create_as_with_parenthesis_around_both():
+    sql = "CREATE TABLE tab1 AS (SELECT * FROM (tab2))"
+    assert_table_lineage_equal(sql, {"tab2"}, {"tab1"}, test_sqlparse=False)
+
+
+# specify columns in CREATE statement, sqlparse would parse my_view as function call
+def test_view_with_subquery_custom_columns():
+    # select as subquery
+    sql = "CREATE VIEW my_view (random1,random2) AS (SELECT col1,col2 FROM tbl)"
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("col1", "tbl"),
+                ColumnQualifierTuple("random1", "my_view"),
+            ),
+            (
+                ColumnQualifierTuple("col2", "tbl"),
+                ColumnQualifierTuple("random2", "my_view"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
+    sql = "CREATE VIEW my_view (random1,random2) AS SELECT col1,col2 FROM tbl"
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("col1", "tbl"),
+                ColumnQualifierTuple("random1", "my_view"),
+            ),
+            (
+                ColumnQualifierTuple("col2", "tbl"),
+                ColumnQualifierTuple("random2", "my_view"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
+
+
+def test_create_view_with_same_columns():
+    sql = "CREATE VIEW my_view (col1,col2) AS (SELECT col1,col2 FROM tbl)"
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("col1", "tbl"),
+                ColumnQualifierTuple("col1", "my_view"),
+            ),
+            (
+                ColumnQualifierTuple("col2", "tbl"),
+                ColumnQualifierTuple("col2", "my_view"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
+
+
+# specify columns in INSERT statement, sqlparse would parse my_view as function call
+def test_insert_into_with_columns():
+    # table lineage works for sqlparse
+    assert_table_lineage_equal(
+        "INSERT INTO tab1 (col1, col2) SELECT * FROM tab2;",
+        {"tab2"},
+        {"tab1"},
+        test_sqlparse=False,
+    )
+
+
+def test_insert_into_with_columns_and_select_union():
+    assert_table_lineage_equal(
+        "INSERT INTO tab1 (col1, col2) SELECT * FROM tab2 UNION SELECT * FROM tab3",
+        {"tab2", "tab3"},
+        {"tab1"},
+        test_sqlparse=False,
+    )
+    assert_table_lineage_equal(
+        "INSERT INTO tab1 (col1, col2) (SELECT * FROM tab2 UNION SELECT * FROM tab3)",
+        {"tab2", "tab3"},
+        {"tab1"},
+        test_sqlparse=False,
+    )
+
+
+def test_insert_with_custom_columns():
+    # test with query as subquery
+    sql = "INSERT INTO tgt_tbl(random1, random2) (SELECT col1,col2 FROM src_tbl)"
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("col1", "src_tbl"),
+                ColumnQualifierTuple("random1", "tgt_tbl"),
+            ),
+            (
+                ColumnQualifierTuple("col2", "src_tbl"),
+                ColumnQualifierTuple("random2", "tgt_tbl"),
+            ),
+        ],
+        test_sqlparse=False,
+    )
+
+    # test with plain query
+    sql = "INSERT INTO tgt_tbl(random1, random2) SELECT col1,col2 FROM src_tbl_new"
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("col1", "src_tbl_new"),
+                ColumnQualifierTuple("random1", "tgt_tbl"),
+            ),
+            (
+                ColumnQualifierTuple("col2", "src_tbl_new"),
+                ColumnQualifierTuple("random2", "tgt_tbl"),
+            ),
+        ],
+        test_sqlparse=False,
+    )