Custom Workflow Attributes (#1275)

Analogous to dbos-inc/dbos-transact-py#720

Author: kraftp

schemas/system_db_schema.ts

@@ -34,6 +34,7 @@ export interface workflow_status {
   delay_until_epoch_ms?: number | null;
   rate_limited?: boolean;
   completed_at?: number | null;
+  attributes?: Record<string, unknown> | null; // Custom key-value attributes attached at creation, stored as JSONB.
 }
 
 export interface notifications {

src/client.ts

@@ -19,6 +19,7 @@ import {
   type WorkflowHandle,
   WorkflowSerializationFormat,
   type WorkflowStatus,
+  validateWorkflowAttributes,
 } from './workflow';
 import { cancellableSleep } from './utils';
 import {
@@ -143,6 +144,11 @@ export interface ClientEnqueueOptions {
    *     and the handle resolves with the original workflow's result.
    */
   duplicationPolicy?: DuplicationPolicy;
+  /**
+   * Custom key-value attributes to attach to the workflow at creation.
+   * Attributes are searchable via the `attributes` filter of `listWorkflows`.
+   */
+  attributes?: Record<string, unknown>;
 }
 
 /**
@@ -292,6 +298,7 @@ export class DBOSClient {
     options: ClientEnqueueOptions,
     ...args: Parameters<T>
   ): Promise<WorkflowHandle<Awaited<ReturnType<T>>>> {
+    validateWorkflowAttributes(options.attributes);
     const { workflowName, workflowClassName, workflowConfigName, queueName, appVersion } = options;
     const workflowUUID = options.workflowID ?? randomUUID();
 
@@ -325,6 +332,7 @@ export class DBOSClient {
       queuePartitionKey: options.queuePartitionKey,
       serialization: serparam.serialization,
       delayUntilEpochMS,
+      attributes: options.attributes,
     };
 
     let finalID: string;
@@ -376,6 +384,7 @@ export class DBOSClient {
     positionalArgs: unknown[],
     namedArgs?: { [key: string]: unknown },
   ): Promise<WorkflowHandle<T>> {
+    validateWorkflowAttributes(options.attributes);
     const { workflowName, workflowClassName, workflowConfigName, queueName, appVersion } = options;
     const workflowUUID = options.workflowID ?? randomUUID();
 
@@ -414,6 +423,7 @@ export class DBOSClient {
       queuePartitionKey: options.queuePartitionKey,
       serialization: serparam.serialization,
       delayUntilEpochMS,
+      attributes: options.attributes,
     };
 
     let finalID: string;

src/conductor/conductor.ts

@@ -284,6 +284,7 @@ export class Conductor {
               queuesOnly: body.queues_only,
               wasForkedFrom: body.was_forked_from,
               hasParent: body.has_parent,
+              attributes: body.attributes,
             };
             let workflowsOutput: protocol.WorkflowsOutput[] = [];
             try {
@@ -323,6 +324,7 @@ export class Conductor {
               executorId: bodyQueued.executor_id,
               wasForkedFrom: bodyQueued.was_forked_from,
               hasParent: bodyQueued.has_parent,
+              attributes: bodyQueued.attributes,
             };
             let queuedWFOutput: protocol.WorkflowsOutput[] = [];
             try {

src/conductor/protocol.ts

@@ -203,6 +203,7 @@ export interface ListWorkflowsBody {
   queues_only?: boolean;
   was_forked_from?: boolean;
   has_parent?: boolean;
+  attributes?: Record<string, unknown>;
 }
 
 export class WorkflowsOutput {
@@ -233,6 +234,7 @@ export class WorkflowsOutput {
   ParentWorkflowID?: string;
   DelayUntilEpochMS?: string;
   CompletedAt?: string;
+  Attributes?: string;
 
   constructor(info: WorkflowStatus) {
     // Mark empty fields as undefined
@@ -264,6 +266,8 @@ export class WorkflowsOutput {
     this.ParentWorkflowID = info.parentWorkflowID;
     this.DelayUntilEpochMS = info.delayUntilEpochMS !== undefined ? String(info.delayUntilEpochMS) : undefined;
     this.CompletedAt = info.completedAt !== undefined ? String(info.completedAt) : undefined;
+    // JSON rather than inspect() so the wire format stays parseable by Conductor.
+    this.Attributes = info.attributes !== undefined ? JSON.stringify(info.attributes) : undefined;
   }
 }
 
@@ -329,6 +333,7 @@ export interface ListQueuedWorkflowsBody {
   executor_id?: string | string[];
   was_forked_from?: boolean;
   has_parent?: boolean;
+  attributes?: Record<string, unknown>;
 }
 
 export class ListQueuedWorkflowsRequest implements BaseMessage {

src/dbos-executor.ts

@@ -519,6 +519,7 @@ export class DBOSExecutor {
       parentWorkflowID: callerID,
       serialization: funcArgs.sername,
       delayUntilEpochMS,
+      attributes: params.workflowAttributes,
     };
 
     if (isTempWorkflow) {

src/dbos.ts

@@ -29,6 +29,7 @@ import {
   WorkflowParams,
   WorkflowSerializationFormat,
   WorkflowStatus,
+  validateWorkflowAttributes,
 } from './workflow';
 import { DLogger, GlobalLogger } from './telemetry/logs';
 import {
@@ -171,6 +172,10 @@ export interface StartWorkflowParams {
   //     Requires `queueName` and `enqueueOptions.deduplicationID`. Arguments passed by the
   //     colliding caller are discarded and the handle resolves with the original workflow's result.
   duplicationPolicy?: DuplicationPolicy;
+  // Custom key-value attributes to attach to the workflow at creation. Recorded in the
+  // workflow status and searchable via the `attributes` filter of `DBOS.listWorkflows`.
+  // Not inherited by child workflows.
+  workflowAttributes?: Record<string, unknown>;
 }
 
 /**
@@ -1265,8 +1270,10 @@ export class DBOS {
       const existing: any = {};
       for (const k of Object.keys(options) as (keyof DBOSContextOptions)[]) {
         if (Object.hasOwn(pctx, k))
+          // Save the current value (not the incoming one) so the finally block can restore it,
+          // letting nested `with*` blocks correctly roll back to the outer value.
           // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
-          existing[k] = options[k];
+          existing[k] = pctx[k];
         // eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-member-access
         (pctx as any)[k] = options[k];
       }
@@ -1773,6 +1780,7 @@ export class DBOS {
     startWfFuncId?: number,
   ): Promise<InternalWFHandle<Return>> {
     ensureDBOSIsLaunched('workflows');
+    validateWorkflowAttributes(params.workflowAttributes);
     const wfId = getNextWFID(params.workflowID);
     const ppctx = getCurrentContextStore();
 
@@ -1839,6 +1847,7 @@ export class DBOS {
           params.timeoutMS === null || ppctx?.workflowTimeoutMS === null ? undefined : ppctx?.deadlineEpochMS,
         enqueueOptions: params.enqueueOptions,
         duplicationPolicy: params.duplicationPolicy,
+        workflowAttributes: params.workflowAttributes,
       };
 
       return await invokeRegOp(wfParams, pwfid, funcId);
@@ -1850,6 +1859,7 @@ export class DBOS {
         configuredInstance: instance,
         timeoutMS,
         duplicationPolicy: params.duplicationPolicy,
+        workflowAttributes: params.workflowAttributes,
       };
 
       return await invokeRegOp(wfParams, undefined, undefined);

src/scheduler/scheduler.ts

@@ -65,6 +65,11 @@ interface ScheduleLoopEntry {
   scheduleId: string;
 }
 
+// During the first minute after startup, poll for schedules every second so schedules
+// registered around launch are picked up promptly rather than after a full polling interval.
+const STARTUP_FAST_POLL_DURATION_MS = 60_000;
+const STARTUP_FAST_POLL_INTERVAL_MS = 1_000;
+
 export class DynamicSchedulerLoop implements DBOSLifecycleCallback {
   readonly #mainController = new AbortController();
   #pollingPromise: Promise<void> | undefined;
@@ -99,14 +104,20 @@ export class DynamicSchedulerLoop implements DBOSLifecycleCallback {
   }
 
   async #pollingLoop(signal: AbortSignal): Promise<void> {
+    const startupDeadline = Date.now() + STARTUP_FAST_POLL_DURATION_MS;
+    const pollTimeout = () =>
+      Date.now() < startupDeadline
+        ? Math.min(STARTUP_FAST_POLL_INTERVAL_MS, this.#pollingIntervalMs)
+        : this.#pollingIntervalMs;
+
     while (!signal.aborted) {
       let schedules: WorkflowScheduleInternal[];
       try {
         const executor = DBOSExecutor.globalInstance!;
         schedules = await executor.systemDatabase.listSchedules();
       } catch (e) {
         DBOS.logger.warn(`Dynamic scheduler: error listing schedules: ${(e as Error).message}`);
-        await interruptibleSleep(this.#pollingIntervalMs, signal);
+        await interruptibleSleep(pollTimeout(), signal);
         continue;
       }
 
@@ -179,7 +190,7 @@ export class DynamicSchedulerLoop implements DBOSLifecycleCallback {
         }
       }
 
-      await interruptibleSleep(this.#pollingIntervalMs, signal);
+      await interruptibleSleep(pollTimeout(), signal);
     }
   }
 

src/sysdb_migrations/internal/migrations.ts

@@ -657,5 +657,15 @@ export function allMigrations(
           ]
         : [],
     },
+    // ADD COLUMN with no default is catalog-only; the GIN index built in the same
+    // transaction covers zero rows, so no CONCURRENTLY is needed. The index serves
+    // containment (@>) filters on workflow attributes; on CockroachDB, USING GIN
+    // creates an inverted index.
+    {
+      pg: [
+        `ALTER TABLE "${schemaName}"."workflow_status" ADD COLUMN IF NOT EXISTS "attributes" JSONB`,
+        `CREATE INDEX IF NOT EXISTS "idx_workflow_status_attributes" ON "${schemaName}"."workflow_status" USING GIN ("attributes") WHERE "attributes" IS NOT NULL`,
+      ],
+    },
   ];
 }

src/system_database.ts

@@ -215,6 +215,8 @@ export interface WorkflowStatusInternal {
   serialization: string | null;
   delayUntilEpochMS?: number;
   completedAt?: number;
+  // Custom key-value attributes attached to the workflow at creation. Not inherited by child workflows.
+  attributes?: Record<string, unknown>;
 }
 
 export interface EnqueueOptions {
@@ -429,6 +431,7 @@ function mapWorkflowStatus(row: workflow_status): WorkflowStatusInternal {
     serialization: row.serialization,
     delayUntilEpochMS: row.delay_until_epoch_ms ? Number(row.delay_until_epoch_ms) : undefined,
     completedAt: row.completed_at ? Number(row.completed_at) : undefined,
+    attributes: row.attributes ?? undefined,
   };
 }
 
@@ -1301,7 +1304,7 @@ export class SystemDatabase {
       const { rows: statusRows } = await client.query<workflow_status>(
         `SELECT workflow_uuid, name, class_name, config_name, application_id,
                 authenticated_user, authenticated_roles, assumed_role, inputs, serialization,
-                request, application_version
+                request, application_version, attributes
          FROM "${this.schemaName}".workflow_status
          WHERE workflow_uuid = ANY($1)`,
         [originalWorkflowIDs],
@@ -1334,6 +1337,7 @@ export class SystemDatabase {
         'queue_partition_key',
         'forked_from',
         'serialization',
+        'attributes',
       ];
       if (options.timeoutMS !== undefined) {
         insertCols.push('workflow_timeout_ms');
@@ -1364,6 +1368,7 @@ export class SystemDatabase {
           options.queuePartitionKey ?? null,
           origID,
           ws.serialization,
+          ws.attributes ? JSON.stringify(ws.attributes) : null,
         );
         if (options.timeoutMS !== undefined) {
           params.push(options.timeoutMS);
@@ -2724,6 +2729,7 @@ export class SystemDatabase {
       'parent_workflow_id',
       'delay_until_epoch_ms',
       'completed_at',
+      'attributes',
     ];
 
     input.loadInput = input.loadInput ?? true;
@@ -2811,6 +2817,14 @@ export class SystemDatabase {
       }
     }
 
+    // Match workflows whose attributes JSONB contains all the given key-value pairs.
+    // The `@>` containment operator is served by the GIN index on the attributes column.
+    if (input.attributes && Object.keys(input.attributes).length > 0) {
+      whereClauses.push(`attributes @> $${paramCounter}::jsonb`);
+      params.push(JSON.stringify(input.attributes));
+      paramCounter++;
+    }
+
     if (input.startTime) {
       whereClauses.push(`created_at >= $${paramCounter}`);
       params.push(new Date(input.startTime).getTime());
@@ -3555,8 +3569,9 @@ export class SystemDatabase {
           parent_workflow_id,
           serialization,
           owner_xid,
-          delay_until_epoch_ms
-        ) VALUES($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $26, $27, $28)
+          delay_until_epoch_ms,
+          attributes
+        ) VALUES($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $26, $27, $28, $29)
         ON CONFLICT (workflow_uuid)
           DO UPDATE SET
             recovery_attempts = CASE
@@ -3601,6 +3616,7 @@ export class SystemDatabase {
           initStatus.serialization,
           ownerXid,
           initStatus.delayUntilEpochMS ?? null,
+          initStatus.attributes ? JSON.stringify(initStatus.attributes) : null,
         ],
       );
       if (rows.length === 0) {

src/workflow.ts

@@ -5,6 +5,40 @@ import { deserializePositionalArgs, registerSerializationRecipe } from './serial
 import { DBOS, resolvePollingIntervalMs, runInternalStep, type PollingOptions } from './dbos';
 import { DuplicationPolicy, EnqueueOptions } from './system_database';
 import { DBOSExecutor } from './dbos-executor';
+import { DBOSError } from './error';
+
+/**
+ * Validate that custom workflow attributes, if provided, are a plain key-value object that
+ * can be serialized to JSON for storage in the JSONB `attributes` column. A key-value object
+ * is required because attributes are queried with the `@>` containment filter; scalars and
+ * arrays would store but never match meaningfully. Called at the workflow-creation entry
+ * points (`startWorkflow`/enqueue) so invalid input fails fast at the call site rather than
+ * surfacing as a database error at insert time.
+ */
+export function validateWorkflowAttributes(attributes: unknown): void {
+  if (attributes === undefined || attributes === null) {
+    return;
+  }
+  if (typeof attributes !== 'object' || Array.isArray(attributes)) {
+    throw new DBOSError(
+      `Invalid workflow attributes: must be a key-value object, got ${
+        Array.isArray(attributes) ? 'array' : typeof attributes
+      }.`,
+    );
+  }
+  let serialized: string;
+  try {
+    serialized = JSON.stringify(attributes);
+  } catch (e) {
+    throw new DBOSError(`Invalid workflow attributes: must be JSON-serializable. ${(e as Error).message}`);
+  }
+  // JSON.stringify silently drops keys whose values are functions or undefined, and returns
+  // "{}" for objects with no serializable keys (e.g. a class instance). Reject these rather
+  // than store an empty object that does not reflect what the caller passed.
+  if (serialized === '{}' && Object.keys(attributes).length > 0) {
+    throw new DBOSError('Invalid workflow attributes: object has no JSON-serializable properties.');
+  }
+}
 
 export interface WorkflowParams {
   workflowUUID?: string;
@@ -18,6 +52,8 @@ export interface WorkflowParams {
   // executor skips its dedup-error pre-recording at the parent's funcID — the wrapper records the
   // child mapping itself after attaching to the existing workflow.
   duplicationPolicy?: DuplicationPolicy;
+  // Custom key-value attributes to attach to the workflow at creation. Not inherited by child workflows.
+  workflowAttributes?: Record<string, unknown>;
 }
 
 export const DEFAULT_MAX_RECOVERY_ATTEMPTS = 100;
@@ -116,6 +152,9 @@ export interface WorkflowStatus {
   // If this workflow was started by another workflow, that workflow's ID.
   readonly parentWorkflowID?: string;
 
+  // Custom key-value attributes attached to the workflow at creation, if any.
+  readonly attributes?: Record<string, unknown>;
+
   // INTERNAL
   // Deprecated field
   readonly applicationID: string;
@@ -145,6 +184,7 @@ export interface GetWorkflowsInput {
   wasForkedFrom?: boolean; // Filter workflows that have (or have not) been forked from.
   parentWorkflowID?: string | string[]; // Get workflows started by this parent workflow ID (or any of these parent workflow IDs).
   hasParent?: boolean; // Filter workflows that have (or do not have) a parent workflow.
+  attributes?: Record<string, unknown>; // Retrieve workflows whose custom attributes contain all of these key-value pairs.
   limit?: number; // Return up to this many workflows IDs. IDs are ordered by workflow creation time.
   offset?: number; // Skip this many workflows IDs. IDs are ordered by workflow creation time.
   sortDesc?: boolean; // Sort the workflows in descending order by creation time (default ascending order).