b5551cd26764_add_captain_column_to_teams.py

"""Add captain column to Teams

Revision ID: b5551cd26764
Revises: 4e4d5a9ea000
Create Date: 2019-04-12 00:29:08.021141

"""
from CTFd.models import db
from alembic import op
import sqlalchemy as sa
from sqlalchemy.sql import text, table, column, and_

# revision identifiers, used by Alembic.
revision = 'b5551cd26764'
down_revision = '4e4d5a9ea000'
branch_labels = None
depends_on = None

teams_table = table('teams',
    column('id', db.Integer),
    column('captain_id', db.Integer),
)

users_table = table('users',
    column('id', db.Integer),
    column('team_id', db.Integer),
)


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('teams', sa.Column('captain_id', sa.Integer(), nullable=True))
    op.create_foreign_key('team_captain_id', 'teams', 'users', ['captain_id'], ['id'])

    connection = op.get_bind()
    for team in connection.execute(teams_table.select()):
        users = connection.execute(
            users_table.select().where(users_table.c.team_id == team.id).order_by(users_table.c.id).limit(1)
        )
        for user in users:
            connection.execute(
                teams_table.update().where(
                    teams_table.c.id == team.id
                ).values(
                    captain_id=user.id
                )
            )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint('team_captain_id', 'teams', type_='foreignkey')
    op.drop_column('teams', 'captain_id')
    # ### end Alembic commands ###