Repository: reata/sqllineage

Issue:
support create table..clone and alter table...swap
CREATE TABLE CLONE:
```sql
create table tab2 CLONE tab1
```
Ref:

- https://cloud.google.com/bigquery/docs/table-clones-create
- https://docs.snowflake.com/en/sql-reference/sql/create-clone

----------

ALTER TABLE SWAP WITH:
```sql
alter table tab1 swap with tab2
```
Ref:

- https://docs.snowflake.com/en/sql-reference/sql/alter-table

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/372?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#372](https://app.codecov.io/gh/reata/sqllineage/pull/372?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (e524e6d) 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     #372   +/-   ##
=======================================
  Coverage   99.39%   99.39%           
=======================================
  Files          42       42           
  Lines        1975     1976    +1     
=======================================
+ Hits         1963     1964    +1     
  Misses         12       12           
```


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


reata: > @reata I tried with the sqlparse as well but the sqlparse in case of `create table tab2 CLONE tab1` query considers the `clone` as the alias of the table and in case of `alter table tab1 swap with tab2` it considers the `with` keyword as start of an CTE so I'm not sure how to proceed as I don't have much experience with sqlparse, it would be great if you can give me a hand here.

Sure, thanks for the info. Let me try if I can find a way to deal with this.

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/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)


Test patch:
diff --git a/tests/test_others_dialect_specific.py b/tests/test_others_dialect_specific.py
index 3ca5c9b..0aac6d0 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)