@tanstack/workflow-store-drizzle-postgres 0.0.1 → 0.0.3

package/src/store.ts

@@ -1,7 +1,13 @@
 import { LogConflictError } from '@tanstack/workflow-core'
 import { sql } from 'drizzle-orm'
+import {
+  getDrizzlePostgresWorkflowStoreSchemaStatements,
+  qualifiedTableName,
+  resolveDrizzlePostgresWorkflowStoreTables,
+} from './schema-contract'
 import type { RunState, WorkflowEvent } from '@tanstack/workflow-core'
 import type { SQL } from 'drizzle-orm'
+import type { DrizzlePostgresWorkflowStoreTables } from './schema-contract'
 import type {
   AppendEventsArgs,
   AppendEventsResult,
@@ -50,17 +56,6 @@ export interface DrizzlePostgresDatabase {
   ) => Promise<TResult>
 }
 
-export interface DrizzlePostgresWorkflowStoreTables {
-  runs: string
-  runStates: string
-  eventLocks: string
-  events: string
-  timers: string
-  signalDeliveries: string
-  schedules: string
-  scheduleBuckets: string
-}
-
 export interface DrizzlePostgresWorkflowStoreOptions {
   db: DrizzlePostgresDatabase
   schema?: string
@@ -72,31 +67,19 @@ export type DrizzlePostgresWorkflowStore = WorkflowExecutionStore &
     ensureSchema: () => Promise<void>
   }
 
