Refactor/extra engine types

src/ast-parser.test.ts

@@ -1,11 +1,12 @@
 import type pg from 'pg';
 import type { PgBoss } from 'pg-boss';
 import { beforeAll, describe, expect, it } from 'vitest';
+import { z } from 'zod';
 import { workflow } from './definition';
 import { WorkflowEngine } from './engine';
 import { getBoss } from './tests/pgboss';
 import { createTestDatabase } from './tests/test-db';
-import { StepType } from './types';
+import { type StepBaseContext, StepType } from './types';
 
 let testBoss: PgBoss;
 let testPool: pg.Pool;
@@ -62,16 +63,20 @@ describe('AST Parser for Workflow Steps', () => {
   });
 
   it('should detect conditional steps', async () => {
-    const conditionalWorkflow = workflow('conditional-workflow', async ({ step, input }) => {
-      await step.run('step-1', async () => 'result-1');
+    const conditionalWorkflow = workflow(
+      'conditional-workflow',
+      async ({ step, input }) => {
+        await step.run('step-1', async () => 'result-1');
 
-      if (input.condition) {
-        await step.run('conditional-step', async () => 'conditional-result');
-      }
+        if (input.condition) {
+          await step.run('conditional-step', async () => 'conditional-result');
+        }
 
-      await step.run('step-3', async () => 'result-3');
-      return 'completed';
-    });
+        await step.run('step-3', async () => 'result-3');
+        return 'completed';
+      },
+      { inputSchema: z.object({ condition: z.boolean() }) },
+    );
 
     const engine = new WorkflowEngine({ pool: testPool, boss: testBoss });
     await engine.registerWorkflow(conditionalWorkflow);
@@ -199,53 +204,96 @@ describe('AST Parser for Workflow Steps', () => {
 
     const engine = new WorkflowEngine({ pool: testPool, boss: testBoss });
     await engine.registerWorkflow(mixedStepWorkflow);
+
+    expect(engine.workflows.get('mixed-step-workflow')?.steps).toEqual([
+      { id: 'step-1', type: StepType.RUN, conditional: false, loop: false, isDynamic: false },
+      { id: 'step-2', type: StepType.WAIT_FOR, conditional: false, loop: false, isDynamic: false },
+      { id: 'step-3', type: StepType.PAUSE, conditional: false, loop: false, isDynamic: false },
+      { id: 'step-4', type: StepType.RUN, conditional: false, loop: false, isDynamic: false },
+    ]);
   });
 
   it('should handle nested conditionals and loops', async () => {
-    const nestedWorkflow = workflow('nested-workflow', async ({ step, input }) => {
-      await step.run('start', async () => 'started');
-
-      for (let i = 0; i < input.outerCount; i++) {
-        if (i % 2 === 0) {
-          await step.run(`even-${i}`, async () => `even-result-${i}`);
-
-          for (let j = 0; j < 2; j++) {
-            await step.run(`nested-${i}-${j}`, async () => `nested-result-${i}-${j}`);
+    const nestedWorkflow = workflow(
+      'nested-workflow',
+      async ({ step, input }) => {
+        await step.run('start', async () => 'started');
+
+        for (let i = 0; i < input.outerCount; i++) {
+          if (i % 2 === 0) {
+            await step.run(`even-${i}`, async () => `even-result-${i}`);
+
+            for (let j = 0; j < 2; j++) {
+              await step.run(`nested-${i}-${j}`, async () => `nested-result-${i}-${j}`);
+            }
+          } else {
+            await step.run(`odd-${i}`, async () => `odd-result-${i}`);
           }
-        } else {
-          await step.run(`odd-${i}`, async () => `odd-result-${i}`);
         }
-      }
 
-      await step.run('end', async () => 'ended');
-      return 'completed';
-    });
+        await step.run('end', async () => 'ended');
+        return 'completed';
+      },
+      { inputSchema: z.object({ outerCount: z.number() }) },
+    );
 
     const engine = new WorkflowEngine({ pool: testPool, boss: testBoss });
     await engine.registerWorkflow(nestedWorkflow);
+
+    expect(engine.workflows.get('nested-workflow')?.steps).toEqual([
+      { id: 'start', type: StepType.RUN, conditional: false, loop: false, isDynamic: false },
+      { id: `even-\${...}`, type: StepType.RUN, conditional: true, loop: true, isDynamic: true },
+      {
+        id: `nested-\${...}-\${...}`,
+        type: StepType.RUN,
+        conditional: true,
+        loop: true,
+        isDynamic: true,
+      },
+      { id: `odd-\${...}`, type: StepType.RUN, conditional: true, loop: true, isDynamic: true },
+      { id: 'end', type: StepType.RUN, conditional: false, loop: false, isDynamic: false },
+    ]);
   });
 
   it('should handle switch statements', async () => {
-    const switchWorkflow = workflow('switch-workflow', async ({ step, input }) => {
-      await step.run('step-1', async () => 'result-1');
-
-      switch (input.type) {
-        case 'A':
-          await step.run('handle-a', async () => 'handled-a');
-          break;
-        case 'B':
-          await step.run('handle-b', async () => 'handled-b');
-          break;
-        default:
-          await step.run('handle-default', async () => 'handled-default');
-      }
+    const switchWorkflow = workflow(
+      'switch-workflow',
+      async ({ step, input }) => {
+        await step.run('step-1', async () => 'result-1');
+
+        switch (input.type) {
+          case 'A':
+            await step.run('handle-a', async () => 'handled-a');
+            break;
+          case 'B':
+            await step.run('handle-b', async () => 'handled-b');
+            break;
+          default:
+            await step.run('handle-default', async () => 'handled-default');
+        }
 
-      await step.run('step-3', async () => 'result-3');
-      return 'completed';
-    });
+        await step.run('step-3', async () => 'result-3');
+        return 'completed';
+      },
+      { inputSchema: z.object({ type: z.string() }) },
+    );
 
     const engine = new WorkflowEngine({ pool: testPool, boss: testBoss });
     await engine.registerWorkflow(switchWorkflow);
+
+    expect(engine.workflows.get('switch-workflow')?.steps).toEqual([
+      { id: 'step-1', type: StepType.RUN, conditional: false, loop: false, isDynamic: false },
+      { id: 'handle-a', type: StepType.RUN, conditional: true, loop: false, isDynamic: false },
+      { id: 'handle-b', type: StepType.RUN, conditional: true, loop: false, isDynamic: false },
+      {
+        id: 'handle-default',
+        type: StepType.RUN,
+        conditional: true,
+        loop: false,
+        isDynamic: false,
+      },
+      { id: 'step-3', type: StepType.RUN, conditional: false, loop: false, isDynamic: false },
+    ]);
   });
 
   it('should throw error for duplicate static step IDs', async () => {
@@ -281,4 +329,22 @@ describe('AST Parser for Workflow Steps', () => {
       "Duplicate step ID detected: 'process-item'. Step IDs must be unique within a workflow.",
     );
   });
+
+  it('does not parse the steps of externally defined workflow functions', async () => {
+    const workflowHandler = async (step: StepBaseContext, _input: unknown) => {
+      await step.run('process-item', async () => 'result-1');
+    };
+
+    const testWorkflow = workflow(
+      'parsing-will-not-work-at-this-workflow',
+      async ({ step, input }) => {
+        workflowHandler(step, input);
+      },
+    );
+
+    const engine = new WorkflowEngine({ pool: testPool, boss: testBoss });
+    await engine.registerWorkflow(testWorkflow);
+
+    expect(engine.workflows.get('parsing-will-not-work-at-this-workflow')?.steps).toEqual([]);
+  });
 });

src/ast-parser.ts

@@ -1,6 +1,6 @@
 import * as ts from 'typescript';
 import type { StepInternalDefinition, WorkflowContext } from './types';
-import { StepType } from './types';
+import { STEP_BASE_METHOD_TYPES } from './types';
 
 type ParseWorkflowHandlerReturnType = {
   steps: StepInternalDefinition[];
@@ -73,21 +73,11 @@ export function parseWorkflowHandler(
       const objectName = propertyAccess.expression.getText(sourceFile);
       const methodName = propertyAccess.name.text;
 
-      if (
-        objectName === 'step' &&
-        (methodName === 'run' ||
-          methodName === 'waitFor' ||
-          methodName === 'pause' ||
-          methodName === 'waitUntil' ||
-          methodName === 'delay' ||
-          methodName === 'sleep' ||
-          methodName === 'poll' ||
-          methodName === 'invokeChildWorkflow')
-      ) {
+      const stepType = objectName === 'step' ? STEP_BASE_METHOD_TYPES.get(methodName) : undefined;
+      if (stepType !== undefined) {
         const firstArg = node.arguments[0];
         if (firstArg) {
           const { id, isDynamic } = extractStepId(firstArg);
-          const stepType = methodName === 'sleep' ? StepType.DELAY : (methodName as StepType);
 
           const stepDefinition: StepInternalDefinition = {
             id,

src/db/migration.test.ts

@@ -0,0 +1,75 @@
+import type pg from 'pg';
+import type { Db } from 'pg-boss';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+import { closeTestDatabase, createTestDatabase } from '../tests/test-db';
+import { runMigrations } from './migration';
+
+const CURRENT_SCHEMA_VERSION = 5;
+
+describe('runMigrations', () => {
+  let pool: pg.Pool;
+  let db: Db;
+
+  beforeEach(async () => {
+    pool = await createTestDatabase();
+    db = {
+      executeSql: (text: string, values?: unknown[]) =>
+        pool.query(text, values) as Promise<{ rows: unknown[] }>,
+    };
+  });
+
+  afterEach(async () => {
+    await closeTestDatabase();
+  });
+
+  const tableExists = async (name: string): Promise<boolean> => {
+    const result = await db.executeSql(
+      `SELECT 1 FROM information_schema.tables
+       WHERE table_schema = current_schema() AND table_name = $1 LIMIT 1`,
+      [name],
+    );
+    return result.rows.length > 0;
+  };
+
+  const schemaVersion = async (): Promise<number> => {
+    const result = await db.executeSql('SELECT version FROM workflow_schema_version LIMIT 1', []);
+    return (result.rows[0] as { version: number }).version;
+  };
+
+  it('migrates a fresh database to the current schema version', async () => {
+    await runMigrations(db);
+
+    expect(await tableExists('workflow_runs')).toBe(true);
+    expect(await tableExists('workflow_schema_version')).toBe(true);
+    expect(await schemaVersion()).toBe(CURRENT_SCHEMA_VERSION);
+  });
+
+  it('is idempotent when run repeatedly', async () => {
+    await runMigrations(db);
+    await runMigrations(db); // second run hits the fast path, must not throw
+    await runMigrations(db);
+
+    expect(await schemaVersion()).toBe(CURRENT_SCHEMA_VERSION);
+  });
+
+  it('throws when a foreign workflow_runs table already exists', async () => {
+    // Simulate a consumer who already has their own unrelated workflow_runs table.
+    await db.executeSql('CREATE TABLE workflow_runs (id integer PRIMARY KEY, my_col text)', []);
+
+    await expect(runMigrations(db)).rejects.toThrow(
+      /already exists in this schema but was not created by pg-workflows/,
+    );
+
+    // The guard must fire before any DDL: no version table, and the foreign
+    // table must be left untouched (still has its original column).
+    expect(await tableExists('workflow_schema_version')).toBe(false);
+    const cols = await db.executeSql(
+      `SELECT column_name FROM information_schema.columns
+       WHERE table_schema = current_schema() AND table_name = 'workflow_runs'`,
+      [],
+    );
+    const columnNames = cols.rows.map((r) => (r as { column_name: string }).column_name);
+    expect(columnNames).toContain('my_col');
+    expect(columnNames).not.toContain('idempotency_key'); // proves no ALTER ran
+  });
+});

src/db/migration.ts

@@ -24,6 +24,38 @@ export async function runMigrations(db: Db): Promise<void> {
   const commands: string[] = [];
 
   if (currentVersion < 1) {
+    // On first run, we run a check to see if there is an existing foreign
+    // `workflow_runs` table. To decide on that, we check if the `workflow_runs`
+    // exists but a `workflow_schema_version` table doesn't. If that is true
+    // then `workflow_runs` is foreign and we should fail lest we start
+    // adding rows there, corrupting it. We also do the reverse check.
+    // We are basing this off the fact that our transaction creates both or no tables
+    // This is not airtight, but serves as an extra safety check for now.
+    const existing = await db.executeSql(
+      `SELECT
+         to_regclass('workflow_runs') IS NOT NULL AS has_runs_table,
+         to_regclass('workflow_schema_version') IS NOT NULL AS has_version_table`,
+      [],
+    );
+
+    const row = existing.rows[0] as
+      | { has_runs_table: boolean; has_version_table: boolean }
+      | undefined;
+
+    if (row?.has_runs_table && !row.has_version_table) {
+      throw new Error(
+        `pg-workflows: a "workflow_runs" table already exists in this schema but was not ` +
+          `created by pg-workflows. Point the workflow engine at a dedicated schema/database.`,
+      );
+    }
+
+    if (!row?.has_runs_table && row?.has_version_table) {
+      throw new Error(
+        `pg-workflows: a "workflow_schema_version" table already exists in this schema but was not ` +
+          `created by pg-workflows. Point the workflow engine at a dedicated schema/database.`,
+      );
+    }
+
     commands.push(`
       CREATE TABLE IF NOT EXISTS workflow_runs (
         id varchar(32) PRIMARY KEY NOT NULL,