@pgflow/core 0.0.5-prealpha.2

package/supabase/migrations/000050_complete_task.sql

@@ -0,0 +1,98 @@
+-- drop function if exists pgflow.complete_task(uuid, text, int, jsonb);
+create or replace function pgflow.complete_task(
+  run_id uuid,
+  step_slug text,
+  task_index int,
+  output jsonb
+)
+returns setof pgflow.step_tasks
+language plpgsql
+volatile
+set search_path to ''
+as $$
+begin
+
+WITH run_lock AS (
+  SELECT * FROM pgflow.runs
+  WHERE pgflow.runs.run_id = complete_task.run_id
+  FOR UPDATE
+),
+step_lock AS (
+  SELECT * FROM pgflow.step_states
+  WHERE pgflow.step_states.run_id = complete_task.run_id
+    AND pgflow.step_states.step_slug = complete_task.step_slug
+  FOR UPDATE
+),
+task AS (
+  UPDATE pgflow.step_tasks
+  SET
+    status = 'completed',
+    output = complete_task.output
+  WHERE pgflow.step_tasks.run_id = complete_task.run_id
+    AND pgflow.step_tasks.step_slug = complete_task.step_slug
+    AND pgflow.step_tasks.task_index = complete_task.task_index
+  RETURNING *
+),
+step_state AS (
+  UPDATE pgflow.step_states
+  SET
+    status = CASE
+    WHEN pgflow.step_states.remaining_tasks = 1 THEN 'completed'  -- Will be 0 after decrement
+    ELSE 'started'
+    END,
+    remaining_tasks = pgflow.step_states.remaining_tasks - 1
+  FROM task
+  WHERE pgflow.step_states.run_id = complete_task.run_id
+    AND pgflow.step_states.step_slug = complete_task.step_slug
+  RETURNING pgflow.step_states.*
+),
+-- Find all dependent steps if the current step was completed
+dependent_steps AS (
+  SELECT d.step_slug AS dependent_step_slug
+  FROM pgflow.deps d
+  JOIN step_state s ON s.status = 'completed' AND d.flow_slug = s.flow_slug
+  WHERE d.dep_slug = complete_task.step_slug
+  ORDER BY d.step_slug  -- Ensure consistent ordering
+),
+-- Lock dependent steps before updating
+dependent_steps_lock AS (
+  SELECT * FROM pgflow.step_states
+  WHERE pgflow.step_states.run_id = complete_task.run_id
+    AND pgflow.step_states.step_slug IN (SELECT dependent_step_slug FROM dependent_steps)
+  FOR UPDATE
+),
+-- Update all dependent steps
+dependent_steps_update AS (
+  UPDATE pgflow.step_states
+  SET remaining_deps = pgflow.step_states.remaining_deps - 1
+  FROM dependent_steps
+  WHERE pgflow.step_states.run_id = complete_task.run_id
+    AND pgflow.step_states.step_slug = dependent_steps.dependent_step_slug
+)
+-- Only decrement remaining_steps, don't update status
+UPDATE pgflow.runs
+SET remaining_steps = pgflow.runs.remaining_steps - 1
+FROM step_state
+WHERE pgflow.runs.run_id = complete_task.run_id
+  AND step_state.status = 'completed';
+
+PERFORM pgmq.archive(
+  queue_name => (SELECT run.flow_slug FROM pgflow.runs AS run WHERE run.run_id = complete_task.run_id),
+  msg_id => (SELECT message_id FROM pgflow.step_tasks AS step_task
+             WHERE step_task.run_id = complete_task.run_id
+             AND step_task.step_slug = complete_task.step_slug
+             AND step_task.task_index = complete_task.task_index)
+);
+
+PERFORM pgflow.start_ready_steps(complete_task.run_id);
+
+PERFORM pgflow.maybe_complete_run(complete_task.run_id);
+
+RETURN QUERY SELECT *
+FROM pgflow.step_tasks AS step_task
+WHERE step_task.run_id = complete_task.run_id
+  AND step_task.step_slug = complete_task.step_slug
+  AND step_task.task_index = complete_task.task_index;
+
+end;
+$$;

package/supabase/migrations/000055_calculate_retry_delay.sql

