vercel__workflow/packages/world-postgres/src/drizzle/schema.ts

import {
  type Event,
  type Hook,
  type SerializedData,
  type Step,
  StepStatusSchema,
  type Wait,
  WaitStatusSchema,
  type WorkflowRun,
  WorkflowRunStatusSchema,
} from '@workflow/world';
import { sql } from 'drizzle-orm';
import {
  boolean,
  customType,
  index,
  integer,
  /** @deprecated: use Cbor instead */
  jsonb,
  pgSchema,
  primaryKey,
  text,
  timestamp,
  uniqueIndex,
  varchar,
} from 'drizzle-orm/pg-core';
import { Cbor, type Cborized } from './cbor.js';

export const schema = pgSchema('workflow');

function mustBeMoreThanOne<T>(t: T[]) {
  return t as [T, ...T[]];
}

export const workflowRunStatus = schema.enum(
  'status',
  mustBeMoreThanOne(WorkflowRunStatusSchema.options)
);

export const stepStatus = schema.enum(
  'step_status',
  mustBeMoreThanOne(StepStatusSchema.options)
);

export const waitStatus = schema.enum(
  'wait_status',
  mustBeMoreThanOne(WaitStatusSchema.options)
);

/**
 * A mapped type that converts all properties of T to Drizzle ORM column definitions,
 * marking them as not nullable if they are not optional in T.
 */
type DrizzlishOfType<T extends object> = {
  [key in keyof T]-?: undefined extends T[key]
    ? { _: { notNull: boolean } }
    : { _: { notNull: true } };
};

/**
 * Serialization currently uses `any[]`.
 */
export type SerializedContent = any[];

export const runs = schema.table(
  'workflow_runs',
  {
    runId: varchar('id').primaryKey(),
    /** @deprecated */
    outputJson: jsonb('output').$type<SerializedContent>(),
    output: Cbor<SerializedContent>()('output_cbor'),
    deploymentId: varchar('deployment_id').notNull(),
    status: workflowRunStatus('status').notNull(),
    workflowName: varchar('name').notNull(),
    specVersion: integer('spec_version'),
    /** @deprecated */
    executionContextJson:
      jsonb('execution_context').$type<Record<string, any>>(),
    executionContext: Cbor<Record<string, any>>()('execution_context_cbor'),
    /** @deprecated */
    inputJson: jsonb('input').$type<SerializedContent>(),
    input: Cbor<SerializedContent>()('input_cbor'),
    /** @deprecated - use error instead (legacy JSON-stringified StructuredError) */
    errorJson: text('error'),
    /**
     * The thrown value from a run_failed event, serialized via the workflow
     * serialization pipeline (dehydrateRunError). Stored as a Uint8Array and
     * wrapped in CBOR for transport.
     */
    error: Cbor<SerializedData>()('error_cbor'),
    /**
     * The high-level error category (USER_ERROR, RUNTIME_ERROR, etc.) from
     * a run_failed event. Plaintext metadata for routing, so it does not
     * require decryption or hydration.
     */
    errorCode: varchar('error_code'),
    /**
     * Plaintext string-string metadata attached to the run via
     * `setAttributes()`. EXPERIMENTAL MVP: stored as JSONB to allow
     * SQL-side merge (`jsonb_set` / `jsonb_strip_nulls`) without a
     * read-modify-write cycle. Defaults to `{}` so existing rows
     * (pre-migration) read as the empty map.
     */
    attributes: jsonb('attributes')
      .$type<Record<string, string>>()
      .default({})
      .notNull(),
    /**
     * The run's X25519 public key (base64), stamped at creation by SDKs that
     * support sealed (`encp`) envelopes. Lets cross-run writers seal payloads
     * to this run without holding its symmetric key. Not secret. The private
     * scalar is re-derived on demand and never stored. Null on runs created by
     * older SDKs, which fall back to the symmetric path.
     */
    encryptionPublicKey: varchar('encryption_public_key'),
    createdAt: timestamp('created_at').defaultNow().notNull(),
    updatedAt: timestamp('updated_at')
      .defaultNow()
      .$onUpdateFn(() => new Date())
      .notNull(),
    completedAt: timestamp('completed_at'),
    startedAt: timestamp('started_at'),
    expiredAt: timestamp('expired_at'),
  } satisfies DrizzlishOfType<
    Cborized<
      Omit<WorkflowRun, 'input'> & { input?: unknown },
      'input' | 'output' | 'executionContext' | 'error'
    >
  >,
  (tb) => [index().on(tb.workflowName), index().on(tb.status)]
);

