pgflow 0.0.16 → 0.0.17-yolo

package/dist/migrations/000030_read_with_poll_backport.sql

@@ -0,0 +1,70 @@
+--------------------------------------------------------------------------------
+-- 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 pgflow.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/dist/migrations/000040_poll_for_tasks.sql

@@ -0,0 +1,100 @@
+create or replace function pgflow.poll_for_tasks(
+  queue_name text,
+  vt integer,
+  qty integer,
+  max_poll_seconds integer default 5,
+  poll_interval_ms integer default 100
+)
+returns setof pgflow.step_task_record
+volatile
+set search_path to ''
+as $$
+
+with read_messages as (
+  select *
+  from pgflow.read_with_poll(
+    queue_name,
+    vt,
+    qty,
+    max_poll_seconds,
+    poll_interval_ms
+  )
+),
+tasks as (
+  select
+    task.flow_slug,
+    task.run_id,
+    task.step_slug,
+    task.task_index,
+    task.message_id
+  from pgflow.step_tasks as task
+  join read_messages as message on message.msg_id = task.message_id
+  where task.message_id = message.msg_id
+    and task.status = 'queued'
+),
+increment_attempts as (
+  update pgflow.step_tasks
+  set attempts_count = attempts_count + 1
+  from tasks
+  where step_tasks.message_id = tasks.message_id
+  and status = 'queued'
+),
+runs as (
+  select
+    r.run_id,
+    r.input
+  from pgflow.runs r
+  where r.run_id in (select run_id from tasks)
+),
+deps as (
+  select
+    st.run_id,
+    st.step_slug,
+    dep.dep_slug,
+    dep_task.output as dep_output
+  from tasks st
+  join pgflow.deps dep on dep.flow_slug = st.flow_slug and dep.step_slug = st.step_slug
+  join pgflow.step_tasks dep_task on
+    dep_task.run_id = st.run_id and
+    dep_task.step_slug = dep.dep_slug and
+    dep_task.status = 'completed'
+),
+deps_outputs as (
+  select
+    d.run_id,
+    d.step_slug,
+    jsonb_object_agg(d.dep_slug, d.dep_output) as deps_output
+  from deps d
+  group by d.run_id, d.step_slug
+),
+timeouts as (
+  select
+    task.message_id,
+    coalesce(step.opt_timeout, flow.opt_timeout) + 2 as vt_delay
+  from tasks task
+  join pgflow.flows flow on flow.flow_slug = task.flow_slug
+  join pgflow.steps step on step.flow_slug = task.flow_slug and step.step_slug = task.step_slug
+)
+
+select
+  st.flow_slug,
+  st.run_id,
+  st.step_slug,
+  jsonb_build_object('run', r.input) ||
+  coalesce(dep_out.deps_output, '{}'::jsonb) as input,
+  st.message_id as msg_id
+from tasks st
+join runs r on st.run_id = r.run_id
+left join deps_outputs dep_out on
+  dep_out.run_id = st.run_id and
+  dep_out.step_slug = st.step_slug
+cross join lateral (
+  -- TODO: this is slow because it calls set_vt for each row, and set_vt
+  --       builds dynamic query from string every time it is called
+  --       implement set_vt_batch(msgs_ids bigint[], vt_delays int[])
+  select pgmq.set_vt(queue_name, st.message_id,
+    (select t.vt_delay from timeouts t where t.message_id = st.message_id)
+  )
+) set_vt;
+
+$$ language sql;

package/dist/migrations/000045_maybe_complete_run.sql

@@ -0,0 +1,30 @@
+create or replace function pgflow.maybe_complete_run(run_id uuid)
+returns void
+language sql
+volatile
+set search_path to ''
+as $$
+  -- Update run status to completed and set output when there are no remaining steps
+  -- All done in a single declarative SQL statement
+  UPDATE pgflow.runs
+  SET
+    status = 'completed',
+    output = (
+      -- Get outputs from final steps (steps that are not dependencies for other steps)
+      SELECT jsonb_object_agg(st.step_slug, st.output)
+      FROM pgflow.step_tasks st
+      JOIN pgflow.step_states ss ON ss.run_id = st.run_id AND ss.step_slug = st.step_slug
+      JOIN pgflow.runs r ON r.run_id = ss.run_id AND r.flow_slug = ss.flow_slug
+      WHERE st.run_id = maybe_complete_run.run_id
+        AND st.status = 'completed'
+        AND NOT EXISTS (
+          SELECT 1
+          FROM pgflow.deps d
+          WHERE d.flow_slug = ss.flow_slug
+            AND d.dep_slug = ss.step_slug
+        )
+    )
+  WHERE pgflow.runs.run_id = maybe_complete_run.run_id
+    AND pgflow.runs.remaining_steps = 0
+    AND pgflow.runs.status != 'completed';
+$$;

package/dist/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/dist/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/dist/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/dist/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/dist/package.json

@@ -0,0 +1,38 @@
+{
+  "name": "pgflow",
+  "version": "0.0.17-yolo",
+  "type": "module",
+  "main": "./dist/index.js",
+  "typings": "./dist/index.d.ts",
+  "bin": "./dist/index.js",
+  "exports": {
+    "./package.json": "./package.json",
+    ".": {
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.js"
+    }
+  },
+  "types": "./dist/index.d.ts",
+  "module": "./dist/index.js",
+  "devDependencies": {
+    "@types/node": "^22.14.1",
+    "tsx": "^4.19.3"
+  },
+  "dependencies": {
+    "@clack/prompts": "^0.10.1",
+    "@commander-js/extra-typings": "^13.1.0",
+    "chalk": "^5.4.1",
+    "commander": "^13.1.0",
+    "toml-patch": "^0.2.3"
+  },
+  "publishConfig": {
+    "access": "public",
+    "directory": "."
+  },
+  "files": [
+    "dist"
+  ],
+  "scripts": {
+    "postinstall": "chmod +x dist/index.js || true"
+  }
+}

package/dist/tsconfig.lib.tsbuildinfo

@@ -0,0 +1 @@
+{"version":"5.6.3"}

package/package.json

@@ -1,19 +1,19 @@
 {
   "name": "pgflow",
-  "version": "0.0.16",
+  "version": "0.0.17-yolo",
   "type": "module",
-  "main": "./dist/src/index.js",
-  "typings": "./dist/src/index.d.ts",
-  "bin": "./dist/src/index.js",
+  "main": "./dist/index.js",
+  "typings": "./dist/index.d.ts",
+  "bin": "./dist/index.js",
   "exports": {
     "./package.json": "./package.json",
     ".": {
-      "types": "./dist/src/index.d.ts",
-      "import": "./dist/src/index.js"
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.js"
     }
   },
-  "types": "./dist/src/index.d.ts",
-  "module": "./dist/src/index.js",
+  "types": "./dist/index.d.ts",
+  "module": "./dist/index.js",
   "devDependencies": {
     "@types/node": "^22.14.1",
     "tsx": "^4.19.3"
@@ -31,5 +31,8 @@
   },
   "files": [
     "dist"
-  ]
-}
+  ],
+  "scripts": {
+    "postinstall": "chmod +x dist/index.js || true"
+  }
+}