Repository: reata/sqllineage

Issue:
BigQuery Specific MERGE statement feature support
Ref: https://cloud.google.com/bigquery/docs/reference/standard-sql/dml-syntax#merge_statement

- The INTO keyword after MERGE is optional
- The column references after INSERT are optional as INSERT ROWS
- The subquery after USING keyword can be a CTE query (this is not specific to bigquery)

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

```diff
@@           Coverage Diff           @@
##           master     #375   +/-   ##
=======================================
  Coverage   99.39%   99.39%           
=======================================
  Files          42       42           
  Lines        1975     1980    +5     
=======================================
+ Hits         1963     1968    +5     
  Misses         12       12           
```


| [Impacted Files](https://app.codecov.io/gh/reata/sqllineage/pull/375?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\_merge\_extractor.py](https://app.codecov.io/gh/reata/sqllineage/pull/375?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL3BhcnNlci9zcWxmbHVmZi9leHRyYWN0b3JzL2RtbF9tZXJnZV9leHRyYWN0b3IucHk=) | `100.00% <100.00%> (ø)` | |
| [sqllineage/core/parser/sqlparse/analyzer.py](https://app.codecov.io/gh/reata/sqllineage/pull/375?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL3BhcnNlci9zcWxwYXJzZS9hbmFseXplci5weQ==) | `100.00% <100.00%> (ø)` | |



Patch:
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_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/handlers/target.py b/sqllineage/core/parser/sqlfluff/handlers/target.py
index e4d3aa4..a1a1052 100644
--- a/sqllineage/core/parser/sqlfluff/handlers/target.py
+++ b/sqllineage/core/parser/sqlfluff/handlers/target.py
@@ -22,7 +22,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 +35,7 @@ class TargetHandler(ConditionalSegmentBaseHandler):
         "DIRECTORY",
     )
 
-    LIKE_KEYWORD = "LIKE"
+    READ_KEYWORDS = {"LIKE", "CLONE"}
 
     FROM_KEYWORD = "FROM"
 
@@ -44,17 +44,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 +81,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)
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]


Test patch:
diff --git a/tests/test_others.py b/tests/test_others.py
index a5e7349..eb618f0 100644
--- a/tests/test_others.py
+++ b/tests/test_others.py
@@ -211,3 +211,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,
+    )