export const events = schema.table(
  'workflow_events',
  {
    eventId: varchar('id').notNull(),
    eventType: varchar('type').$type<Event['eventType']>().notNull(),
    correlationId: varchar('correlation_id'),
    createdAt: timestamp('created_at').defaultNow().notNull(),
    runId: varchar('run_id').notNull(),
    /** @deprecated */
    eventDataJson: jsonb('payload'),
    eventData: Cbor<unknown>()('payload_cbor'),
    specVersion: integer('spec_version'),
    // `resumeId` is omitted deliberately: world-postgres does not advertise
    // lazy hook-resume dedup and stays on the sequential single-writer path,
    // so it never needs to persist the resume idempotency key (no column, no
    // migration).
  } satisfies DrizzlishOfType<
    Cborized<
      Omit<Event, 'occurredAt' | 'resumeId'> & { eventData?: undefined },
      'eventData'
    >
  >,
  (tb) => [
    // Event ids are per-run slot positions, so `evnt_…0001` exists once per
    // run and is only unique together with the run it belongs to. Runs
    // created before slots keep globally-unique ULIDs, which this key also
    // admits.
    primaryKey({ columns: [tb.runId, tb.eventId] }),
    // No standalone index on `runId`: the primary key leads with it, so every
    // by-run lookup and range scan is served by that index already. Keeping one
    // would cost a second write per event on the table's hottest path.
    index().on(tb.correlationId),
    // Runtime-correlated one-shot events must be unique per (run, correlation).
    // Without
    // this, two concurrent invocations producing identical correlationIds
    // (e.g. the snapshot runtime's deterministic ULIDs across replays) can
    // both insert events, causing duplicate operations in the log.
    // The unique violation is caught in events.create and translated to
    // EntityConflictError, matching the runtime's expected dedup contract.
    uniqueIndex('workflow_events_entity_creation_unique')
      .on(tb.runId, tb.correlationId, tb.eventType)
      .where(
        sql`${tb.eventType} IN ('step_created', 'hook_created', 'wait_created', 'attr_set')`
      ),
  ]
);

/**
 * Which runs are slot-numbered. A row exists iff the run is, so its absence is
 * exactly the "this run predates slots, keep minting ULIDs" signal, so no scan
 * of the event log is needed to tell the two schemes apart.
 *
 * A marker, not a counter. Positions are allocated by the insert that occupies
 * them (`MAX(slot) + 1` read from the log inside the INSERT), so nothing is
 * handed out ahead of the write that uses it and a write that fails leaves the
 * position free for the next one. A counter here would instead burn a position
 * per failed write, and every such hole is permanent.
 */
export const eventSlots = schema.table('workflow_event_slots', {
  runId: varchar('run_id').primaryKey(),
});

