"""Add loyalty_accepted to payments

Revision ID: 76b9df869bbe
Revises: 364536eb33d5
Create Date: 2016-12-03 11:58:04.152000

"""
import pickle

from alembic import op
from sqlalchemy import (
    String,
    Integer,
    PickleType,
    Boolean,
    DateTime,
)
from sqlalchemy.sql import column, table


# revision identifiers, used by Alembic.
revision = '76b9df869bbe'
down_revision = '364536eb33d5'
branch_labels = None
depends_on = None

first_agent_version = '17.3.17'


def upgrade():
    """
    Use batch-operations to support sqlite

    > with op.batch_alter_table('table_name') as batch_op:
    >     pass
    """
    conn = op.get_bind()
    tickets_table = table(
        'tickets',
        column('id', String(32)),
        column('payments', PickleType),
    )
    held_tickets_table = table(
        'held_tickets',
        column('id', String(32)),
        column('payments', PickleType),
    )

    for my_table in (tickets_table, held_tickets_table):
        cursor = conn.execute("SELECT id, payments FROM {}".format(my_table.name))
        rows = cursor.fetchall()

        tickets_payments = [
            (row[0], pickle.loads(row[1]))
            for row in rows
        ]
        for ticket_id, payments in tickets_payments:
            for payment in payments:
                payment.loyalty_accepted = None

        for ticket_id, payments in tickets_payments:
            op.execute(
                my_table
                .update()
                .where(my_table.c.id == ticket_id)
                .values({'payments': payments})
            )


def downgrade():
    """
    Use batch-operations to support sqlite

    > with op.batch_alter_table('table_name') as batch_op:
    >     pass
    """

    pass
