"""Add 'merged' and 'closed' status values to tasks table.

Revision ID: 0008
Revises: 0007
"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op

revision: str = "0008"
down_revision: Union[str, None] = "0007"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

_NEW_CHECK = (
    "status IN ('grooming', 'todo', 'in_progress', 'done', 'cancelled',"
    " 'timeout', 'ready-to-deploy', 'merged', 'closed')"
)

_OLD_CHECK = (
    "status IN ('grooming', 'todo', 'in_progress', 'done', 'cancelled',"
    " 'timeout', 'ready-to-deploy')"
)


def upgrade() -> None:
    conn = op.get_bind()

    # Find all check constraints on the tasks table that cover the status column.
    rows = conn.execute(
        sa.text(
            "SELECT conname FROM pg_constraint"
            " WHERE conrelid = 'tasks'::regclass AND contype = 'c'"
            " AND pg_get_constraintdef(oid) ILIKE '%status%'"
        )
    ).fetchall()

    for (conname,) in rows:
        conn.execute(sa.text(f"ALTER TABLE tasks DROP CONSTRAINT IF EXISTS \"{conname}\""))

    conn.execute(
        sa.text(
            f"ALTER TABLE tasks ADD CONSTRAINT tasks_status_check CHECK ({_NEW_CHECK})"
        )
    )


def downgrade() -> None:
    conn = op.get_bind()

    rows = conn.execute(
        sa.text(
            "SELECT conname FROM pg_constraint"
            " WHERE conrelid = 'tasks'::regclass AND contype = 'c'"
            " AND pg_get_constraintdef(oid) ILIKE '%status%'"
        )
    ).fetchall()

    for (conname,) in rows:
        conn.execute(sa.text(f"ALTER TABLE tasks DROP CONSTRAINT IF EXISTS \"{conname}\""))

    conn.execute(
        sa.text(
            f"ALTER TABLE tasks ADD CONSTRAINT tasks_status_check CHECK ({_OLD_CHECK})"
        )
    )