export const steps = schema.table(
  'workflow_steps',
  {
    runId: varchar('run_id').notNull(),
    stepId: varchar('step_id').primaryKey(),
    stepName: varchar('step_name').notNull(),
    status: stepStatus('status').notNull(),
    /** @deprecated */
    inputJson: jsonb('input').$type<SerializedContent>(),
    input: Cbor<SerializedContent>()('input_cbor'),
    /** @deprecated we stream binary data */
    outputJson: jsonb('output').$type<SerializedContent>(),
    output: Cbor<SerializedContent>()('output_cbor'),
    /** @deprecated - use error instead (legacy JSON-stringified StructuredError) */
    errorJson: text('error'),
    /**
     * The thrown value from a step_failed / step_retrying event, serialized
     * via the workflow serialization pipeline (dehydrateStepError). Stored
     * as a Uint8Array and wrapped in CBOR for transport.
     */
    error: Cbor<SerializedData>()('error_cbor'),
    attempt: integer('attempt').notNull(),
    /** Maps to startedAt in Step interface */
    startedAt: timestamp('started_at'),
    completedAt: timestamp('completed_at'),
    createdAt: timestamp('created_at').defaultNow().notNull(),
    updatedAt: timestamp('updated_at')
      .defaultNow()
      .$onUpdateFn(() => new Date())
      .notNull(),
    retryAfter: timestamp('retry_after'),
    specVersion: integer('spec_version'),
  } satisfies DrizzlishOfType<
    Cborized<
      Omit<Step, 'input'> & {
        input?: unknown;
      },
      'output' | 'input' | 'error'
    >
  >,
  (tb) => [index().on(tb.runId), index().on(tb.status)]
);

export const hooks = schema.table(
  'workflow_hooks',
  {
    runId: varchar('run_id').notNull(),
    hookId: varchar('hook_id').primaryKey(),
    token: varchar('token').notNull(),
    ownerId: varchar('owner_id').notNull(),
    projectId: varchar('project_id').notNull(),
    environment: varchar('environment').notNull(),
    createdAt: timestamp('created_at').defaultNow().notNull(),
    tokenRetentionUntil: timestamp('token_retention_until', {
      withTimezone: true,
    }),
    /** @deprecated */
    metadataJson: jsonb('metadata').$type<SerializedContent>(),
    metadata: Cbor<SerializedContent>()('metadata_cbor'),
    specVersion: integer('spec_version'),
    isWebhook: boolean('is_webhook').default(true),
    isSystem: boolean('is_system').default(false),
    // Server-synthesized resume slice. Not carried by the hook_created event,
    // so this backend leaves it null; reads fall back to runs.get.
    resumeContext: Cbor<NonNullable<Hook['resumeContext']>>()('resume_context'),
    // `resumeCapabilities` is deliberately response-only (attested fresh on
    // each by-token lookup, never persisted), so it must not become a column.
  } satisfies DrizzlishOfType<
    Cborized<Omit<Hook, 'resumeCapabilities'>, 'metadata'>
  >,
  (tb) => [index().on(tb.runId), index().on(tb.token)]
);

export const waits = schema.table(
  'workflow_waits',
  {
    waitId: varchar('wait_id').primaryKey(),
    runId: varchar('run_id').notNull(),
    status: waitStatus('status').notNull(),
    resumeAt: timestamp('resume_at'),
    completedAt: timestamp('completed_at'),
    createdAt: timestamp('created_at').defaultNow().notNull(),
    updatedAt: timestamp('updated_at')
      .defaultNow()
      .$onUpdateFn(() => new Date())
      .notNull(),
    specVersion: integer('spec_version'),
  } satisfies DrizzlishOfType<Wait>,
  (tb) => [index().on(tb.runId)]
);

const bytea = customType<{ data: Buffer; notNull: false; default: false }>({
  dataType() {
    return 'bytea';
  },
});

export const streams = schema.table(
  'workflow_stream_chunks',
  {
    chunkId: varchar('id').$type<`chnk_${string}`>().notNull(),
    streamId: varchar('stream_id').notNull(),
    runId: varchar('run_id'),
    chunkData: bytea('data').notNull(),
    createdAt: timestamp('created_at').defaultNow().notNull(),
    eof: boolean('eof').notNull(),
  },
  (tb) => [
    primaryKey({ columns: [tb.streamId, tb.chunkId] }),
    index().on(tb.runId),
  ]
);