@@ -0,0 +1,11 @@
+create or replace function pgflow.calculate_retry_delay(
+  base_delay numeric,
+  attempts_count int
+)
+returns int
+language sql
+immutable
+parallel safe
+as $$
+  select floor(base_delay * power(2, attempts_count))::int
+$$;

package/supabase/migrations/000060_fail_task.sql

@@ -0,0 +1,124 @@
+create or replace function pgflow.fail_task(
+  run_id uuid,
+  step_slug text,
+  task_index int,
+  error_message text
+)
+returns setof pgflow.step_tasks
+language plpgsql
+volatile
+set search_path to ''
+as $$
+begin
+
+WITH run_lock AS (
+  SELECT * FROM pgflow.runs
+  WHERE pgflow.runs.run_id = fail_task.run_id
+  FOR UPDATE
+),
+step_lock AS (
+  SELECT * FROM pgflow.step_states
+  WHERE pgflow.step_states.run_id = fail_task.run_id
+    AND pgflow.step_states.step_slug = fail_task.step_slug
+  FOR UPDATE
+),
+flow_info AS (
+  SELECT r.flow_slug
+  FROM pgflow.runs r
+  WHERE r.run_id = fail_task.run_id
+),
+config AS (
+  SELECT
+    COALESCE(s.opt_max_attempts, f.opt_max_attempts) AS opt_max_attempts,
+    COALESCE(s.opt_base_delay, f.opt_base_delay) AS opt_base_delay
+  FROM pgflow.steps s
+  JOIN pgflow.flows f ON f.flow_slug = s.flow_slug
+  JOIN flow_info fi ON fi.flow_slug = s.flow_slug
+  WHERE s.flow_slug = fi.flow_slug AND s.step_slug = fail_task.step_slug
+),
+
+fail_or_retry_task as (
+  UPDATE pgflow.step_tasks as task
+  SET
+    status = CASE
+      WHEN task.attempts_count < (SELECT opt_max_attempts FROM config) THEN 'queued'
+      ELSE 'failed'
+    END,
+    error_message = fail_task.error_message
+  WHERE task.run_id = fail_task.run_id
+    AND task.step_slug = fail_task.step_slug
+    AND task.task_index = fail_task.task_index
+    AND task.status = 'queued'
+  RETURNING *
+),
+maybe_fail_step AS (
+  UPDATE pgflow.step_states
+  SET
+    status = CASE
+             WHEN (select fail_or_retry_task.status from fail_or_retry_task) = 'failed' THEN 'failed'
+             ELSE pgflow.step_states.status
+             END
+  FROM fail_or_retry_task
+  WHERE pgflow.step_states.run_id = fail_task.run_id
+    AND pgflow.step_states.step_slug = fail_task.step_slug
+  RETURNING pgflow.step_states.*
+)
+UPDATE pgflow.runs
+SET status = CASE
+              WHEN (select status from maybe_fail_step) = 'failed' THEN 'failed'
+              ELSE status
+              END
+WHERE pgflow.runs.run_id = fail_task.run_id;
+
+-- For queued tasks: delay the message for retry with exponential backoff
+PERFORM (
+  WITH retry_config AS (
+    SELECT
+      COALESCE(s.opt_base_delay, f.opt_base_delay) AS base_delay
+    FROM pgflow.steps s
+    JOIN pgflow.flows f ON f.flow_slug = s.flow_slug
+    JOIN pgflow.runs r ON r.flow_slug = f.flow_slug
+    WHERE r.run_id = fail_task.run_id
+      AND s.step_slug = fail_task.step_slug
+  ),
+  queued_tasks AS (
+    SELECT
+      r.flow_slug,
+      st.message_id,
+      pgflow.calculate_retry_delay((SELECT base_delay FROM retry_config), st.attempts_count) AS calculated_delay
+    FROM pgflow.step_tasks st
+    JOIN pgflow.runs r ON st.run_id = r.run_id
+    WHERE st.run_id = fail_task.run_id
+      AND st.step_slug = fail_task.step_slug
+      AND st.task_index = fail_task.task_index
+      AND st.status = 'queued'
+  )
+  SELECT pgmq.set_vt(qt.flow_slug, qt.message_id, qt.calculated_delay)
+  FROM queued_tasks qt
+  WHERE EXISTS (SELECT 1 FROM queued_tasks)
+);
+
+-- For failed tasks: archive the message
+PERFORM (
+  WITH failed_tasks AS (
+    SELECT r.flow_slug, st.message_id
+    FROM pgflow.step_tasks st
+    JOIN pgflow.runs r ON st.run_id = r.run_id
+    WHERE st.run_id = fail_task.run_id
+      AND st.step_slug = fail_task.step_slug
+      AND st.task_index = fail_task.task_index
+      AND st.status = 'failed'
+  )
+  SELECT pgmq.archive(ft.flow_slug, ft.message_id)
+  FROM failed_tasks ft
+  WHERE EXISTS (SELECT 1 FROM failed_tasks)
+);
+
+return query select *
+from pgflow.step_tasks st
+where st.run_id = fail_task.run_id
+  and st.step_slug = fail_task.step_slug
+  and st.task_index = fail_task.task_index;
+
+end;
+$$;

