Approach (brief):- Define what “idempotent” means for this migration: applying it again does not change schema or data after first successful run.- Use a disposable PostgreSQL test database (TEST_DATABASE_URL) created per test.- Apply migration SQL as a transaction when possible. Verify idempotency by comparing schema+critical rows before/after a second apply.- Simulate partial failure by splitting migration into multiple statements, injecting a failing statement in the middle (Postgres: PERFORM raise_exception via DO block or use INVALID SQL). Verify that after a failed apply the DB is rolled back to pre-migration state.pytest test (uses psycopg2 and SQLAlchemy inspector for schema comparison):python
import os
import psycopg2
import sqlalchemy as sa
import pytest
from sqlalchemy.engine import reflection
TEST_DATABASE_URL = os.getenv("TEST_DATABASE_URL") # e.g. postgres://user:pass@localhost:5432/db
def apply_migration(conn, sql):
"""Apply raw SQL migration within a transaction."""
with conn.cursor() as cur:
cur.execute("BEGIN;")
try:
cur.execute(sql)
cur.execute("COMMIT;")
except Exception:
cur.execute("ROLLBACK;")
raise
def get_schema_snapshot(engine):
"""Return a normalized schema snapshot: tables -> sorted columns and types."""
insp = reflection.Inspector.from_engine(engine)
schema = {}
for table in sorted(insp.get_table_names()):
cols = insp.get_columns(table)
schema[table] = tuple((c['name'], str(c['type'])) for c in cols)
return schema
def get_table_rows(conn, table):
with conn.cursor() as cur:
cur.execute(f"SELECT * FROM {table} ORDER BY 1;")
return cur.fetchall()
@pytest.fixture
def db_engine():
engine = sa.create_engine(TEST_DATABASE_URL)
yield engine
engine.dispose()
@pytest.fixture
def db_conn():
conn = psycopg2.connect(TEST_DATABASE_URL)
conn.autocommit = False
yield conn
conn.close()
def test_migration_idempotent_and_rollback(db_engine, db_conn):
# Initial setup: create a baseline table and seed a row
with db_conn.cursor() as cur:
cur.execute("DROP TABLE IF EXISTS users;")
cur.execute("CREATE TABLE users (id serial primary key, name text);")
cur.execute("INSERT INTO users (name) VALUES ('alice');")
db_conn.commit()
before_schema = get_schema_snapshot(db_engine)
before_rows = get_table_rows(db_conn, "users")
# Example migration: add column if not exists, populate it
migration_sql = """
ALTER TABLE users ADD COLUMN IF NOT EXISTS email text;
UPDATE users SET email = name || '@example.com' WHERE email IS NULL;
"""
# Apply migration twice to confirm idempotency
apply_migration(db_conn, migration_sql)
after_first_schema = get_schema_snapshot(db_engine)
after_first_rows = get_table_rows(db_conn, "users")
apply_migration(db_conn, migration_sql)
after_second_schema = get_schema_snapshot(db_engine)
after_second_rows = get_table_rows(db_conn, "users")
assert after_first_schema == after_second_schema, "Schema changed on second apply"
assert after_first_rows == after_second_rows, "Data changed on second apply"
# Simulate partial failure: create a migration with a failing statement in the middle.
# Use an invalid SQL command to force PostgreSQL error. The whole migration should rollback.
failing_migration = """
ALTER TABLE users ADD COLUMN IF NOT EXISTS bio text;
-- force a runtime error (invalid SQL)
THIS_WILL_FAIL;
ALTER TABLE users ADD COLUMN IF NOT EXISTS created_at timestamptz DEFAULT now();
"""
# Capture state before failing migration
state_schema_before_fail = get_schema_snapshot(db_engine)
state_rows_before_fail = get_table_rows(db_conn, "users")
with pytest.raises(Exception):
apply_migration(db_conn, failing_migration)
# Verify state is unchanged (rolled back)
state_schema_after_fail = get_schema_snapshot(db_engine)
state_rows_after_fail = get_table_rows(db_conn, "users")
assert state_schema_before_fail == state_schema_after_fail, "Schema changed despite failed migration"
assert state_rows_before_fail == state_rows_after_fail, "Data changed despite failed migration"
