feat: support AUTOINCREMENT when creating table

resolves #184

Author: tekumara

fakesnow/cursor.py

@@ -320,12 +320,14 @@ 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.
 
-        # Try merge transform first
+        auto_seq_result = transforms.create_table_autoincrement(expression)
+        if len(auto_seq_result) > 1:
+            return auto_seq_result
+
         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

fakesnow/transforms/__init__.py

@@ -3,6 +3,7 @@
 from fakesnow.transforms.ddl import (
     alter_table_add_multiple_columns as alter_table_add_multiple_columns,
     alter_table_strip_cluster_by as alter_table_strip_cluster_by,
+    create_table_autoincrement as create_table_autoincrement,
 )
 from fakesnow.transforms.merge import merge as merge
 from fakesnow.transforms.show import (

fakesnow/transforms/ddl.py

@@ -13,12 +13,17 @@
 
 from __future__ import annotations
 
+import secrets
+
+import sqlglot
 from sqlglot import exp
 
 from fakesnow.transforms.transforms import SUCCESS_NOP
 
 
-def alter_table_add_multiple_columns(expression: exp.Expression) -> list[exp.Expression]:
+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 IF EXISTS tab1 ADD COLUMN IF NOT EXISTS col1 INT, col2 VARCHAR(50), col3 BOOLEAN;
@@ -82,3 +87,99 @@ def alter_table_strip_cluster_by(expression: exp.Expression) -> exp.Expression:
     ):
         return SUCCESS_NOP
     return expression
+
+
+def create_table_autoincrement(
+    expression: exp.Expression,
+) -> list[exp.Expression]:
+    """Split CREATE TABLE with AUTOINCREMENT into CREATE SEQUENCE + CREATE TABLE with DEFAULT NEXTVAL.
+
+    Example transform:
+        CREATE TABLE test_table (id NUMERIC NOT NULL AUTOINCREMENT, name VARCHAR)
+        ->
+        CREATE SEQUENCE test_table_id_seq START 1;
+        CREATE TABLE test_table (id NUMERIC NOT NULL DEFAULT NEXTVAL('test_table_id_seq'), name VARCHAR)
+    """
+
+    if not (
+        isinstance(expression, exp.Create)
+        and expression.kind == "TABLE"
+        and (schema := expression.this)
+        and (table := schema.this)
+        and isinstance(schema, exp.Schema)
+        and isinstance(table, exp.Table)
+        # Find AUTOINCREMENT/IDENTITY columns
+        and (
+            auto_cols := [
+                cd
+                for cd in (schema.expressions or [])
+                if isinstance(cd, exp.ColumnDef)
+                and (cd.find(exp.AutoIncrementColumnConstraint) or cd.find(exp.GeneratedAsIdentityColumnConstraint))
+            ]
+        )
+    ):
+        return [expression]
+
+    if len(auto_cols) > 1:
+        raise NotImplementedError("Multiple AUTOINCREMENT columns")
+
+    auto = auto_cols[0]
+    col_name = auto.this.name
+    table_name = table.name
+    # When recreating the same table with a sequence, we need to give the sequence a unique name to avoid
+    # Dependency Error: Cannot drop entry "_FS_SEQ_..." because there are entries that depend on it.
+    random_suffix = secrets.token_hex(4)
+    seq_name = f"_fs_seq_{table_name}_{col_name}_{random_suffix}"
+
+    # Build CREATE SEQUENCE, using START/INCREMENT if provided
+    start_val = "1"
+    increment_val = "1"
+
+    identity = auto.find(exp.GeneratedAsIdentityColumnConstraint)
+    if identity:
+        s = identity.args.get("start")
+        i = identity.args.get("increment")
+        if isinstance(s, exp.Literal):
+            start_val = s.this
+        if isinstance(i, exp.Literal):
+            increment_val = i.this
+
+    seq_stmt = sqlglot.parse_one(
+        f"CREATE SEQUENCE {seq_name} START WITH {start_val} INCREMENT BY {increment_val}",
+        read="duckdb",
+    )
+
+    # Build modified CREATE TABLE with DEFAULT NEXTVAL('<seq_name>') and without AUTOINCREMENT
+    new_create: exp.Create = expression.copy()
+    new_schema: exp.Schema = new_create.this
+
+    # Find the corresponding column in the copied schema
+    target_col = next(
+        cd for cd in (new_schema.expressions or []) if isinstance(cd, exp.ColumnDef) and cd.this.name == col_name
+    )
+
+    existing_constraints: list[exp.Expression] = target_col.args.get("constraints", []) or []
+
+    # Replace the AUTOINCREMENT/IDENTITY constraint in-place with DEFAULT NEXTVAL('<seq_name>')
+    for c in existing_constraints:
+        if isinstance(c, exp.ColumnConstraint) and isinstance(
+            c.args.get("kind"),
+            (
+                exp.AutoIncrementColumnConstraint,
+                exp.GeneratedAsIdentityColumnConstraint,
+            ),
+        ):
+            c.set(
+                "kind",
+                exp.DefaultColumnConstraint(
+                    this=exp.Anonymous(
+                        this="NEXTVAL",
+                        expressions=[exp.Literal(this=seq_name, is_string=True)],
+                    )
+                ),
+            )
+            break
+
+    target_col.set("constraints", existing_constraints)
+
+    return [seq_stmt, new_create]

tests/test_sequence.py

@@ -1,7 +1,41 @@
+import re
+
 import snowflake.connector.cursor
 import sqlglot
 
-from fakesnow.transforms import sequence_nextval
+from fakesnow.transforms import create_table_autoincrement, sequence_nextval
+
+
+def test_autoincrement(cur: snowflake.connector.cursor.SnowflakeCursor):
+    cur.execute("CREATE TABLE test_table (id NUMERIC NOT NULL AUTOINCREMENT, name VARCHAR)")
+    cur.execute("insert into test_table(name) values ('foo'), ('bar')")
+    cur.execute("select * from test_table")
+    assert cur.fetchall() == [(1, "foo"), (2, "bar")]
+
+    # recreate the table with a different sequence
+    cur.execute(
+        "CREATE or replace TABLE test_table(id NUMERIC NOT NULL IDENTITY start 10 increment 5 order, name VARCHAR)"
+    )
+    cur.execute("insert into test_table(name) values ('foo'), ('bar')")
+    cur.execute("select * from test_table")
+    assert cur.fetchall() == [(10, "foo"), (15, "bar")]
+
+
+def test_autoincrement_transform() -> None:
+    expr = sqlglot.parse_one(
+        "create table test_table (id numeric autoincrement start 10 increment 5 order)",
+        dialect="snowflake",
+    )
+    seq_stmt, table_stmt = create_table_autoincrement(expr)
+    seq_sql = seq_stmt.sql()
+    table_sql = table_stmt.sql()
+
+    m_seq = re.search(r"CREATE SEQUENCE (\S+) START WITH 10 INCREMENT BY 5", seq_sql)
+    assert m_seq, f"Unexpected SEQUENCE SQL: {seq_sql}"
+    seq_name = m_seq.group(1)
+    assert re.match(r"^_fs_seq_test_table_id_[0-9a-f]{8}$", seq_name)
+
+    assert table_sql == f"CREATE TABLE test_table (id DECIMAL(38, 0) DEFAULT NEXTVAL('{seq_name}'))"
 
 
 def test_sequence(dcur: snowflake.connector.cursor.SnowflakeCursor):