Skip to content

Commit

Permalink
feat: alter table cluster by
Browse files Browse the repository at this point in the history
  • Loading branch information
tekumara committed Jul 17, 2024
1 parent 72b791b commit 9a78fc2
Show file tree
Hide file tree
Showing 4 changed files with 33 additions and 9 deletions.
1 change: 1 addition & 0 deletions fakesnow/fakes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
14 changes: 13 additions & 1 deletion fakesnow/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand Down
11 changes: 7 additions & 4 deletions tests/test_fakes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
16 changes: 12 additions & 4 deletions tests/test_transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand All @@ -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")


Expand Down

0 comments on commit 9a78fc2

Please sign in to comment.