feat: Add support for multi-column ALTER TABLE ADD COLUMN statements (#247)

# feat: Add support for multi-column ALTER TABLE ADD COLUMN statements

## Overview

This PR introduces support for Snowflake's multi-column `ALTER TABLE ADD
COLUMN` syntax, allowing users to add multiple columns in a single
statement. The implementation automatically splits these statements into
individual DuckDB-compatible `ALTER TABLE` commands.

Resolves #244

Co-authored-by: Chetan Agrawal <chetan.agrawal@moveworks.ai>
Co-authored-by: Oliver Mannion <125105+tekumara@users.noreply.github.com>

Author: 3 people

fakesnow/cursor.py

@@ -317,7 +317,19 @@ def _transform(self, expression: exp.Expression, params: MutableParams | None) -
     def _transform_explode(self, expression: exp.Expression) -> list[exp.Expression]:
         # Applies transformations that require splitting the expression into multiple expressions
         # Split transforms have limited support at the moment.
-        return transforms.merge(expression)
+
+        # Try merge transform first
+        merge_result = transforms.merge(expression)
+        if len(merge_result) > 1:
+            return merge_result
+
+        # If merge didn't split the expression, try alter_table_add_multiple_columns
+        alter_result = transforms.alter_table_add_multiple_columns(expression)
+        if len(alter_result) > 1:
+            return alter_result
+
+        # Return original expression if no transform applied
+        return [expression]
 
     def _execute(self, transformed: exp.Expression, params: MutableParams | None = None) -> None:
         self._arrow_table = None

fakesnow/transforms/__init__.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+from fakesnow.transforms.ddl import (
+    alter_table_add_multiple_columns as alter_table_add_multiple_columns,
+)
 from fakesnow.transforms.merge import merge as merge
 from fakesnow.transforms.show import (
     show_columns as show_columns,

fakesnow/transforms/ddl.py

@@ -0,0 +1,55 @@
+"""DDL (Data Definition Language) transforms for fakesnow.
+
+This module contains transformations for DDL statements like ALTER TABLE, CREATE TABLE, etc.
+
+Future DDL transforms can be added here:
+- CREATE TABLE enhancements
+- DROP TABLE transformations
+- ALTER TABLE modifications (beyond ADD COLUMN)
+- INDEX operations
+- VIEW operations
+- etc.
+"""
+
+from __future__ import annotations
+
+from sqlglot import exp
+
+
+def alter_table_add_multiple_columns(expression: exp.Expression) -> list[exp.Expression]:
+    """Transform ALTER TABLE ADD COLUMN with multiple columns into separate statements.
+
+    Snowflake supports: ALTER TABLE tab1 ADD COLUMN col1 INT, col2 VARCHAR(50), col3 BOOLEAN;
+    DuckDB requires separate statements for each column.
+
+    Args:
+        expression: The expression to potentially transform
+
+    Returns:
+        List of expressions - multiple ALTER statements if transformation applied,
+        otherwise single original expression
+    """
+    if not isinstance(expression, exp.Alter):
+        return [expression]
+
+    actions = expression.args.get("actions")
+    if not (
+        actions
+        and isinstance(actions, list)
+        and all(isinstance(action, exp.ColumnDef) for action in actions)
+        and len(actions) > 1
+    ):
+        return [expression]
+
+    # Create separate ALTER statements for each column
+    alter_statements = []
+    for action in actions:
+        new_alter = exp.Alter(
+            this=expression.this,
+            kind=expression.args.get("kind"),
+            exists=expression.args.get("exists"),
+            actions=[action],
+        )
+        alter_statements.append(new_alter)
+
+    return alter_statements

tests/test_ddl.py

@@ -0,0 +1,22 @@
+import sqlglot
+
+from fakesnow.transforms.ddl import alter_table_add_multiple_columns
+
+
+def test_alter_table_add_multiple_columns() -> None:
+    # Test multiple columns are split into separate statements
+    sql = "alter table tab1 add column col1 int, col2 varchar(50), col3 boolean"
+    result = alter_table_add_multiple_columns(sqlglot.parse_one(sql, dialect="snowflake"))
+
+    assert len(result) == 3
+    assert result[0].sql() == "ALTER TABLE tab1 ADD COLUMN col1 INT"
+    assert result[1].sql() == "ALTER TABLE tab1 ADD COLUMN col2 VARCHAR(50)"
+    assert result[2].sql() == "ALTER TABLE tab1 ADD COLUMN col3 BOOLEAN"
+
+    # Test multiple columns with IF EXISTS clause
+    sql_if_exists = "alter table if exists tab1 add column col1 int, col2 varchar(50)"
+    result_if_exists = alter_table_add_multiple_columns(sqlglot.parse_one(sql_if_exists, dialect="snowflake"))
+
+    assert len(result_if_exists) == 2
+    assert result_if_exists[0].sql() == "ALTER TABLE IF EXISTS tab1 ADD COLUMN col1 INT"
+    assert result_if_exists[1].sql() == "ALTER TABLE IF EXISTS tab1 ADD COLUMN col2 VARCHAR(50)"

tests/test_fakes.py

@@ -39,7 +39,8 @@ def test_alias_on_join(conn: snowflake.connector.SnowflakeConnection):
 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")
+    dcur.execute("alter table if exists table1 add column col1 int, col2 varchar(50)")
+    dcur.execute("select id,name,col1,col2 from table1")
     assert dcur.execute("alter table table1 cluster by (name)").fetchall() == [
         {"status": "Statement executed successfully."}
     ]