"""add snapshots to run storage

Revision ID: c63a27054f08
Revises: 9fe9e746268c
Create Date: 2020-04-09 05:57:20.639458

"""
import sqlalchemy as sa
from alembic import op

from dagster.core.storage.migration.utils import has_column, has_table

# alembic magic breaks pylint
# pylint: disable=no-member

# revision identifiers, used by Alembic.
revision = 'c63a27054f08'
down_revision = '9fe9e746268c'
branch_labels = None
depends_on = None


def upgrade():
    if not has_table('snapshots'):
        op.create_table(
            'snapshots',
            sa.Column('id', sa.Integer, primary_key=True, autoincrement=True, nullable=False),
            sa.Column('snapshot_id', sa.String(255), unique=True, nullable=False),
            sa.Column('snapshot_body', sa.LargeBinary, nullable=False),
            sa.Column('snapshot_type', sa.String(63), nullable=False),
        )

    if not has_column('runs', 'snapshot_id'):
        # Sqlite does not support adding foreign keys to existing
        # tables, so we are forced to fallback on this witchcraft.
        # See https://alembic.sqlalchemy.org/en/latest/batch.html#dealing-with-referencing-foreign-keys
        # for additional context
        with op.batch_alter_table('runs') as batch_op:
            batch_op.execute('PRAGMA foreign_keys = OFF;')
            batch_op.add_column(
                sa.Column(
                    'snapshot_id',
                    sa.String(255),
                    sa.ForeignKey(
                        'snapshots.snapshot_id', name='fk_runs_snapshot_id_snapshots_snapshot_id'
                    ),
                ),
            )
        op.execute('PRAGMA foreign_keys = ON;')


def downgrade():
    if has_column('runs', 'snapshot_id'):
        with op.batch_alter_table('runs') as batch_op:
            batch_op.drop_column('snapshot_id')

    if has_table('snapshots'):
        op.drop_table('snapshots')
