"""Task queue improvements

Revision ID: 6b24c66979ab
Revises: 49cfda6eb87e
Create Date: 2024-04-04 12:07:28.055915

"""

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision = "6b24c66979ab"
down_revision = "49cfda6eb87e"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column("task_queue", sa.Column("sort_date", postgresql.TIMESTAMP(timezone=True), nullable=True))
    op.create_index(
        "ix_task_queue_sort",
        "task_queue",
        [sa.text("priority DESC"), sa.text("sort_date ASC"), sa.text("id ASC")],
        unique=False,
    )

    op.execute(
        sa.text(
            """
        UPDATE task_queue
        SET sort_date = br.created_on
        FROM base_record br
        WHERE task_queue.record_id = br.id
        """
        )
    )

    op.alter_column("task_queue", "sort_date", nullable=False)

    op.drop_index("ix_task_queue_required_programs", table_name="task_queue")
    op.create_index(
        "ix_task_queue_required_programs", "task_queue", [sa.text("required_programs")], postgresql_using="gin"
    )

    op.drop_index("ix_base_record_created_on", table_name="base_record")
    op.create_index("ix_base_record_created_on", "base_record", ["created_on"])
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index("ix_base_record_created_on", table_name="base_record")
    op.create_index("ix_base_record_created_on", "base_record", ["created_on"], postgresql_using="brin")

    op.drop_index("ix_task_queue_required_programs", table_name="task_queue")
    op.create_index("ix_task_queue_required_programs", "task_queue", [sa.text("required_programs")])

    op.drop_index("ix_task_queue_sort", table_name="task_queue")
    op.drop_column("task_queue", "sort_date")

    # ### end Alembic commands ###