-export const defaultDrizzlePostgresWorkflowStoreTables: DrizzlePostgresWorkflowStoreTables =
-  {
-    runs: 'workflow_runs',
-    runStates: 'workflow_run_states',
-    eventLocks: 'workflow_event_locks',
-    events: 'workflow_events',
-    timers: 'workflow_timers',
-    signalDeliveries: 'workflow_signal_deliveries',
-    schedules: 'workflow_schedules',
-    scheduleBuckets: 'workflow_schedule_buckets',
-  }
-
 export function createDrizzlePostgresWorkflowStore(
   options: DrizzlePostgresWorkflowStoreOptions,
 ): DrizzlePostgresWorkflowStore {
-  const tableNames = {
-    ...defaultDrizzlePostgresWorkflowStoreTables,
-    ...options.tables,
-  }
+  const tableNames = resolveDrizzlePostgresWorkflowStoreTables(options.tables)
   const tableSql = tableSqls(options.schema, tableNames)
   const db = options.db
 
   return {
     async ensureSchema() {
-      for (const statement of schemaStatements(options.schema, tableNames)) {
+      for (const statement of getDrizzlePostgresWorkflowStoreSchemaStatements({
+        schema: options.schema,
+        tables: tableNames,
+      })) {
         await db.execute(sql.raw(statement))
       }
     },
@@ -165,6 +148,7 @@ export function createDrizzlePostgresWorkflowStore(
             input,
             output,
             error,
+            awaiting,
             waiting_for,
             pending_approval,
             created_at,
@@ -178,6 +162,7 @@ export function createDrizzlePostgresWorkflowStore(
             ${encodeJson(state.input)}::jsonb,
             ${encodeJsonOrNull(state.output)}::jsonb,
             ${encodeJsonOrNull(state.error)}::jsonb,
+            ${encodeJsonOrNull(state.awaiting)}::jsonb,
             ${encodeJsonOrNull(state.waitingFor)}::jsonb,
             ${encodeJsonOrNull(state.pendingApproval)}::jsonb,
             ${state.createdAt},
@@ -190,6 +175,7 @@ export function createDrizzlePostgresWorkflowStore(
             input = excluded.input,
             output = excluded.output,
             error = excluded.error,
+            awaiting = excluded.awaiting,
             waiting_for = excluded.waiting_for,
             pending_approval = excluded.pending_approval,
             created_at = excluded.created_at,
@@ -204,6 +190,7 @@ export function createDrizzlePostgresWorkflowStore(
             input,
             output,
             error,
+            awaiting,
             waiting_for,
             pending_approval,
             wake_at,
@@ -218,6 +205,7 @@ export function createDrizzlePostgresWorkflowStore(
             ${encodeJson(state.input)}::jsonb,
             ${encodeJsonOrNull(state.output)}::jsonb,
             ${encodeJsonOrNull(state.error)}::jsonb,
+            ${encodeJsonOrNull(state.awaiting)}::jsonb,
             ${encodeJsonOrNull(state.waitingFor)}::jsonb,
             ${encodeJsonOrNull(state.pendingApproval)}::jsonb,
             ${
@@ -235,6 +223,7 @@ export function createDrizzlePostgresWorkflowStore(
             input = excluded.input,
             output = excluded.output,
             error = excluded.error,
+            awaiting = excluded.awaiting,
             waiting_for = excluded.waiting_for,
             pending_approval = excluded.pending_approval,
             wake_at = excluded.wake_at,
@@ -402,6 +391,7 @@ export function createDrizzlePostgresWorkflowStore(
         update ${tableSql.runs}
         set
           status = 'paused',
+          awaiting = ${encodeJsonOrNull(args.awaiting)}::jsonb,
           waiting_for = ${encodeJsonOrNull(args.waitingFor)}::jsonb,
           pending_approval = ${encodeJsonOrNull(args.pendingApproval)}::jsonb,
           wake_at = ${args.wakeAt ?? null},
@@ -418,6 +408,7 @@ export function createDrizzlePostgresWorkflowStore(
         set
           status = 'finished',
           output = ${encodeJson(args.output)}::jsonb,
+          awaiting = null,
           waiting_for = null,
           pending_approval = null,
           wake_at = null,
@@ -435,6 +426,7 @@ export function createDrizzlePostgresWorkflowStore(
         set
           status = 'errored',
           error = ${encodeJson(args.error)}::jsonb,
+          awaiting = null,
           waiting_for = null,
           pending_approval = null,
           wake_at = null,
@@ -524,7 +516,7 @@ export function createDrizzlePostgresWorkflowStore(
         )
         if (existingDelivery) return { kind: 'duplicate', run }
 
-        if (run.waitingFor?.signalName !== args.delivery.name) {
+        if (!isRunWaitingForSignal(run, args.delivery)) {
           return { kind: 'not-waiting', run }
         }
 
@@ -549,6 +541,7 @@ export function createDrizzlePostgresWorkflowStore(
             update ${tableSql.runs}
             set
               status = 'queued',
+              awaiting = null,
               waiting_for = null,
               pending_approval = null,
               wake_at = null,
@@ -578,7 +571,7 @@ export function createDrizzlePostgresWorkflowStore(
         )
         if (existingDelivery) return { kind: 'duplicate', run }
 
-        if (run.pendingApproval?.approvalId !== args.approval.approvalId) {
+        if (!isRunWaitingForApproval(run, args.approval)) {
           return { kind: 'not-waiting', run }
         }
 
@@ -597,6 +590,7 @@ export function createDrizzlePostgresWorkflowStore(
             update ${tableSql.runs}
             set
               status = 'queued',
+              awaiting = null,
               waiting_for = null,
               pending_approval = null,
               wake_at = null,
@@ -650,12 +644,26 @@ export function createDrizzlePostgresWorkflowStore(
       const schedules = await queryRows<ScheduleRow>(
         db,
         sql`
-          select *
-          from ${tableSql.schedules}
-          where enabled = true
-            and next_fire_at is not null
-            and next_fire_at <= ${args.now}
-          order by next_fire_at asc, schedule_id asc
+          select schedule.*
+          from ${tableSql.schedules} schedule
+          left join ${tableSql.scheduleBuckets} bucket
+            on bucket.schedule_id = schedule.schedule_id
+            and bucket.bucket_id = schedule.next_fire_at::text
+          where schedule.enabled = true
+            and schedule.next_fire_at is not null
+            and schedule.next_fire_at <= ${args.now}
+            and (
+              bucket.schedule_id is null
+              or (
+                bucket.status <> 'started'
+                and (
+                  bucket.lease_owner is null
+                  or bucket.lease_owner = ${args.leaseOwner}
+                  or bucket.lease_expires_at <= ${args.now}
+                )
+              )
+            )
+          order by schedule.next_fire_at asc, schedule.schedule_id asc
           limit ${args.limit}
         `,
       )
@@ -810,6 +818,7 @@ interface RunRow {
   input: unknown
   output: unknown
   error: unknown
+  awaiting: unknown
   waiting_for: unknown
   pending_approval: unknown
   wake_at: number | string | null
@@ -827,6 +836,7 @@ interface RunStateRow {
   input: unknown
   output: unknown
   error: unknown
+  awaiting: unknown
   waiting_for: unknown
   pending_approval: unknown
   created_at: number | string
@@ -890,122 +900,6 @@ function tableSqls(
   }
 }
 
-function schemaStatements(
-  schema: string | undefined,
-  tables: DrizzlePostgresWorkflowStoreTables,
-): Array<string> {
-  const runs = qualifiedTableName(schema, tables.runs)
-  const runStates = qualifiedTableName(schema, tables.runStates)
-  const eventLocks = qualifiedTableName(schema, tables.eventLocks)
-  const events = qualifiedTableName(schema, tables.events)
-  const timers = qualifiedTableName(schema, tables.timers)
-  const signalDeliveries = qualifiedTableName(schema, tables.signalDeliveries)
-  const schedules = qualifiedTableName(schema, tables.schedules)
-  const scheduleBuckets = qualifiedTableName(schema, tables.scheduleBuckets)
-
-  return [
-    ...(schema ? [`create schema if not exists ${quoteIdent(schema)}`] : []),
-    `create table if not exists ${runs} (
-      run_id text primary key,
-      workflow_id text not null,
-      workflow_version text,
-      status text not null,
-      input jsonb not null,
-      output jsonb,
-      error jsonb,
-      waiting_for jsonb,
-      pending_approval jsonb,
-      wake_at bigint,
-      lease_owner text,
-      lease_expires_at bigint,
-      created_at bigint not null,
-      updated_at bigint not null
-    )`,
-    `create index if not exists ${quoteIdent(`${tables.runs}_status_idx`)}
-      on ${runs} (status, updated_at)`,
-    `create index if not exists ${quoteIdent(`${tables.runs}_lease_idx`)}
-      on ${runs} (status, lease_expires_at)`,
-    `create table if not exists ${runStates} (
-      run_id text primary key,
-      workflow_id text not null,
-      workflow_version text,
-      status text not null,
-      input jsonb not null,
-      output jsonb,
-      error jsonb,
-      waiting_for jsonb,
-      pending_approval jsonb,
-      created_at bigint not null,
-      updated_at bigint not null
-    )`,
-    `create table if not exists ${eventLocks} (
-      run_id text primary key,
-      created_at bigint not null
-    )`,
-    `create table if not exists ${events} (
-      run_id text not null,
-      event_index integer not null,
-      event_type text not null,
-      step_id text,
-      event jsonb not null,
-      created_at bigint not null,
-      primary key (run_id, event_index)
-    )`,
-    `create index if not exists ${quoteIdent(`${tables.events}_type_idx`)}
-      on ${events} (run_id, event_type)`,
-    `create table if not exists ${timers} (
-      run_id text not null,
-      signal_id text not null,
-      workflow_id text not null,
-      workflow_version text,
-      wake_at bigint not null,
-      lease_owner text,
-      lease_expires_at bigint,
-      primary key (run_id, signal_id)
-    )`,
-    `create index if not exists ${quoteIdent(`${tables.timers}_due_idx`)}
-      on ${timers} (wake_at, lease_expires_at)`,
-    `create table if not exists ${signalDeliveries} (
-      run_id text not null,
-      signal_id text not null,
-      created_at bigint not null,
-      primary key (run_id, signal_id)
-    )`,
-    `create table if not exists ${schedules} (
-      schedule_id text primary key,
-      workflow_id text not null,
-      workflow_version text,
-      schedule jsonb not null,
-      overlap_policy text not null,
-      input jsonb,
-      next_fire_at bigint,
-      enabled boolean not null,
-      updated_at bigint not null
-    )`,
-    `create index if not exists ${quoteIdent(`${tables.schedules}_due_idx`)}
-      on ${schedules} (enabled, next_fire_at)`,
-    `create table if not exists ${scheduleBuckets} (
-      schedule_id text not null,
-      bucket_id text not null,
-      workflow_id text not null,
-      workflow_version text,
-      run_id text not null,
-      fire_at bigint not null,
-      input jsonb,
-      overlap_policy text not null,
-      status text not null,
-      lease_owner text,
-      lease_expires_at bigint,
-      started_at bigint,
-      primary key (schedule_id, bucket_id)
-    )`,
-    `create index if not exists ${quoteIdent(
-      `${tables.scheduleBuckets}_lease_idx`,
-    )}
-      on ${scheduleBuckets} (status, fire_at, lease_expires_at)`,
-  ]
-}
-
 async function loadRunById(
   db: DrizzlePostgresDatabase,
   tables: TableSqls,
@@ -1144,6 +1038,7 @@ function runFromRow(row: RunRow): WorkflowExecution {
     input: decodeJson(row.input),
     output: decodeJsonOrUndefined(row.output),
     error: decodeJsonOrUndefined(row.error),
+    awaiting: decodeJsonOrUndefined(row.awaiting),
     waitingFor: decodeJsonOrUndefined(row.waiting_for),
     pendingApproval: decodeJsonOrUndefined(row.pending_approval),
     wakeAt: numberOrUndefined(row.wake_at),
@@ -1168,6 +1063,7 @@ function runStateFromRow(row: RunStateRow): RunState {
     input: decodeJson(row.input),
     output: decodeJsonOrUndefined(row.output),
     error: decodeJsonOrUndefined(row.error),
+    awaiting: decodeJsonOrUndefined(row.awaiting),
     waitingFor: decodeJsonOrUndefined(row.waiting_for),
     pendingApproval: decodeJsonOrUndefined(row.pending_approval),
     createdAt: Number(row.created_at),
@@ -1196,6 +1092,54 @@ function timerFromRow(row: TimerRow): TimerWakeup {
   }
 }
 
+function isRunWaitingForSignal(
+  run: WorkflowExecution,
+  delivery: DeliverSignalArgs['delivery'],
+) {
+  return (
+    signalAwaitableMatches(run.waitingFor, delivery) ||
+    run.awaiting?.some(
+      (awaitable) =>
+        awaitable.type === 'signal' &&
+        signalAwaitableMatches(awaitable, delivery),
+    ) === true
+  )
+}
+
+function signalAwaitableMatches(
+  awaitable:
+    | NonNullable<WorkflowExecution['waitingFor']>
+    | Extract<
+        NonNullable<WorkflowExecution['awaiting']>[number],
+        {
+          type: 'signal'
+        }
+      >
+    | undefined,
+  delivery: DeliverSignalArgs['delivery'],
+) {
+  return (
+    awaitable?.signalName === delivery.name &&
+    (delivery.stepId === undefined ||
+      awaitable.stepId === undefined ||
+      awaitable.stepId === delivery.stepId)
+  )
+}
+
+function isRunWaitingForApproval(
+  run: WorkflowExecution,
+  approval: DeliverApprovalArgs['approval'],
+) {
+  return (
+    run.pendingApproval?.approvalId === approval.approvalId ||
+    run.awaiting?.some(
+      (awaitable) =>
+        awaitable.type === 'approval' &&
+        awaitable.approvalId === approval.approvalId,
+    ) === true
+  )
+}
+
 function scheduleFromRow(row: ScheduleRow) {
   return {
     scheduleId: row.schedule_id satisfies ScheduleId,
@@ -1227,6 +1171,7 @@ function toRunSummary(row: RunRow): RunSummary {
     workflowId: run.workflowId,
     workflowVersion: run.workflowVersion,
     status: run.status,
+    awaiting: run.awaiting,
     waitingFor: run.waitingFor,
     pendingApproval: run.pendingApproval,
     wakeAt: run.wakeAt,
@@ -1266,13 +1211,3 @@ function numberOrUndefined(value: number | string | null) {
 function getStepId(event: WorkflowEvent) {
   return 'stepId' in event ? event.stepId : undefined
 }
-
-function qualifiedTableName(schema: string | undefined, table: string) {
-  return schema
-    ? `${quoteIdent(schema)}.${quoteIdent(table)}`
-    : quoteIdent(table)
-}
-
-function quoteIdent(identifier: string) {
-  return `"${identifier.replaceAll('"', '""')}"`
-}

package/src/tables.ts

@@ -0,0 +1,149 @@
+import {
+  bigint,
+  boolean,
+  index,
+  integer,
+  jsonb,
+  pgTable,
+  primaryKey,
+  text,
+} from 'drizzle-orm/pg-core'
+
+export const workflowSchemaMigrations = pgTable('workflow_schema_migrations', {
+  migrationId: text('migration_id').primaryKey(),
+  packageName: text('package_name').notNull(),
+  packageVersion: text('package_version'),
+  appliedAt: bigint('applied_at', { mode: 'number' }).notNull(),
+})
+
+export const workflowRuns = pgTable(
+  'workflow_runs',
+  {
+    runId: text('run_id').primaryKey(),
+    workflowId: text('workflow_id').notNull(),
+    workflowVersion: text('workflow_version'),
+    status: text('status').notNull(),
+    input: jsonb('input').notNull(),
+    output: jsonb('output'),
+    error: jsonb('error'),
+    awaiting: jsonb('awaiting'),
+    waitingFor: jsonb('waiting_for'),
+    pendingApproval: jsonb('pending_approval'),
+    wakeAt: bigint('wake_at', { mode: 'number' }),
+    leaseOwner: text('lease_owner'),
+    leaseExpiresAt: bigint('lease_expires_at', { mode: 'number' }),
+    createdAt: bigint('created_at', { mode: 'number' }).notNull(),
+    updatedAt: bigint('updated_at', { mode: 'number' }).notNull(),
+  },
+  (table) => [
+    index('workflow_runs_status_idx').on(table.status, table.updatedAt),
+    index('workflow_runs_lease_idx').on(table.status, table.leaseExpiresAt),
+  ],
+)
+
+export const workflowRunStates = pgTable('workflow_run_states', {
+  runId: text('run_id').primaryKey(),
+  workflowId: text('workflow_id').notNull(),
+  workflowVersion: text('workflow_version'),
+  status: text('status').notNull(),
+  input: jsonb('input').notNull(),
+  output: jsonb('output'),
+  error: jsonb('error'),
+  awaiting: jsonb('awaiting'),
+  waitingFor: jsonb('waiting_for'),
+  pendingApproval: jsonb('pending_approval'),
+  createdAt: bigint('created_at', { mode: 'number' }).notNull(),
+  updatedAt: bigint('updated_at', { mode: 'number' }).notNull(),
+})
+
+export const workflowEventLocks = pgTable('workflow_event_locks', {
+  runId: text('run_id').primaryKey(),
+  createdAt: bigint('created_at', { mode: 'number' }).notNull(),
+})
+
+export const workflowEvents = pgTable(
+  'workflow_events',
+  {
+    runId: text('run_id').notNull(),
+    eventIndex: integer('event_index').notNull(),
+    eventType: text('event_type').notNull(),
+    stepId: text('step_id'),
+    event: jsonb('event').notNull(),
+    createdAt: bigint('created_at', { mode: 'number' }).notNull(),
+  },
+  (table) => [
+    primaryKey({ columns: [table.runId, table.eventIndex] }),
+    index('workflow_events_type_idx').on(table.runId, table.eventType),
+  ],
+)
+
+export const workflowTimers = pgTable(
+  'workflow_timers',
+  {
+    runId: text('run_id').notNull(),
+    signalId: text('signal_id').notNull(),
+    workflowId: text('workflow_id').notNull(),
+    workflowVersion: text('workflow_version'),
+    wakeAt: bigint('wake_at', { mode: 'number' }).notNull(),
+    leaseOwner: text('lease_owner'),
+    leaseExpiresAt: bigint('lease_expires_at', { mode: 'number' }),
+  },
+  (table) => [
+    primaryKey({ columns: [table.runId, table.signalId] }),
+    index('workflow_timers_due_idx').on(table.wakeAt, table.leaseExpiresAt),
+  ],
+)
+
+export const workflowSignalDeliveries = pgTable(
+  'workflow_signal_deliveries',
+  {
+    runId: text('run_id').notNull(),
+    signalId: text('signal_id').notNull(),
+    createdAt: bigint('created_at', { mode: 'number' }).notNull(),
+  },
+  (table) => [primaryKey({ columns: [table.runId, table.signalId] })],
+)
+
+export const workflowSchedules = pgTable(
+  'workflow_schedules',
+  {
+    scheduleId: text('schedule_id').primaryKey(),
+    workflowId: text('workflow_id').notNull(),
+    workflowVersion: text('workflow_version'),
+    schedule: jsonb('schedule').notNull(),
+    overlapPolicy: text('overlap_policy').notNull(),
+    input: jsonb('input'),
+    nextFireAt: bigint('next_fire_at', { mode: 'number' }),
+    enabled: boolean('enabled').notNull(),
+    updatedAt: bigint('updated_at', { mode: 'number' }).notNull(),
+  },
+  (table) => [
+    index('workflow_schedules_due_idx').on(table.enabled, table.nextFireAt),
+  ],
+)
+
+export const workflowScheduleBuckets = pgTable(
+  'workflow_schedule_buckets',
+  {
+    scheduleId: text('schedule_id').notNull(),
+    bucketId: text('bucket_id').notNull(),
+    workflowId: text('workflow_id').notNull(),
+    workflowVersion: text('workflow_version'),
+    runId: text('run_id').notNull(),
+    fireAt: bigint('fire_at', { mode: 'number' }).notNull(),
+    input: jsonb('input'),
+    overlapPolicy: text('overlap_policy').notNull(),
+    status: text('status').notNull(),
+    leaseOwner: text('lease_owner'),
+    leaseExpiresAt: bigint('lease_expires_at', { mode: 'number' }),
+    startedAt: bigint('started_at', { mode: 'number' }),
+  },
+  (table) => [
+    primaryKey({ columns: [table.scheduleId, table.bucketId] }),
+    index('workflow_schedule_buckets_lease_idx').on(
+      table.status,
+      table.fireAt,
+      table.leaseExpiresAt,
+    ),
+  ],
+)