archimedes_migrations 0.3.0

Migrations package for archimedes, a high performance Rust/PostgreSQL job queue
Documentation
pub const M000005_MIGRATION: &[&str] = &[
    r#"
        alter table :ARCHIMEDES_SCHEMA.jobs add column revision int default 0 not null;
    "#,
    r#"
        alter table :ARCHIMEDES_SCHEMA.jobs add column flags jsonb default null;
    "#,
    r#"
        drop function :ARCHIMEDES_SCHEMA.add_job(text, json, text, timestamptz, int, text, int);
    "#,
    r#"
        create function :ARCHIMEDES_SCHEMA.add_job(
            identifier text,
            payload json = null,
            queue_name text = null,
            run_at timestamptz = null,
            max_attempts int = null,
            job_key text = null,
            priority int = null,
            flags text[] = null
        ) returns :ARCHIMEDES_SCHEMA.jobs as $$
        declare
            v_job :ARCHIMEDES_SCHEMA.jobs;
        begin
            -- Apply rationality checks
            if length(identifier) > 128 then
                raise exception 'Task identifier is too long (max length: 128).' using errcode = 'GWBID';
            end if;
            if queue_name is not null and length(queue_name) > 128 then
                raise exception 'Job queue name is too long (max length: 128).' using errcode = 'GWBQN';
            end if;
            if job_key is not null and length(job_key) > 512 then
                raise exception 'Job key is too long (max length: 512).' using errcode = 'GWBJK';
            end if;
            if max_attempts < 1 then
                raise exception 'Job maximum attempts must be at least 1' using errcode = 'GWBMA';
            end if;

            if job_key is not null then
                -- Upsert job
                insert into :ARCHIMEDES_SCHEMA.jobs (
                    task_identifier,
                    payload,
                    queue_name,
                    run_at,
                    max_attempts,
                    key,
                    priority,
                    flags
                )
                    values(
                        identifier,
                        coalesce(payload, '{}'::json),
                        queue_name,
                        coalesce(run_at, now()),
                        coalesce(max_attempts, 25),
                        job_key,
                        coalesce(priority, 0),
                        (
                          select jsonb_object_agg(flag, true)
                          from unnest(flags) as item(flag)
                        )
                    )
                    on conflict (key) do update set
                        task_identifier=excluded.task_identifier,
                        payload=excluded.payload,
                        queue_name=excluded.queue_name,
                        max_attempts=excluded.max_attempts,
                        run_at=excluded.run_at,
                        priority=excluded.priority,
                        revision=jobs.revision + 1,
                        flags=excluded.flags,

                        -- always reset error/retry state
                        attempts=0,
                        last_error=null
                    where jobs.locked_at is null
                    returning *
                    into v_job;

                -- If upsert succeeded (insert or update), return early
                if not (v_job is null) then
                    return v_job;
                end if;

                -- Upsert failed -> there must be an existing job that is locked. Remove
                -- existing key to allow a new one to be inserted, and prevent any
                -- subsequent retries by bumping attempts to the max allowed.
                update :ARCHIMEDES_SCHEMA.jobs
                    set
                        key = null,
                        attempts = jobs.max_attempts
                    where key = job_key;
            end if;

            -- insert the new job. Assume no conflicts due to the update above
            insert into :ARCHIMEDES_SCHEMA.jobs(
                task_identifier,
                payload,
                queue_name,
                run_at,
                max_attempts,
                key,
                priority,
                flags
            )
            values(
                identifier,
                coalesce(payload, '{}'::json),
                queue_name,
                coalesce(run_at, now()),
                coalesce(max_attempts, 25),
                job_key,
                coalesce(priority, 0),
                (
                    select jsonb_object_agg(flag, true)
                    from unnest(flags) as item(flag)
                )
            )
            returning *
            into v_job;

            return v_job;
        end;
        $$ language plpgsql volatile;
    "#,
    r#"
        drop function :ARCHIMEDES_SCHEMA.get_job(text, text[], interval);
    "#,
    r#"
        create function :ARCHIMEDES_SCHEMA.get_job(
            worker_id text,
            task_identifiers text[] = null,
            job_expiry interval = interval '4 hours',
            forbidden_flags text[] = null
        ) returns :ARCHIMEDES_SCHEMA.jobs as $$
        declare
            v_job_id bigint;
            v_queue_name text;
            v_row :ARCHIMEDES_SCHEMA.jobs;
            v_now timestamptz = now();
        begin
            if worker_id is null or length(worker_id) < 10 then
                raise exception 'invalid worker id';
            end if;

            select jobs.queue_name, jobs.id into v_queue_name, v_job_id
                from :ARCHIMEDES_SCHEMA.jobs
                where (jobs.locked_at is null or jobs.locked_at < (v_now - job_expiry))
                and (
                    jobs.queue_name is null
                    or
                    exists (
                        select 1
                        from :ARCHIMEDES_SCHEMA.job_queues
                        where job_queues.queue_name = jobs.queue_name
                        and (job_queues.locked_at is null or job_queues.locked_at < (v_now - job_expiry))
                        for update
                        skip locked
                    )
                )
                and run_at <= v_now
                and attempts < max_attempts
                and (task_identifiers is null or task_identifier = any(task_identifiers))
                and (forbidden_flags is null or (flags ?| forbidden_flags) is not true)
                order by priority asc, run_at asc, id asc
                limit 1
                for update
                skip locked;

            if v_job_id is null then
                return null;
            end if;

            if v_queue_name is not null then
                update :ARCHIMEDES_SCHEMA.job_queues
                set
                    locked_by = worker_id,
                    locked_at = v_now
                where job_queues.queue_name = v_queue_name;
            end if;

            update :ARCHIMEDES_SCHEMA.jobs
                set
                    attempts = attempts + 1,
                    locked_by = worker_id,
                    locked_at = v_now
                where id = v_job_id
                returning * into v_row;
    
            return v_row;
        end;
        $$ language plpgsql volatile;
    "#,
];