"""Add BlockSchema capabilities

Revision ID: 724e6dcc6b5d
Revises: dc7a3c6fd3e9
Create Date: 2022-05-26 13:57:43.152931

"""
import sqlalchemy as sa
from alembic import op

import prefect

# revision identifiers, used by Alembic.
revision = "724e6dcc6b5d"
down_revision = "dc7a3c6fd3e9"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "block_schema",
        sa.Column(
            "capabilities",
            prefect.server.utilities.database.JSON(astext_type=sa.Text()),
            server_default="[]",
            nullable=False,
        ),
    )

    connection = op.get_bind()
    meta_data = sa.MetaData(bind=connection)
    meta_data.reflect()
    BLOCK_SCHEMA = meta_data.tables["block_schema"]

    results = connection.execute(sa.select([BLOCK_SCHEMA.c.id, BLOCK_SCHEMA.c.type]))

    for id, type in results:
        if type == "STORAGE":
            connection.execute(
                sa.update(BLOCK_SCHEMA)
                .where(BLOCK_SCHEMA.c.id == id)
                .values(capabilities=["writeable", "readable", "storage"])
            )

    op.drop_index("ix_block_schema__type", table_name="block_schema")
    op.drop_column("block_schema", "type")
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "block_schema",
        sa.Column("type", sa.VARCHAR(), autoincrement=False, nullable=True),
    )
    op.create_index("ix_block_schema__type", "block_schema", ["type"], unique=False)

    connection = op.get_bind()
    meta_data = sa.MetaData(bind=connection)
    meta_data.reflect()
    BLOCK_SCHEMA = meta_data.tables["block_schema"]

    results = connection.execute(
        sa.select([BLOCK_SCHEMA.c.id, BLOCK_SCHEMA.c.capabilities])
    )

    for id, capabilities in results:
        if "storage" in capabilities:
            connection.execute(
                sa.update(BLOCK_SCHEMA)
                .where(BLOCK_SCHEMA.c.id == id)
                .values(type="STORAGE")
            )

    op.drop_column("block_schema", "capabilities")
    # ### end Alembic commands ###
