add mark_worker_stopped() SQL function for graceful shutdown signaling

Author: jumski

pkgs/core/schemas/0058_function_mark_worker_stopped.sql

@@ -0,0 +1,17 @@
+-- Mark Worker Stopped
+-- Sets stopped_at timestamp on a worker row for graceful shutdown signaling
+
+drop function if exists pgflow.mark_worker_stopped(uuid);
+
+create or replace function pgflow.mark_worker_stopped(
+  worker_id uuid
+) returns void
+language sql
+as $$
+  update pgflow.workers
+  set stopped_at = clock_timestamp()
+  where workers.worker_id = mark_worker_stopped.worker_id;
+$$;
+
+comment on function pgflow.mark_worker_stopped(uuid) is
+'Marks a worker as stopped for graceful shutdown. Called by workers on beforeunload.';

pkgs/core/src/database-types.ts

@@ -515,6 +515,7 @@ export type Database = {
       get_run_with_states: { Args: { run_id: string }; Returns: Json }
       is_local: { Args: never; Returns: boolean }
       is_valid_slug: { Args: { slug: string }; Returns: boolean }
+      mark_worker_stopped: { Args: { worker_id: string }; Returns: undefined }
       maybe_complete_run: { Args: { run_id: string }; Returns: undefined }
       poll_for_tasks: {
         Args: {

pkgs/core/supabase/migrations/20251204165231_pgflow_temp_mark_worker_stopped.sql

@@ -0,0 +1,8 @@
+-- Create "mark_worker_stopped" function
+CREATE FUNCTION "pgflow"."mark_worker_stopped" ("worker_id" uuid) RETURNS void LANGUAGE sql AS $$
+update pgflow.workers
+  set stopped_at = clock_timestamp()
+  where workers.worker_id = mark_worker_stopped.worker_id;
+$$;
+-- Set comment to function: "mark_worker_stopped"
+COMMENT ON FUNCTION "pgflow"."mark_worker_stopped" IS 'Marks a worker as stopped for graceful shutdown. Called by workers on beforeunload.';

pkgs/core/supabase/migrations/atlas.sum

@@ -1,4 +1,4 @@
-h1:eMTY2CZPFm2BibVAHjtIMhCWIFcmTcqQq9VRXfs1/1A=
+h1:lFHz8CiL5ustrpb+hKVK2Ae5/jVpsOq1AX5vVgxj9Yo=
 20250429164909_pgflow_initial.sql h1:I3n/tQIg5Q5nLg7RDoU3BzqHvFVjmumQxVNbXTPG15s=
 20250517072017_pgflow_fix_poll_for_tasks_to_use_separate_statement_for_polling.sql h1:wTuXuwMxVniCr3ONCpodpVWJcHktoQZIbqMZ3sUHKMY=
 20250609105135_pgflow_add_start_tasks_and_started_status.sql h1:ggGanW4Wyt8Kv6TWjnZ00/qVb3sm+/eFVDjGfT8qyPg=
@@ -20,3 +20,4 @@ h1:eMTY2CZPFm2BibVAHjtIMhCWIFcmTcqQq9VRXfs1/1A=
 20251204142050_pgflow_temp_ensure_flow_compiled_auto_detect.sql h1:VwqZiOcVaCahb6BZ918ioFLgwQcF/sy1TR9a4lSnVvs=
 20251204145037_pgflow_temp_worker_functions_schema.sql h1:5DJJEP0jcg7yapTe7t6FX2ypZIj92lGvJ12AX8g5fz4=
 20251204164612_pgflow_temp_track_worker_function.sql h1:SC8Z4Un37A2XVaIIvyZJ0eQ/7r0JENSi6RcBx4GaUCE=
+20251204165231_pgflow_temp_mark_worker_stopped.sql h1:5SQ353rHkcRMZ/EPTnm4+ZkC7NSQ4YhftU+Wyq00Cr0=

pkgs/core/supabase/tests/mark_worker_stopped/sets_stopped_at.test.sql

@@ -0,0 +1,35 @@
+-- Test: mark_worker_stopped() sets stopped_at on the worker row
+begin;
+select plan(3);
+
+-- Setup: Create a worker entry
+insert into pgflow.workers (worker_id, queue_name, function_name, started_at, last_heartbeat_at)
+values ('11111111-1111-1111-1111-111111111111', 'test_queue', 'test_function', now(), now());
+
+-- Verify worker exists and stopped_at is NULL
+select is(
+  (select stopped_at from pgflow.workers where worker_id = '11111111-1111-1111-1111-111111111111'),
+  null,
+  'Worker should have NULL stopped_at initially'
+);
+
+-- Execute: Mark worker as stopped
+select pgflow.mark_worker_stopped('11111111-1111-1111-1111-111111111111'::uuid);
+
+-- Test: stopped_at should be set
+select isnt(
+  (select stopped_at from pgflow.workers where worker_id = '11111111-1111-1111-1111-111111111111'),
+  null,
+  'Worker should have stopped_at set after marking stopped'
+);
+
+-- Test: stopped_at should be recent (within last second)
+select ok(
+  (select stopped_at > clock_timestamp() - interval '1 second'
+   from pgflow.workers
+   where worker_id = '11111111-1111-1111-1111-111111111111'),
+  'stopped_at should be recent (within last second)'
+);
+
+select finish();
+rollback;