PostgreSQL durable queue claiming

create schema if not exists app;

create type app.job_state as enum (
  'pending',
  'running',
  'retry_wait',
  'succeeded',
  'dead'
);

create table app.jobs (
  id uuid primary key default gen_random_uuid(),
  tenant_id uuid not null,
  operation_key text not null,
  kind text not null,
  payload jsonb not null check (jsonb_typeof(payload) = 'object'),
  state app.job_state not null default 'pending',
  attempts integer not null default 0 check (attempts >= 0),
  max_attempts integer not null default 6 check (max_attempts between 1 and 20),
  available_at timestamptz not null default clock_timestamp(),
  lease_owner text,
  lease_token uuid,
  lease_expires_at timestamptz,
  last_error jsonb,
  result jsonb,
  created_at timestamptz not null default clock_timestamp(),
  updated_at timestamptz not null default clock_timestamp(),
  completed_at timestamptz,
  constraint jobs_operation_key_unique unique (tenant_id, operation_key),
  constraint jobs_lease_shape check (
    (state = 'running' and lease_owner is not null and lease_token is not null and lease_expires_at is not null)
    or
    (state <> 'running' and lease_owner is null and lease_token is null and lease_expires_at is null)
  )
);

create index jobs_claimable_idx
on app.jobs (tenant_id, available_at, created_at, id)
where state in ('pending', 'retry_wait');

create index jobs_expired_lease_idx
on app.jobs (tenant_id, lease_expires_at, id)
where state = 'running';

create table app.job_events (
  sequence bigint generated always as identity primary key,
  id uuid not null default gen_random_uuid(),
  tenant_id uuid not null,
  job_id uuid not null references app.jobs(id) on delete cascade,
  event_type text not null check (event_type in (
    'created', 'claimed', 'retry_scheduled',
    'succeeded', 'dead', 'lease_expired'
  )),
  actor text not null,
  details jsonb not null default '{}'::jsonb check (jsonb_typeof(details) = 'object'),
  created_at timestamptz not null default clock_timestamp(),
  constraint job_events_id_unique unique (tenant_id, id)
);

create index job_events_job_timeline_idx
on app.job_events (tenant_id, job_id, sequence);

alter table app.jobs enable row level security;
alter table app.jobs force row level security;
alter table app.job_events enable row level security;
alter table app.job_events force row level security;

create or replace function app.current_tenant_id()
returns uuid
language sql
stable
set search_path = ''
as $$
  select nullif(current_setting('app.tenant_id', true), '')::uuid
$$;

create policy jobs_tenant_select on app.jobs
for select
using (tenant_id = app.current_tenant_id());

create policy job_events_tenant_select on app.job_events
for select
using (tenant_id = app.current_tenant_id());

revoke all on schema app from public;
revoke all on all tables in schema app from public;
revoke all on all functions in schema app from public;

create or replace function app.claim_jobs(
  p_tenant_id uuid,
  p_owner text,
  p_limit integer default 10,
  p_lease interval default interval '30 seconds'
)
returns setof app.jobs
language plpgsql
security definer
set search_path = ''
as $$
declare
  v_now timestamptz := clock_timestamp();
  v_job app.jobs;
begin
  if p_tenant_id is null or nullif(btrim(p_owner), '') is null then
    raise exception using errcode = '22023', message = 'Tenant/owner required.';
  end if;
  if p_limit not between 1 and 100 or p_lease <= interval '0 seconds' then
    raise exception using errcode = '22023', message = 'Invalid claim bounds.';
  end if;

  for v_job in
    with candidates as materialized (
      select id
      from app.jobs
      where tenant_id = p_tenant_id
        and state in ('pending', 'retry_wait')
        and available_at <= v_now
      order by available_at, created_at, id
      for update skip locked
      limit p_limit
    ), claimed as (
      update app.jobs as jobs
      set state = 'running',
          attempts = jobs.attempts + 1,
          lease_owner = p_owner,
          lease_token = gen_random_uuid(),
          lease_expires_at = v_now + p_lease,
          updated_at = v_now
      from candidates
      where jobs.id = candidates.id
        and jobs.tenant_id = p_tenant_id
      returning jobs.*
    )
    select * from claimed
  loop
    insert into app.job_events (
      tenant_id, job_id, event_type, actor, details, created_at
    ) values (
      v_job.tenant_id,
      v_job.id,
      'claimed',
      p_owner,
      jsonb_build_object(
        'attempt', v_job.attempts,
        'leaseToken', v_job.lease_token,
        'leaseExpiresAt', v_job.lease_expires_at
      ),
      v_now
    );
    return next v_job;
  end loop;
end;
$$;

create or replace function app.renew_job_lease(
  p_tenant_id uuid,
  p_job_id uuid,
  p_owner text,
  p_token uuid,
  p_lease interval default interval '30 seconds'
)
returns boolean
language plpgsql
security definer
set search_path = ''
as $$
declare
  v_now timestamptz := clock_timestamp();
  v_renewed boolean := false;
