ALTER TABLE generation_tasks
    ADD COLUMN IF NOT EXISTS lease_token VARCHAR(64),
    ADD COLUMN IF NOT EXISTS lease_owner VARCHAR(128),
    ADD COLUMN IF NOT EXISTS lease_expires_at TIMESTAMPTZ,
    ADD COLUMN IF NOT EXISTS attempt_count INT NOT NULL DEFAULT 0,
    ADD COLUMN IF NOT EXISTS last_heartbeat_at TIMESTAMPTZ;

CREATE INDEX IF NOT EXISTS idx_generation_tasks_running_lease
    ON generation_tasks(lease_expires_at, id)
    WHERE status = 'running';

CREATE INDEX IF NOT EXISTS idx_generation_tasks_queued_stale
    ON generation_tasks(updated_at, id)
    WHERE status = 'queued';