Alembic data migration

migrate_complex_uuid_data_to_user_table.py

"""[4] Migrate complex_uuid data to user table

Revision ID: dc33b414403d
Revises: 336cf482107d
Create Date: 2017-06-21 15:24:46.691582

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
from sqlalchemy.orm import sessionmaker

# revision identifiers, used by Alembic.
revision = 'dc33b414403d'
down_revision = '336cf482107d'
branch_labels = None
depends_on = None

Session = sessionmaker()


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('user_table', sa.Column('complex_uuid', postgresql.UUID(as_uuid=True), nullable=True))
    op.create_unique_constraint(None, 'user_table', ['complex_uuid'])
    # ### end Alembic commands ###

    copy_complex_uuid_to_user_table()

    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint('agent_complex_uuid_key', 'agent', type_='unique')
    op.drop_column('agent', 'complex_uuid')
    # ### end Alembic commands ###


def copy_complex_uuid_to_user_table():
    # create local table objects from the current state of the DB
    agent_table = sa.Table('agent', sa.MetaData(bind=op.get_bind()), autoload=True)
    user_table = sa.Table('user_table', sa.MetaData(bind=op.get_bind()), autoload=True)

    # copy complex_uuid data
    op.get_bind().execute(
        user_table.update().values(
            complex_uuid=agent_table.c.uuid
        ).where(
            user_table.c.id == agent_table.c.user_id
        )
    )


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(None, 'user_table', type_='unique')
    op.drop_column('user_table', 'complex_uuid')
    # ### end Alembic commands ###

    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('agent', sa.Column('complex_uuid', postgresql.UUID(), autoincrement=False, nullable=True))
    op.create_unique_constraint('agent_complex_uuid_key', 'agent', ['complex_uuid'])
    # ### end Alembic commands ###