|
| 1 | +"""Edit replace id integers in all models to use UUID instead |
| 2 | +
|
| 3 | +Revision ID: d98dd8ec85a3 |
| 4 | +Revises: 9c0a54914c78 |
| 5 | +Create Date: 2024-07-19 04:08:04.000976 |
| 6 | +
|
| 7 | +""" |
| 8 | +from alembic import op |
| 9 | +import sqlalchemy as sa |
| 10 | +import sqlmodel.sql.sqltypes |
| 11 | +from sqlalchemy.dialects import postgresql |
| 12 | + |
| 13 | + |
| 14 | +# revision identifiers, used by Alembic. |
| 15 | +revision = 'd98dd8ec85a3' |
| 16 | +down_revision = '9c0a54914c78' |
| 17 | +branch_labels = None |
| 18 | +depends_on = None |
| 19 | + |
| 20 | + |
| 21 | +def upgrade(): |
| 22 | + # Ensure uuid-ossp extension is available |
| 23 | + op.execute('CREATE EXTENSION IF NOT EXISTS "uuid-ossp"') |
| 24 | + |
| 25 | + # Create a new UUID column with a default UUID value |
| 26 | + op.add_column('user', sa.Column('new_id', postgresql.UUID(as_uuid=True), default=sa.text('uuid_generate_v4()'))) |
| 27 | + op.add_column('item', sa.Column('new_id', postgresql.UUID(as_uuid=True), default=sa.text('uuid_generate_v4()'))) |
| 28 | + op.add_column('item', sa.Column('new_owner_id', postgresql.UUID(as_uuid=True), nullable=True)) |
| 29 | + |
| 30 | + # Populate the new columns with UUIDs |
| 31 | + op.execute('UPDATE "user" SET new_id = uuid_generate_v4()') |
| 32 | + op.execute('UPDATE item SET new_id = uuid_generate_v4()') |
| 33 | + op.execute('UPDATE item SET new_owner_id = (SELECT new_id FROM "user" WHERE "user".id = item.owner_id)') |
| 34 | + |
| 35 | + # Set the new_id as not nullable |
| 36 | + op.alter_column('user', 'new_id', nullable=False) |
| 37 | + op.alter_column('item', 'new_id', nullable=False) |
| 38 | + |
| 39 | + # Drop old columns and rename new columns |
| 40 | + op.drop_constraint('item_owner_id_fkey', 'item', type_='foreignkey') |
| 41 | + op.drop_column('item', 'owner_id') |
| 42 | + op.alter_column('item', 'new_owner_id', new_column_name='owner_id') |
| 43 | + |
| 44 | + op.drop_column('user', 'id') |
| 45 | + op.alter_column('user', 'new_id', new_column_name='id') |
| 46 | + |
| 47 | + op.drop_column('item', 'id') |
| 48 | + op.alter_column('item', 'new_id', new_column_name='id') |
| 49 | + |
| 50 | + # Create primary key constraint |
| 51 | + op.create_primary_key('user_pkey', 'user', ['id']) |
| 52 | + op.create_primary_key('item_pkey', 'item', ['id']) |
| 53 | + |
| 54 | + # Recreate foreign key constraint |
| 55 | + op.create_foreign_key('item_owner_id_fkey', 'item', 'user', ['owner_id'], ['id']) |
| 56 | + |
| 57 | +def downgrade(): |
| 58 | + # Reverse the upgrade process |
| 59 | + op.add_column('user', sa.Column('old_id', sa.Integer, autoincrement=True)) |
| 60 | + op.add_column('item', sa.Column('old_id', sa.Integer, autoincrement=True)) |
| 61 | + op.add_column('item', sa.Column('old_owner_id', sa.Integer, nullable=True)) |
| 62 | + |
| 63 | + # Populate the old columns with default values |
| 64 | + # Generate sequences for the integer IDs if not exist |
| 65 | + op.execute('CREATE SEQUENCE IF NOT EXISTS user_id_seq AS INTEGER OWNED BY "user".old_id') |
| 66 | + op.execute('CREATE SEQUENCE IF NOT EXISTS item_id_seq AS INTEGER OWNED BY item.old_id') |
| 67 | + |
| 68 | + op.execute('SELECT setval(\'user_id_seq\', COALESCE((SELECT MAX(old_id) + 1 FROM "user"), 1), false)') |
| 69 | + op.execute('SELECT setval(\'item_id_seq\', COALESCE((SELECT MAX(old_id) + 1 FROM item), 1), false)') |
| 70 | + |
| 71 | + op.execute('UPDATE "user" SET old_id = nextval(\'user_id_seq\')') |
| 72 | + op.execute('UPDATE item SET old_id = nextval(\'item_id_seq\'), old_owner_id = (SELECT old_id FROM "user" WHERE "user".id = item.owner_id)') |
| 73 | + |
| 74 | + # Drop new columns and rename old columns back |
| 75 | + op.drop_constraint('item_owner_id_fkey', 'item', type_='foreignkey') |
| 76 | + op.drop_column('item', 'owner_id') |
| 77 | + op.alter_column('item', 'old_owner_id', new_column_name='owner_id') |
| 78 | + |
| 79 | + op.drop_column('user', 'id') |
| 80 | + op.alter_column('user', 'old_id', new_column_name='id') |
| 81 | + |
| 82 | + op.drop_column('item', 'id') |
| 83 | + op.alter_column('item', 'old_id', new_column_name='id') |
| 84 | + |
| 85 | + # Create primary key constraint |
| 86 | + op.create_primary_key('user_pkey', 'user', ['id']) |
| 87 | + op.create_primary_key('item_pkey', 'item', ['id']) |
| 88 | + |
| 89 | + # Recreate foreign key constraint |
| 90 | + op.create_foreign_key('item_owner_id_fkey', 'item', 'user', ['owner_id'], ['id']) |
0 commit comments