package/supabase/migrations/000_edge_worker_initial.sql

@@ -0,0 +1,86 @@
+create extension if not exists pgmq version '1.4.4';
+
+create schema if not exists edge_worker;
+
+-------------------------------------------------------------------------------
+-- Workers Table --------------------------------------------------------------
+-------------------------------------------------------------------------------
+create table if not exists edge_worker.workers (
+  worker_id uuid not null primary key,
+  queue_name text not null,
+  function_name text not null,
+  started_at timestamptz not null default now(),
+  stopped_at timestamptz,
+  last_heartbeat_at timestamptz not null default now()
+);
+
+--------------------------------------------------------------------------------
+-- Read With Poll --------------------------------------------------------------
+--                                                                            --
+-- This is a backport of the pgmq.read_with_poll function from version 1.5.0  --
+-- It is required because it fixes a bug with high CPU usage and Supabase     --
+-- is still using version 1.4.4.                                              --
+--                                                                            --
+-- It is slightly modified (removed headers which are not available in 1.4.1) --
+--                                                                            --
+-- This will be removed once Supabase upgrades to 1.5.0 or higher.            --
+--------------------------------------------------------------------------------
+create function edge_worker.read_with_poll(
+  queue_name text,
+  vt integer,
+  qty integer,
+  max_poll_seconds integer default 5,
+  poll_interval_ms integer default 100,
+  conditional jsonb default '{}'
+)
+returns setof pgmq.message_record as $$
+DECLARE
+    r pgmq.message_record;
+    stop_at timestamp;
+    sql text;
+    qtable text := pgmq.format_table_name(queue_name, 'q');
+BEGIN
+    stop_at := clock_timestamp() + make_interval(secs => max_poll_seconds);
+    LOOP
+      IF (SELECT clock_timestamp() >= stop_at) THEN
+        RETURN;
+      END IF;
+
+      sql := FORMAT(
+          $QUERY$
+          WITH cte AS
+          (
+              SELECT msg_id
+              FROM pgmq.%I
+              WHERE vt <= clock_timestamp() AND CASE
+                  WHEN %L != '{}'::jsonb THEN (message @> %2$L)::integer
+                  ELSE 1
+              END = 1
+              ORDER BY msg_id ASC
+              LIMIT $1
+              FOR UPDATE SKIP LOCKED
+          )
+          UPDATE pgmq.%I m
+          SET
+              vt = clock_timestamp() + %L,
+              read_ct = read_ct + 1
+          FROM cte
+          WHERE m.msg_id = cte.msg_id
+          RETURNING m.msg_id, m.read_ct, m.enqueued_at, m.vt, m.message;
+          $QUERY$,
+          qtable, conditional, qtable, make_interval(secs => vt)
+      );
+
+      FOR r IN
+        EXECUTE sql USING qty
+      LOOP
+        RETURN NEXT r;
+      END LOOP;
+      IF FOUND THEN
+        RETURN;
+      ELSE
+        PERFORM pg_sleep(poll_interval_ms::numeric / 1000);
+      END IF;
+    END LOOP;
+END;
+$$ language plpgsql;

package/supabase/seed.sql

