"""partial indexes so GPU-job leasing stays O(batch), not O(completed)

The lease claims the lowest-id pending (or expired-leased) jobs. With only a
plain `status` index, `... ORDER BY id LIMIT n` walked the primary-key index from
the start, skipping the entire prefix of already-done/error rows before reaching
pending ones — so leasing slowed to a crawl as `done` piled up (the whole reason
throughput fell off a cliff mid-run and /status stalled). Two partial indexes fix
it: the pending one is id-ordered so the hot path reads just the first n entries,
and the leased-expiry one keeps the crash-recovery reclaim + the orphan sweep
cheap. They cover only the small live slice of the table, so they stay tiny even
as the done/error history grows to millions.

Revision ID: 0070
Revises: 0069
Create Date: 2026-06-30
"""
from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op

revision: str = "0070"
down_revision: Union[str, None] = "0069"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # Hot path: lowest-id pending jobs. Index on id, restricted to pending, so
    # `WHERE status='pending' ORDER BY id LIMIT n` is a short index-order scan.
    op.create_index(
        "ix_gpu_job_pending", "gpu_job", ["id"],
        postgresql_where=sa.text("status = 'pending'"),
    )
    # Crash-recovery: expired leases, for the lease backstop + recover_orphaned.
    op.create_index(
        "ix_gpu_job_leased_expires", "gpu_job", ["lease_expires_at"],
        postgresql_where=sa.text("status = 'leased'"),
    )


def downgrade() -> None:
    op.drop_index("ix_gpu_job_leased_expires", table_name="gpu_job")
    op.drop_index("ix_gpu_job_pending", table_name="gpu_job")
