state-hub/migrations/versions/e8f9a0b1c2d3_task_wait_qualifiers.py

"""task wait qualifiers (CUST-WP-0074-T02)

Adds ``tasks.decision_id`` (record id of the decision a human-gated wait hangs
on) and ``workplan_dependencies.from_task_id`` so that task-block
``depends_on`` edges carry the waiting task as the from side. The two existing
partial unique indexes are narrowed to frontmatter edges (``from_task_id IS
NULL``) and two task-origin counterparts are added, so two tasks in one
workplan may wait on the same target.

Revision ID: e8f9a0b1c2d3
Revises: d7e8f9a0b1c2
"""
import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql

revision = "e8f9a0b1c2d3"
down_revision = "d7e8f9a0b1c2"
branch_labels = None
depends_on = None


def upgrade() -> None:
    op.add_column("tasks", sa.Column("decision_id", sa.String(length=120), nullable=True))

    op.add_column(
        "workplan_dependencies",
        sa.Column("from_task_id", postgresql.UUID(as_uuid=True), nullable=True),
    )
    op.create_foreign_key(
        "fk_workplan_dependencies_from_task_id",
        "workplan_dependencies",
        "tasks",
        ["from_task_id"],
        ["id"],
        ondelete="CASCADE",
        onupdate="CASCADE",
    )
    op.create_index(
        "ix_workplan_dependencies_from_task_id",
        "workplan_dependencies",
        ["from_task_id"],
    )

    op.drop_index("uq_wp_dep_workplan_target", table_name="workplan_dependencies")
    op.drop_index("uq_wp_dep_task_target", table_name="workplan_dependencies")
    op.create_index(
        "uq_wp_dep_workplan_target",
        "workplan_dependencies",
        ["from_workplan_id", "to_workplan_id", "relationship_type"],
        unique=True,
        postgresql_where=sa.text("to_workplan_id IS NOT NULL AND from_task_id IS NULL"),
    )
    op.create_index(
        "uq_wp_dep_task_target",
        "workplan_dependencies",
        ["from_workplan_id", "to_task_id", "relationship_type"],
        unique=True,
        postgresql_where=sa.text("to_task_id IS NOT NULL AND from_task_id IS NULL"),
    )
    op.create_index(
        "uq_wp_dep_from_task_workplan_target",
        "workplan_dependencies",
        ["from_task_id", "to_workplan_id", "relationship_type"],
        unique=True,
        postgresql_where=sa.text("to_workplan_id IS NOT NULL AND from_task_id IS NOT NULL"),
    )
    op.create_index(
        "uq_wp_dep_from_task_task_target",
        "workplan_dependencies",
        ["from_task_id", "to_task_id", "relationship_type"],
        unique=True,
        postgresql_where=sa.text("to_task_id IS NOT NULL AND from_task_id IS NOT NULL"),
    )


def downgrade() -> None:
    op.drop_index("uq_wp_dep_from_task_task_target", table_name="workplan_dependencies")
    op.drop_index("uq_wp_dep_from_task_workplan_target", table_name="workplan_dependencies")
    op.drop_index("uq_wp_dep_task_target", table_name="workplan_dependencies")
    op.drop_index("uq_wp_dep_workplan_target", table_name="workplan_dependencies")
    # Task-origin rows would collide under the frontmatter-only indexes.
    op.execute("DELETE FROM workplan_dependencies WHERE from_task_id IS NOT NULL")
    op.create_index(
        "uq_wp_dep_workplan_target",
        "workplan_dependencies",
        ["from_workplan_id", "to_workplan_id", "relationship_type"],
        unique=True,
        postgresql_where=sa.text("to_workplan_id IS NOT NULL"),
    )
    op.create_index(
        "uq_wp_dep_task_target",
        "workplan_dependencies",
        ["from_workplan_id", "to_task_id", "relationship_type"],
        unique=True,
        postgresql_where=sa.text("to_task_id IS NOT NULL"),
    )
    op.drop_index("ix_workplan_dependencies_from_task_id", table_name="workplan_dependencies")
    op.drop_constraint(
        "fk_workplan_dependencies_from_task_id", "workplan_dependencies", type_="foreignkey"
    )
    op.drop_column("workplan_dependencies", "from_task_id")
    op.drop_column("tasks", "decision_id")