@@ -0,0 +1,202 @@
+create schema if not exists pgflow_tests;
+
+--------------------------------------------------------------------------------
+--------- reset_db - clears all tables and drops all queues --------------------
+--------------------------------------------------------------------------------
+create or replace function pgflow_tests.reset_db() returns void as $$
+  DELETE FROM pgflow.step_tasks;
+  DELETE FROM pgflow.step_states;
+  DELETE FROM pgflow.runs;
+  DELETE FROM pgflow.deps;
+  DELETE FROM pgflow.steps;
+  DELETE FROM pgflow.flows;
+
+  SELECT pgmq.drop_queue(queue_name) FROM pgmq.list_queues();
+$$ language sql;
+
+--------------------------------------------------------------------------------
+--------- setup_flow - creates a predefined flow and adds steps to it ----------
+--------------------------------------------------------------------------------
+create or replace function pgflow_tests.setup_flow(
+  flow_slug text
+) returns void as $$
+begin
+
+if flow_slug = 'sequential' then
+  PERFORM pgflow.create_flow('sequential', timeout => 1);
+  PERFORM pgflow.add_step('sequential', 'first');
+  PERFORM pgflow.add_step('sequential', 'second', ARRAY['first']);
+  PERFORM pgflow.add_step('sequential', 'last', ARRAY['second']);
+elsif flow_slug = 'sequential_other' then
+  PERFORM pgflow.create_flow('other', timeout => 1);
+  PERFORM pgflow.add_step('other', 'first');
+  PERFORM pgflow.add_step('other', 'second', ARRAY['first']);
+  PERFORM pgflow.add_step('other', 'last', ARRAY['second']);
+elsif flow_slug = 'two_roots' then
+  PERFORM pgflow.create_flow('two_roots', timeout => 1);
+  PERFORM pgflow.add_step('two_roots', 'root_a');
+  PERFORM pgflow.add_step('two_roots', 'root_b');
+  PERFORM pgflow.add_step('two_roots', 'last', ARRAY['root_a', 'root_b']);
+elsif flow_slug = 'two_roots_left_right' then
+  PERFORM pgflow.create_flow('two_roots_left_right', timeout => 1);
+  PERFORM pgflow.add_step('two_roots_left_right', 'connected_root');
+  PERFORM pgflow.add_step('two_roots_left_right', 'disconnected_root');
+  PERFORM pgflow.add_step('two_roots_left_right', 'left', ARRAY['connected_root']);
+  PERFORM pgflow.add_step('two_roots_left_right', 'right', ARRAY['connected_root']);
+else
+  RAISE EXCEPTION 'Unknown test flow: %', flow_slug;
+end if;
+
+end;
+$$ language plpgsql;
+
+--------------------------------------------------------------------------------
+------- poll_and_fail - polls for a task and fails it immediately --------------
+--------------------------------------------------------------------------------
+create or replace function pgflow_tests.poll_and_fail(
+  flow_slug text,
+  vt integer default 1,
+  qty integer default 1
+) returns setof pgflow.step_tasks as $$
+  -- Poll for a task and complete it in one step
+  WITH task AS (
+    SELECT * FROM pgflow.poll_for_tasks(flow_slug, vt, qty) LIMIT 1
+  )
+  SELECT pgflow.fail_task(
+    (SELECT run_id FROM task),
+    (SELECT step_slug FROM task),
+    0,
+    concat(task.step_slug, ' FAILED')
+  )
+  FROM task;
+$$ language sql;
+
+--------------------------------------------------------------------------------
+------- poll_and_complete - polls for a task and completes it immediately ------
+--------------------------------------------------------------------------------
+create or replace function pgflow_tests.poll_and_complete(
+  flow_slug text,
+  vt integer default 1,
+  qty integer default 1
+) returns setof pgflow.step_tasks as $$
+  -- Poll for a task and complete it in one step
+  WITH task AS (
+    SELECT * FROM pgflow.poll_for_tasks(flow_slug, vt, qty) LIMIT 1
+  )
+  SELECT pgflow.complete_task(
+    (SELECT run_id FROM task),
+    (SELECT step_slug FROM task),
+    0,
+    jsonb_build_object('input', task.input)
+  )
+  FROM task;
+$$ language sql;
+
+--------------------------------------------------------------------------------
+------- message_timing - returns messages with added vt_seconds int ------------
+--------------------------------------------------------------------------------
+create or replace function pgflow_tests.message_timing(step_slug text, queue_name text)
+returns table (
+  msg_id bigint,
+  read_ct int,
+  enqueued_at timestamptz,
+  vt timestamptz,
+  message jsonb,
+  vt_seconds int
+)
+language plpgsql
+as $$
+DECLARE
+  qtable TEXT;
+  query TEXT;
+BEGIN
+  qtable := pgmq.format_table_name(queue_name, 'q');
+
+  query := format('
+    SELECT
+      q.msg_id,
+      q.read_ct,
+      q.enqueued_at,
+      q.vt,
+      q.message,
+      extract(epoch from (q.vt - q.enqueued_at))::int as vt_seconds
+    FROM pgmq.%s q
+    JOIN pgflow.step_tasks st ON st.message_id = q.msg_id
+    WHERE st.step_slug = $1', qtable);
+
+  RETURN QUERY EXECUTE query USING step_slug;
+END;
+$$;
+
+--------------------------------------------------------------------------------
+------- reset_message_visibility -----------------------------------------------
+--------------------------------------------------------------------------------
+--
+-- Makes all hidden messages in a queue immediately visible by setting their
+-- visibility time (vt) to the current timestamp.
+--
+-- This is a test utility that allows testing retry logic without waiting for
+-- actual delays to expire. It directly modifies the pgmq queue table.
+--
+-- @param queue_name The name of the queue to modify
+-- @return The number of messages that were made visible
+--
+create or replace function pgflow_tests.reset_message_visibility(
+  queue_name text
+) returns integer as $$
+DECLARE
+  qtable TEXT;
+  query TEXT;
+  updated_count INTEGER;
+BEGIN
+  -- Get the formatted table name for the queue
+  qtable := pgmq.format_table_name(queue_name, 'q');
+
+  -- Construct and execute the query to update all messages' visibility time
+  query := format('
+    UPDATE pgmq.%s
+    SET vt = clock_timestamp()
+    WHERE vt > clock_timestamp()
+    RETURNING 1', qtable);
+
+  -- Execute the query and count the number of updated rows
+  EXECUTE query INTO updated_count;
+
+  -- Return the number of messages that were made visible
+  RETURN COALESCE(updated_count, 0);
+END;
+$$ language plpgsql;
+
+
+--------------------------------------------------------------------------------
+------- assert_retry_delay -----------------------------------------------------
+--------------------------------------------------------------------------------
+--
+-- Asserts that the calculated retry delay matches the expected value.
+--
+-- @param step_slug The slug of the step to check
+-- @param queue_name The name of the queue to check
+-- @param expected_delay The expected delay value
+-- @param description A description of the test case
+-- @return TEXT result from the is() function
+--
+create or replace function pgflow_tests.assert_retry_delay(
+  queue_name text,
+  step_slug text,
+  expected_delay integer,
+  description text
+) returns text as $$
+DECLARE
+  actual_delay INTEGER;
+BEGIN
+  SELECT vt_seconds INTO actual_delay
+  FROM pgflow_tests.message_timing(step_slug, queue_name)
+  LIMIT 1;
+
+  RETURN is(
+    actual_delay,
+    expected_delay,
+    description
+  );
+END;
+$$ language plpgsql;

package/supabase/tests/add_step/basic_step_addition.test.sql

@@ -0,0 +1,29 @@
+begin;
+select plan(3);
+select pgflow_tests.reset_db();
+
+-- Setup
+select pgflow.create_flow('test_flow');
+
+-- Test
+select pgflow.add_step('test_flow', 'first_step');
+select results_eq(
+  $$ SELECT step_slug FROM pgflow.steps WHERE flow_slug = 'test_flow' $$,
+  array['first_step']::text [],
+  'Step should be added to the steps table'
+);
+select is_empty(
+  $$ SELECT * FROM pgflow.deps WHERE flow_slug = 'test_flow' $$,
+  'No dependencies should be added for step with no dependencies'
+);
+select is(
+  (
+    select deps_count::int from pgflow.steps
+    where flow_slug = 'test_flow'
+  ),
+  0::int,
+  'deps_count should be 0 because there are no dependencies'
+);
+
+select * from finish();
+rollback;

package/supabase/tests/add_step/circular_dependency.test.sql

@@ -0,0 +1,21 @@
+begin;
+select plan(1);
+select pgflow_tests.reset_db();
+
+-- Setup
+select pgflow.create_flow('test_flow');
+select pgflow.add_step('test_flow', 'first_step');
+select pgflow.add_step('test_flow', 'second_step', array['first_step']);
+select pgflow.add_step('test_flow', 'third_step', array['second_step']);
+select
+  pgflow.add_step('test_flow', 'fourth_step', array['second_step', 'third_step']);
+
+-- Test
+select throws_ok(
+  $$ SELECT pgflow.add_step('test_flow', 'circular_step', ARRAY['fourth_step', 'circular_step']) $$,
+  'new row for relation "deps" violates check constraint "deps_check"',
+  'Should not allow self-depending steps'
+);
+
+select * from finish();
+rollback;

package/supabase/tests/add_step/flow_isolation.test.sql

@@ -0,0 +1,26 @@
+begin;
+select plan(1);
+select pgflow_tests.reset_db();
+
+-- Setup
+select pgflow.create_flow('test_flow');
+select pgflow.add_step('test_flow', 'first_step');
+select pgflow.create_flow('another_flow');
+
+-- Test
+select pgflow.add_step('another_flow', 'first_step');
+select pgflow.add_step('another_flow', 'another_step', array['first_step']);
+select set_eq(
+  $$
+      SELECT flow_slug, step_slug
+      FROM pgflow.steps WHERE flow_slug = 'another_flow'
+    $$,
+  $$ VALUES
+       ('another_flow', 'another_step'),
+       ('another_flow', 'first_step')
+    $$,
+  'Steps in second flow should be isolated from first flow'
+);
+
+select * from finish();
+rollback;

package/supabase/tests/add_step/idempotent_step_addition.test.sql

@@ -0,0 +1,20 @@
+begin;
+select plan(1);
+select pgflow_tests.reset_db();
+
+-- Setup
+select pgflow.create_flow('test_flow');
+select pgflow.add_step('test_flow', 'first_step');
+
+-- Test
+select pgflow.add_step('test_flow', 'first_step');
+select results_eq(
+  $$
+      SELECT count(*)::int FROM pgflow.steps WHERE flow_slug = 'test_flow' AND step_slug = 'first_step'
+    $$,
+  array[1]::int [],
+  'Calling add_step again for same step does not create a duplicate'
+);
+
+select * from finish();
+rollback;

package/supabase/tests/add_step/invalid_step_slug.test.sql

@@ -0,0 +1,16 @@
+begin;
+select plan(1);
+select pgflow_tests.reset_db();
+
+-- Setup
+select pgflow.create_flow('test_flow');
+
+-- Test
+select throws_ok(
+  $$ SELECT pgflow.add_step('test_flow', '1invalid-slug') $$,
+  'new row for relation "steps" violates check constraint "steps_step_slug_check"',
+  'Should detect and prevent invalid step slug'
+);
+
+select * from finish();
+rollback;

package/supabase/tests/add_step/nonexistent_dependency.test.sql

@@ -0,0 +1,16 @@
+begin;
+select plan(1);
+select pgflow_tests.reset_db();
+
+-- Setup
+select pgflow.create_flow('test_flow');
+
+-- Test
+select throws_ok(
+  $$ SELECT pgflow.add_step('test_flow', 'invalid_dep_step', ARRAY['nonexistent_step']) $$,
+  'insert or update on table "deps" violates foreign key constraint "deps_flow_slug_dep_slug_fkey"',
+  'Should detect and prevent dependency on non-existent step'
+);
+
+select * from finish();
+rollback;

package/supabase/tests/add_step/nonexistent_flow.test.sql

@@ -0,0 +1,13 @@
+begin;
+select plan(1);
+select pgflow_tests.reset_db();
+
+-- Test
+select throws_ok(
+  $$ SELECT pgflow.add_step('nonexistent_flow', 'some_step') $$,
+  'insert or update on table "steps" violates foreign key constraint "steps_flow_slug_fkey"',
+  'Should not allow adding step to non-existent flow'
+);
+
+select * from finish();
+rollback;