Add Options to SQL Enqueue (#1260)

Ports dbos-inc/dbos-transact-py#684 Fixes
#1258

Author: kraftp

src/sysdb_migrations/internal/migrations.ts

@@ -421,9 +421,9 @@ export function allMigrations(
         `CREATE TABLE "${schemaName}"."queues" (
           "queue_id" TEXT PRIMARY KEY DEFAULT gen_random_uuid()::TEXT,
           "name" TEXT NOT NULL UNIQUE,
-          "concurrency" INTEGER,
-          "worker_concurrency" INTEGER,
-          "rate_limit_max" INTEGER,
+          "concurrency" INT4,
+          "worker_concurrency" INT4,
+          "rate_limit_max" INT4,
           "rate_limit_period_sec" DOUBLE PRECISION,
           "priority_enabled" BOOLEAN NOT NULL DEFAULT FALSE,
           "partition_queue" BOOLEAN NOT NULL DEFAULT FALSE,
@@ -527,5 +527,108 @@ export function allMigrations(
         `CREATE INDEX ${c} IF NOT EXISTS "idx_workflow_status_started_at" ON "${schemaName}"."workflow_status" ("started_at_epoch_ms") WHERE "started_at_epoch_ms" IS NOT NULL`,
       ],
     },
+    {
+      pg: [
+        `DROP FUNCTION IF EXISTS "${schemaName}".enqueue_workflow(
+            TEXT, TEXT, JSON[], JSON, TEXT, TEXT, TEXT, TEXT, BIGINT, BIGINT, TEXT, INTEGER, TEXT
+        )`,
+        `CREATE OR REPLACE FUNCTION "${schemaName}".enqueue_workflow(
+            workflow_name TEXT,
+            queue_name TEXT,
+            positional_args JSON[] DEFAULT ARRAY[]::JSON[],
+            named_args JSON DEFAULT '{}'::JSON,
+            class_name TEXT DEFAULT NULL,
+            config_name TEXT DEFAULT NULL,
+            workflow_id TEXT DEFAULT NULL,
+            app_version TEXT DEFAULT NULL,
+            timeout_ms BIGINT DEFAULT NULL,
+            deadline_epoch_ms BIGINT DEFAULT NULL,
+            deduplication_id TEXT DEFAULT NULL,
+            priority INT4 DEFAULT NULL,
+            queue_partition_key TEXT DEFAULT NULL,
+            authenticated_user TEXT DEFAULT NULL,
+            authenticated_roles TEXT DEFAULT NULL,
+            delay_until_epoch_ms BIGINT DEFAULT NULL
+        ) RETURNS TEXT AS $$
+        DECLARE
+            v_workflow_id TEXT;
+            v_serialized_inputs TEXT;
+            v_owner_xid TEXT;
+            v_now BIGINT;
+            v_recovery_attempts INT4 := 0;
+            v_priority INT4;
+            v_status TEXT;
+        BEGIN
+
+            -- Validate required parameters
+            IF workflow_name IS NULL OR workflow_name = '' THEN
+                RAISE EXCEPTION 'Workflow name cannot be null or empty';
+            END IF;
+            IF queue_name IS NULL OR queue_name = '' THEN
+                RAISE EXCEPTION 'Queue name cannot be null or empty';
+            END IF;
+            IF named_args IS NOT NULL AND jsonb_typeof(named_args::jsonb) != 'object' THEN
+                RAISE EXCEPTION 'Named args must be a JSON object';
+            END IF;
+            IF workflow_id IS NOT NULL AND workflow_id = '' THEN
+                RAISE EXCEPTION 'Workflow ID cannot be an empty string if provided.';
+            END IF;
+            IF delay_until_epoch_ms IS NOT NULL AND delay_until_epoch_ms < 0 THEN
+                RAISE EXCEPTION 'delay_until_epoch_ms must be >= 0';
+            END IF;
+
+            v_workflow_id := COALESCE(workflow_id, gen_random_uuid()::TEXT);
+            v_owner_xid := gen_random_uuid()::TEXT;
+            v_priority := COALESCE(priority, 0);
+            v_serialized_inputs := json_build_object(
+                'positionalArgs', positional_args,
+                'namedArgs', named_args
+            )::TEXT;
+            v_now := EXTRACT(epoch FROM now()) * 1000;
+            v_status := CASE WHEN delay_until_epoch_ms IS NULL THEN 'ENQUEUED' ELSE 'DELAYED' END;
+
+            INSERT INTO "${schemaName}".workflow_status (
+                workflow_uuid, status, inputs,
+                name, class_name, config_name,
+                queue_name, deduplication_id, priority, queue_partition_key,
+                application_version,
+                created_at, updated_at, recovery_attempts,
+                workflow_timeout_ms, workflow_deadline_epoch_ms,
+                parent_workflow_id, owner_xid, serialization,
+                authenticated_user, authenticated_roles,
+                delay_until_epoch_ms
+            ) VALUES (
+                v_workflow_id, v_status, v_serialized_inputs,
+                workflow_name, class_name, config_name,
+                queue_name, deduplication_id, v_priority, queue_partition_key,
+                app_version,
+                v_now, v_now, v_recovery_attempts,
+                timeout_ms, deadline_epoch_ms,
+                NULL, v_owner_xid, 'portable_json',
+                authenticated_user, authenticated_roles,
+                delay_until_epoch_ms
+            )
+            ON CONFLICT (workflow_uuid)
+            DO UPDATE SET
+                updated_at = EXCLUDED.updated_at;
+
+            RETURN v_workflow_id;
+
+        EXCEPTION
+            WHEN unique_violation THEN
+                RAISE EXCEPTION 'DBOS queue duplicated'
+                   USING DETAIL = format('Workflow %s with queue %s and deduplication ID %s already exists', v_workflow_id, queue_name, deduplication_id),
+                        ERRCODE = 'unique_violation';
+        END;
+        $$ LANGUAGE plpgsql;`,
+        ...(isCockroach
+          ? []
+          : [
+              `ALTER FUNCTION "${schemaName}".enqueue_workflow(
+            TEXT, TEXT, JSON[], JSON, TEXT, TEXT, TEXT, TEXT, BIGINT, BIGINT, TEXT, INT4, TEXT, TEXT, TEXT, BIGINT
+        ) SET search_path = pg_catalog, pg_temp`,
+            ]),
+      ],
+    },
   ];
 }

tests/pg-client-functions.test.ts

@@ -153,6 +153,132 @@ describe('PostgreSQL Client Functions', () => {
     }
   });
 
+  test('pg-enqueue-with-auth-metadata', async () => {
+    const wfid = `pg-auth-test-${Date.now()}`;
+    const user = 'alice';
+    const roles = ['admin', 'reader'];
+    const dbClient = new Client(poolConfig);
+
+    try {
+      await dbClient.connect();
+
+      await dbClient.query(
+        `
+        SELECT dbos.enqueue_workflow(
+          workflow_name => 'enqueueTest',
+          queue_name => $1,
+          class_name => 'ClientTest',
+          workflow_id => $2,
+          positional_args => ARRAY[$3::JSON, $4::JSON, $5::JSON],
+          authenticated_user => $6,
+          authenticated_roles => $7
+        )
+      `,
+        [
+          testQueueName,
+          wfid,
+          JSON.stringify(42),
+          JSON.stringify('test'),
+          JSON.stringify({ first: 'John', last: 'Doe', age: 30 }),
+          user,
+          JSON.stringify(roles),
+        ],
+      );
+    } finally {
+      await dbClient.end();
+    }
+
+    try {
+      await DBOS.launch();
+      await DBOS.registerQueue(testQueueName, { onConflict: 'always_update' });
+
+      const handle = DBOS.retrieveWorkflow<string>(wfid);
+      const status = await handle.getStatus();
+      expect(status).toBeDefined();
+      expect(status!.authenticatedUser).toBe(user);
+      expect(status!.authenticatedRoles).toEqual(roles);
+
+      const result = await handle.getResult();
+      expect(result).toBe('42-test-{"first":"John","last":"Doe","age":30}');
+    } finally {
+      await DBOS.shutdown();
+    }
+  });
+
+  test('pg-enqueue-with-delay', async () => {
+    const wfid = `pg-delay-test-${Date.now()}`;
+    // Far enough in the future that the workflow won't be promoted during the test
+    const delayUntil = Date.now() + 60_000;
+    const dbClient = new Client(poolConfig);
+
+    try {
+      await dbClient.connect();
+
+      await dbClient.query(
+        `
+        SELECT dbos.enqueue_workflow(
+          workflow_name => 'enqueueTest',
+          queue_name => $1,
+          class_name => 'ClientTest',
+          workflow_id => $2,
+          positional_args => ARRAY[$3::JSON, $4::JSON, $5::JSON],
+          delay_until_epoch_ms => $6
+        )
+      `,
+        [
+          testQueueName,
+          wfid,
+          JSON.stringify(42),
+          JSON.stringify('test'),
+          JSON.stringify({ first: 'John', last: 'Doe', age: 30 }),
+          delayUntil,
+        ],
+      );
+    } finally {
+      await dbClient.end();
+    }
+
+    try {
+      await DBOS.launch();
+      await DBOS.registerQueue(testQueueName, { onConflict: 'always_update' });
+
+      const handle = DBOS.retrieveWorkflow<string>(wfid);
+      const status = await handle.getStatus();
+      expect(status).toBeDefined();
+      expect(status!.status).toBe('DELAYED');
+      expect(status!.delayUntilEpochMS).toBe(delayUntil);
+
+      await DBOS.cancelWorkflow(wfid);
+    } finally {
+      await DBOS.shutdown();
+    }
+  });
+
+  test('pg-enqueue-negative-delay-rejected', async () => {
+    const wfid = `pg-negative-delay-${Date.now()}`;
+    const dbClient = new Client(poolConfig);
+
+    try {
+      await dbClient.connect();
+
+      await expect(
+        dbClient.query(
+          `
+          SELECT dbos.enqueue_workflow(
+            workflow_name => 'enqueueTest',
+            queue_name => $1,
+            workflow_id => $2,
+            delay_until_epoch_ms => -1
+          )
+        `,
+          [testQueueName, wfid],
+        ),
+      ).rejects.toThrow(/delay_until_epoch_ms must be >= 0/);
+    } finally {
+      await dbClient.end();
+    }
+  });
+
   test('pg-enqueue-gen-wfid', async () => {
     const dbClient = new Client(poolConfig);
     let wfid: string;