From 9a78fc2658d3da8838c93381e10b4142405fb2b6 Mon Sep 17 00:00:00 2001
From: Oliver Mannion <125105+tekumara@users.noreply.github.com>
Date: Wed, 17 Jul 2024 16:27:08 +1000
Subject: [PATCH] feat: alter table cluster by

---
 fakesnow/fakes.py        |  1 +
 fakesnow/transforms.py   | 14 +++++++++++++-
 tests/test_fakes.py      | 11 +++++++----
 tests/test_transforms.py | 16 ++++++++++++----
 4 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/fakesnow/fakes.py b/fakesnow/fakes.py
index c453512..fcd6a1d 100644
--- a/fakesnow/fakes.py
+++ b/fakesnow/fakes.py
@@ -203,6 +203,7 @@ def _transform(self, expression: exp.Expression) -> exp.Expression:
             .transform(transforms.sha256)
             .transform(transforms.create_clone)
             .transform(transforms.alias_in_join)
+            .transform(transforms.alter_table_strip_cluster_by)
         )
 
     def _execute(
diff --git a/fakesnow/transforms.py b/fakesnow/transforms.py
index 6aca772..29a848d 100644
--- a/fakesnow/transforms.py
+++ b/fakesnow/transforms.py
@@ -11,7 +11,7 @@
 from fakesnow.variables import Variables
 
 MISSING_DATABASE = "missing_database"
-SUCCESS_NOP = sqlglot.parse_one("SELECT 'Statement executed successfully.'")
+SUCCESS_NOP = sqlglot.parse_one("SELECT 'Statement executed successfully.' as status")
 
 
 def alias_in_join(expression: exp.Expression) -> exp.Expression:
@@ -33,6 +33,18 @@ def alias_in_join(expression: exp.Expression) -> exp.Expression:
     return expression
 
 
+def alter_table_strip_cluster_by(expression: exp.Expression) -> exp.Expression:
+    """Turn alter table cluster by into a no-op"""
+    if (
+        isinstance(expression, exp.AlterTable)
+        and (actions := expression.args.get("actions"))
+        and len(actions) == 1
+        and (isinstance(actions[0], exp.Cluster))
+    ):
+        return SUCCESS_NOP
+    return expression
+
+
 def array_size(expression: exp.Expression) -> exp.Expression:
     if isinstance(expression, exp.ArraySize):
         # case is used to convert 0 to null, because null is returned by duckdb when no case matches
diff --git a/tests/test_fakes.py b/tests/test_fakes.py
index 7289aba..ad35e99 100644
--- a/tests/test_fakes.py
+++ b/tests/test_fakes.py
@@ -41,10 +41,13 @@ def test_alias_on_join(conn: snowflake.connector.SnowflakeConnection):
     assert cur.fetchall() == [("VARCHAR1", "CHAR1", "JOIN"), ("VARCHAR2", "CHAR2", None)]
 
 
-def test_alter_table(cur: snowflake.connector.cursor.SnowflakeCursor):
-    cur.execute("create table table1 (id int)")
-    cur.execute("alter table table1 add column name varchar(20)")
-    cur.execute("select name from table1")
+def test_alter_table(dcur: snowflake.connector.cursor.SnowflakeCursor):
+    dcur.execute("create table table1 (id int)")
+    dcur.execute("alter table table1 add column name varchar(20)")
+    dcur.execute("select name from table1")
+    assert dcur.execute("alter table table1 cluster by (name)").fetchall() == [
+        {"status": "Statement executed successfully."}
+    ]
 
 
 def test_array_size(cur: snowflake.connector.cursor.SnowflakeCursor):
diff --git a/tests/test_transforms.py b/tests/test_transforms.py
index 2935af4..257179c 100644
--- a/tests/test_transforms.py
+++ b/tests/test_transforms.py
@@ -8,6 +8,7 @@
     SUCCESS_NOP,
     _get_to_number_args,
     alias_in_join,
+    alter_table_strip_cluster_by,
     array_agg,
     array_agg_within_group,
     array_size,
@@ -74,6 +75,13 @@ def test_alias_in_join() -> None:
     )
 
 
+def test_alter_table_strip_cluster_by() -> None:
+    assert (
+        sqlglot.parse_one("alter table table1 cluster by (name)").transform(alter_table_strip_cluster_by).sql()
+        == "SELECT 'Statement executed successfully.' AS status"
+    )
+
+
 def test_array_size() -> None:
     assert (
         sqlglot.parse_one("""select array_size(parse_json('["a","b"]'))""").transform(array_size).sql(dialect="duckdb")
@@ -342,7 +350,7 @@ def test_extract_comment_on_columns() -> None:
     e = sqlglot.parse_one("ALTER TABLE ingredients ALTER amount COMMENT 'tablespoons'").transform(
         extract_comment_on_columns
     )
-    assert e.sql() == "SELECT 'Statement executed successfully.'"
+    assert e.sql() == "SELECT 'Statement executed successfully.' AS status"
     assert e.args["col_comments"] == [("amount", "tablespoons")]
 
     # TODO
@@ -365,19 +373,19 @@ def test_extract_comment_on_table() -> None:
     assert e.args["table_comment"] == (table1, "foobar")
 
     e = sqlglot.parse_one("COMMENT ON TABLE table1 IS 'comment1'").transform(extract_comment_on_table)
-    assert e.sql() == "SELECT 'Statement executed successfully.'"
+    assert e.sql() == "SELECT 'Statement executed successfully.' AS status"
     assert e.args["table_comment"] == (table1, "comment1")
 
     e = sqlglot.parse_one("COMMENT ON TABLE table1 IS $$comment2$$", read="snowflake").transform(
         extract_comment_on_table
     )
-    assert e.sql() == "SELECT 'Statement executed successfully.'"
+    assert e.sql() == "SELECT 'Statement executed successfully.' AS status"
     assert e.args["table_comment"] == (table1, "comment2")
 
     e = sqlglot.parse_one("ALTER TABLE table1 SET COMMENT = 'comment1'", read="snowflake").transform(
         extract_comment_on_table
     )
-    assert e.sql() == "SELECT 'Statement executed successfully.'"
+    assert e.sql() == "SELECT 'Statement executed successfully.' AS status"
     assert e.args["table_comment"] == (table1, "comment1")