"""Add indices for coalesced `start time` sorts

Revision ID: 4f90ad6349bd
Revises: a0284438370e
Create Date: 2022-11-10 16:59:21.884051

"""
import sqlalchemy as sa
from alembic import op

# revision identifiers, used by Alembic.
revision = "4f90ad6349bd"
down_revision = "a0284438370e"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("flow_run", schema=None) as batch_op:
        batch_op.drop_index(
            "ix_flow_run__scheduler_deployment_id_auto_scheduled_next_scheduled_start_time"
        )
        batch_op.create_index(
            "ix_flow_run__coalesce_start_time_expected_start_time_asc",
            [sa.text("coalesce('start_time', 'expected_start_time') ASC")],
            unique=False,
        )
        batch_op.create_index(
            "ix_flow_run__coalesce_start_time_expected_start_time_desc",
            [sa.text("coalesce('start_time', 'expected_start_time') DESC")],
            unique=False,
        )

    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("flow_run", schema=None) as batch_op:
        batch_op.execute(
            "DROP INDEX IF EXISTS"
            " ix_flow_run__coalesce_start_time_expected_start_time_desc"
        )
        batch_op.execute(
            "DROP INDEX IF EXISTS"
            " ix_flow_run__coalesce_start_time_expected_start_time_asc"
        )
        batch_op.create_index(
            "ix_flow_run__scheduler_deployment_id_auto_scheduled_next_scheduled_start_time",
            ["deployment_id", "auto_scheduled", "next_scheduled_start_time"],
            unique=False,
        )
    # ### end Alembic commands ###