begin
  update app.jobs
  set lease_expires_at = v_now + p_lease,
      updated_at = v_now
  where tenant_id = p_tenant_id
    and id = p_job_id
    and state = 'running'
    and lease_owner = p_owner
    and lease_token = p_token
    and lease_expires_at > v_now
  returning true into v_renewed;

  return coalesce(v_renewed, false);
end;
$$;

create or replace function app.complete_job(
  p_tenant_id uuid,
  p_job_id uuid,
  p_owner text,
  p_token uuid,
  p_result jsonb
)
returns boolean
language plpgsql
security definer
set search_path = ''
as $$
declare
  v_now timestamptz := clock_timestamp();
  v_completed boolean := false;
begin
  update app.jobs
  set state = 'succeeded',
      result = p_result,
      lease_owner = null,
      lease_token = null,
      lease_expires_at = null,
      completed_at = v_now,
      updated_at = v_now
  where tenant_id = p_tenant_id
    and id = p_job_id
    and state = 'running'
    and lease_owner = p_owner
    and lease_token = p_token
    and lease_expires_at > v_now
  returning true into v_completed;

  if coalesce(v_completed, false) then
    insert into app.job_events (
      tenant_id, job_id, event_type, actor, details, created_at
    ) values (
      p_tenant_id, p_job_id, 'succeeded', p_owner,
      jsonb_build_object('leaseToken', p_token), v_now
    );
  end if;
  return coalesce(v_completed, false);
end;
$$;

create or replace function app.fail_job(
  p_tenant_id uuid,
  p_job_id uuid,
  p_owner text,
  p_token uuid,
  p_error jsonb,
  p_retry_delay interval default interval '5 seconds'
)
returns app.job_state
language plpgsql
security definer
set search_path = ''
as $$
declare
  v_now timestamptz := clock_timestamp();
  v_state app.job_state;
begin
  update app.jobs
  set state = case when attempts >= max_attempts then 'dead'::app.job_state else 'retry_wait'::app.job_state end,
      available_at = case when attempts >= max_attempts then available_at else v_now + p_retry_delay end,
      last_error = p_error,
      lease_owner = null,
      lease_token = null,
      lease_expires_at = null,
      completed_at = case when attempts >= max_attempts then v_now else null end,
      updated_at = v_now
  where tenant_id = p_tenant_id
    and id = p_job_id
    and state = 'running'
    and lease_owner = p_owner
    and lease_token = p_token
    and lease_expires_at > v_now
  returning state into v_state;

  if v_state is not null then
    insert into app.job_events (
      tenant_id, job_id, event_type, actor, details, created_at
    ) values (
      p_tenant_id,
      p_job_id,
      case when v_state = 'dead' then 'dead' else 'retry_scheduled' end,
      p_owner,
      jsonb_build_object('leaseToken', p_token, 'error', p_error),
      v_now
    );
  end if;
  return v_state;
end;
$$;

create or replace function app.release_expired_jobs(
  p_tenant_id uuid,
  p_limit integer default 100
)
returns integer
language plpgsql
security definer
set search_path = ''
as $$
declare
  v_now timestamptz := clock_timestamp();
  v_job record;
  v_count integer := 0;
begin
  for v_job in
    with expired as materialized (
      select id
      from app.jobs
      where tenant_id = p_tenant_id
        and state = 'running'
        and lease_expires_at <= v_now
      order by lease_expires_at, id
      for update skip locked
      limit p_limit
    )
    update app.jobs as jobs
    set state = case when jobs.attempts >= jobs.max_attempts then 'dead'::app.job_state else 'retry_wait'::app.job_state end,
        available_at = case when jobs.attempts >= jobs.max_attempts then jobs.available_at else v_now end,
        last_error = jsonb_build_object('code', 'LEASE_EXPIRED', 'retryable', jobs.attempts < jobs.max_attempts),
        lease_owner = null,
        lease_token = null,
        lease_expires_at = null,
        completed_at = case when jobs.attempts >= jobs.max_attempts then v_now else null end,
        updated_at = v_now
    from expired
    where jobs.id = expired.id
      and jobs.tenant_id = p_tenant_id
    returning jobs.id, jobs.state
  loop
    v_count := v_count + 1;
    insert into app.job_events (
      tenant_id, job_id, event_type, actor, details, created_at
    ) values (
      p_tenant_id,
      v_job.id,
      case when v_job.state = 'dead' then 'dead' else 'lease_expired' end,
      'watchdog',
      jsonb_build_object('recoveredState', v_job.state),
      v_now
    );
  end loop;
  return v_count;
end;
$$;

grant usage on schema app to authenticated;
grant select on app.jobs, app.job_events to authenticated;
grant execute on function app.claim_jobs(uuid, text, integer, interval) to service_role;
grant execute on function app.renew_job_lease(uuid, uuid, text, uuid, interval) to service_role;
grant execute on function app.complete_job(uuid, uuid, text, uuid, jsonb) to service_role;
grant execute on function app.fail_job(uuid, uuid, text, uuid, jsonb, interval) to service_role;
grant execute on function app.release_expired_jobs(uuid, integer) to service_role;