Russell Ballestrini
2639a2f99c
modified: app.py new file: migrations/README new file: migrations/alembic.ini new file: migrations/env.py new file: migrations/script.py.mako new file: migrations/versions/a9e886c56482_create_room_table.py
113 lines
4 KiB
Python
113 lines
4 KiB
Python
from alembic import op
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.sql import table, column, select
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision = 'a9e886c56482'
|
|
down_revision = None
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
def upgrade():
|
|
# Create room table
|
|
op.create_table('room',
|
|
sa.Column('id', sa.Integer(), nullable=False),
|
|
sa.Column('name', sa.String(length=128), nullable=False),
|
|
sa.Column('title', sa.String(length=128), nullable=True),
|
|
sa.PrimaryKeyConstraint('id'),
|
|
sa.UniqueConstraint('name')
|
|
)
|
|
|
|
# Add room_id column to message table
|
|
op.add_column('message', sa.Column('room_id', sa.Integer(), nullable=True))
|
|
|
|
# Temporary table objects
|
|
message_table = table('message',
|
|
column('id', sa.Integer),
|
|
column('username', sa.String),
|
|
column('content', sa.String),
|
|
column('room', sa.String),
|
|
column('room_id', sa.Integer),
|
|
)
|
|
room_table = table('room',
|
|
column('id', sa.Integer),
|
|
column('name', sa.String)
|
|
)
|
|
|
|
# Execution context
|
|
conn = op.get_bind()
|
|
|
|
# Insert distinct rooms into room table and create mapping
|
|
distinct_rooms = conn.execute(select(message_table.c.room).distinct())
|
|
room_name_to_id = {}
|
|
for room_name, in distinct_rooms:
|
|
conn.execute(room_table.insert().values(name=room_name))
|
|
room_id = conn.execute(select(room_table.c.id).where(room_table.c.name == room_name)).scalar()
|
|
room_name_to_id[room_name] = room_id
|
|
|
|
# Update message table with room_id
|
|
for room_name, room_id in room_name_to_id.items():
|
|
conn.execute(message_table.update().where(message_table.c.room == room_name).values(room_id=room_id))
|
|
|
|
# Create new_message table
|
|
new_message_table = op.create_table('new_message',
|
|
sa.Column('id', sa.Integer(), nullable=False),
|
|
sa.Column('username', sa.String(length=128), nullable=False),
|
|
sa.Column('content', sa.String(length=1024), nullable=False),
|
|
sa.Column('room_id', sa.Integer(), nullable=False),
|
|
sa.ForeignKeyConstraint(['room_id'], ['room.id']),
|
|
sa.PrimaryKeyConstraint('id')
|
|
)
|
|
|
|
# Copy data from old message table to new_message table
|
|
old_messages = conn.execute(sa.select(message_table)).fetchall()
|
|
for old_message in old_messages:
|
|
conn.execute(new_message_table.insert().values(
|
|
id=old_message.id,
|
|
username=old_message.username,
|
|
content=old_message.content,
|
|
room_id=old_message.room_id
|
|
))
|
|
|
|
# Drop old message table and rename new_message to message
|
|
op.drop_table('message')
|
|
op.rename_table('new_message', 'message')
|
|
|
|
def downgrade():
|
|
# Recreate old_message table with 'room' column
|
|
old_message_table = op.create_table('old_message',
|
|
sa.Column('id', sa.Integer(), nullable=False),
|
|
sa.Column('username', sa.String(length=128), nullable=False),
|
|
sa.Column('content', sa.String(length=1024), nullable=False),
|
|
sa.Column('room', sa.String(length=128), nullable=False),
|
|
sa.PrimaryKeyConstraint('id')
|
|
)
|
|
|
|
# Copy data back from message to old_message
|
|
message_table = table('message',
|
|
column('id', sa.Integer),
|
|
column('username', sa.String),
|
|
column('content', sa.String),
|
|
column('room_id', sa.Integer)
|
|
)
|
|
room_table = table('room',
|
|
column('id', sa.Integer),
|
|
column('name', sa.String)
|
|
)
|
|
|
|
conn = op.get_bind()
|
|
messages = conn.execute(select(message_table)).fetchall()
|
|
for message in messages:
|
|
room_name = conn.execute(select(room_table.c.name).where(room_table.c.id == message.room_id)).scalar()
|
|
conn.execute(old_message_table.insert().values(
|
|
id=message.id,
|
|
username=message.username,
|
|
content=message.content,
|
|
room=room_name
|
|
))
|
|
|
|
# Drop current message table and rename old_message to message
|
|
op.drop_table('message')
|
|
op.rename_table('old_message', 'message')
|
|
op.drop_table('room')
|
|
|