nyaa/migrations/versions/2bceb2cb4d7c_add_comment_count_to_torrent.py

"""Add comment_count to Torrent

Revision ID: 2bceb2cb4d7c
Revises: d0eeb8049623
Create Date: 2017-05-26 15:07:21.114331

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '2bceb2cb4d7c'
down_revision = 'd0eeb8049623'
branch_labels = None
depends_on = None

COMMENT_UPDATE_SQL = '''UPDATE {0}_torrents
                        SET comment_count = (
                            SELECT COUNT(*) FROM {0}_comments
                            WHERE {0}_torrents.id = {0}_comments.torrent_id
                        );'''


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('nyaa_torrents', sa.Column('comment_count', sa.Integer(), nullable=False))
    op.create_index(op.f('ix_nyaa_torrents_comment_count'), 'nyaa_torrents', ['comment_count'], unique=False)

    op.add_column('sukebei_torrents', sa.Column('comment_count', sa.Integer(), nullable=False))
    op.create_index(op.f('ix_sukebei_torrents_comment_count'), 'sukebei_torrents', ['comment_count'], unique=False)
    # ### end Alembic commands ###

    connection = op.get_bind()

    print('Updating comment counts on nyaa_torrents...')
    connection.execute(sa.sql.text(COMMENT_UPDATE_SQL.format('nyaa')))
    print('Done.')

    print('Updating comment counts on sukebei_torrents...')
    connection.execute(sa.sql.text(COMMENT_UPDATE_SQL.format('sukebei')))
    print('Done.')


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f('ix_nyaa_torrents_comment_count'), table_name='nyaa_torrents')
    op.drop_column('nyaa_torrents', 'comment_count')

    op.drop_index(op.f('ix_sukebei_torrents_comment_count'), table_name='sukebei_torrents')
    op.drop_column('sukebei_torrents', 'comment_count')
    # ### end Alembic commands ###