DataOps-platform/migrations/versions/20260723_190_rule_cleanup_leases.py

"""Add expiring claims for governed-rule cleanup workers."""

from alembic import op

revision = "20260723_190"
down_revision = "20260723_180"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # Revision-180 cleanup workers must be stopped during this schema upgrade.
    # Their owner-only claims have no lease timestamp, so making those claims
    # immediately expired is the only deterministic, fail-safe takeover path.
    op.execute(
        """
        ALTER TABLE public.rule_violation_samples
            ADD COLUMN cleanup_claim_expires_at TIMESTAMPTZ;

        UPDATE public.rule_violation_samples
        SET cleanup_claim_expires_at = CURRENT_TIMESTAMP
        WHERE cleanup_claim IS NOT NULL
          AND cleanup_claim_expires_at IS NULL;

        CREATE INDEX idx_rule_violation_sample_cleanup_lease
            ON public.rule_violation_samples(
                cleanup_claim_expires_at, expires_at, id
            );

        ALTER TABLE public.rule_sql_staging_receipts
            ADD COLUMN cleanup_claim_expires_at TIMESTAMPTZ;

        UPDATE public.rule_sql_staging_receipts
        SET cleanup_claim_expires_at = CURRENT_TIMESTAMP
        WHERE cleanup_claim IS NOT NULL
          AND cleanup_claim_expires_at IS NULL;

        CREATE INDEX idx_rule_sql_receipt_cleanup_lease
            ON public.rule_sql_staging_receipts(
                cleanup_claim_expires_at, expires_at, id
            );
        """
    )


def downgrade() -> None:
    raise RuntimeError("rule cleanup leases are forward-only")