@sonamu-kit/tasks 0.1.0 → 0.1.2

package/src/database/backend.ts

@@ -1,3 +1,4 @@
+import { getLogger } from "@logtape/logtape";
 import { camelize } from "inflection";
 import knex, { type Knex } from "knex";
 import {
@@ -36,22 +37,82 @@ interface BackendPostgresOptions {
   usePubSub?: boolean;
 }
 
+const logger = getLogger(["sonamu", "internal", "tasks"]);
+
 /**
  * Manages a connection to a Postgres database for workflow operations.
  */
 export class BackendPostgres implements Backend {
-  private knex: Knex;
+  private config: Knex.Config;
   private namespaceId: string;
   private usePubSub: boolean;
   private pubsub: PostgresPubSub | null = null;
+  private initialized: boolean = false;
+  private runMigrations: boolean;
+
+  private _knex: Knex | null = null;
+  private get knex(): Knex {
+    if (!this._knex) {
+      this._knex = knex(this.config);
+    }
+
+    return this._knex;
+  }
+
+  constructor(config: Knex.Config, options?: BackendPostgresOptions) {
+    this.config = {
+      ...config,
+      postProcessResponse: (result, _queryContext) => {
+        if (result === null || result === undefined) {
+          return result;
+        }
+
+        if (config?.postProcessResponse) {
+          result = config.postProcessResponse(result, _queryContext);
+        }
+
+        const camelizeRow = (row: Record<string, unknown>) =>
+          Object.fromEntries(
+            Object.entries(row).map(([key, value]) => [camelize(key, true), value]),
+          );
+
+        if (Array.isArray(result)) {
+          return result.map(camelizeRow);
+        }
+
+        return camelizeRow(result);
+      },
+    };
+
+    const { namespaceId, usePubSub, runMigrations } = {
+      namespaceId: DEFAULT_NAMESPACE_ID,
+      usePubSub: true,
+      runMigrations: true,
+      ...options,
+    };
 
-  private constructor(knex: Knex, namespaceId: string, usePubSub: boolean) {
-    this.knex = knex;
     this.namespaceId = namespaceId;
     this.usePubSub = usePubSub;
+    this.runMigrations = runMigrations;
+  }
+
+  async initialize() {
+    if (this.initialized) {
+      return;
+    }
+
+    if (this.runMigrations) {
+      await migrate(this.config, DEFAULT_SCHEMA);
+    }
+
+    this.initialized = true;
   }
 
   async subscribe(callback: OnSubscribed) {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     if (!this.usePubSub) {
       return;
     }
@@ -64,6 +125,10 @@ export class BackendPostgres implements Backend {
   }
 
   async publish(payload?: string): Promise<void> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     if (!this.usePubSub) {
       return;
     }
@@ -75,56 +140,21 @@ export class BackendPostgres implements Backend {
     );
   }
 
-  /**
-   * Create and initialize a new BackendPostgres instance. This will
-   * automatically run migrations on startup unless `runMigrations` is set to
-   * false.
-   */
-  static async connect(
-    dbConf: Knex.Config,
-    options?: BackendPostgresOptions,
-  ): Promise<BackendPostgres> {
-    const postProcessResponse: Knex.Config["postProcessResponse"] = (result, _queryContext) => {
-      if (result === null || result === undefined) {
-        return result;
-      }
-
-      if (dbConf?.postProcessResponse) {
-        result = dbConf.postProcessResponse(result, _queryContext);
-      }
-
-      const camelizeRow = (row: Record<string, unknown>) =>
-        Object.fromEntries(Object.entries(row).map(([key, value]) => [camelize(key, true), value]));
-
-      if (Array.isArray(result)) {
-        return result.map(camelizeRow);
-      }
-
-      return camelizeRow(result);
-    };
-
-    const { namespaceId, runMigrations, usePubSub } = {
-      namespaceId: DEFAULT_NAMESPACE_ID,
-      runMigrations: true,
-      usePubSub: true,
-      ...options,
-    };
-
-    const knexInstance = knex({ ...dbConf, postProcessResponse });
-    if (runMigrations) {
-      await migrate(knexInstance, DEFAULT_SCHEMA);
+  async stop(): Promise<void> {
+    if (!this.initialized) {
+      return;
     }
 
-    return new BackendPostgres(knexInstance, namespaceId, usePubSub);
-  }
-
-  async stop(): Promise<void> {
     await this.pubsub?.destroy();
     this.pubsub = null;
     await this.knex.destroy();
   }
 
   async createWorkflowRun(params: CreateWorkflowRunParams): Promise<WorkflowRun> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     const qb = this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("workflow_runs")
@@ -148,6 +178,7 @@ export class BackendPostgres implements Backend {
 
     const workflowRun = await qb;
     if (!workflowRun[0]) {
+      logger.error("Failed to create workflow run: {params}", { params });
       throw new Error("Failed to create workflow run");
     }
 
@@ -155,6 +186,10 @@ export class BackendPostgres implements Backend {
   }
 
   async getWorkflowRun(params: GetWorkflowRunParams): Promise<WorkflowRun | null> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     const workflowRun = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("workflow_runs")
@@ -189,6 +224,10 @@ export class BackendPostgres implements Backend {
   }
 
   async listWorkflowRuns(params: ListWorkflowRunsParams): Promise<PaginatedResponse<WorkflowRun>> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     const limit = params.limit ?? DEFAULT_PAGINATION_PAGE_SIZE;
     const { after, before } = params;
 
@@ -232,6 +271,10 @@ export class BackendPostgres implements Backend {
   }
 
   async claimWorkflowRun(params: ClaimWorkflowRunParams): Promise<WorkflowRun | null> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     const claimed = await this.knex
       .with("expired", (qb) =>
         qb
@@ -288,6 +331,10 @@ export class BackendPostgres implements Backend {
   }
 
   async extendWorkflowRunLease(params: ExtendWorkflowRunLeaseParams): Promise<WorkflowRun> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     const [updated] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("workflow_runs")
@@ -302,6 +349,7 @@ export class BackendPostgres implements Backend {
       .returning("*");
 
     if (!updated) {
+      logger.error("Failed to extend lease for workflow run: {params}", { params });
       throw new Error("Failed to extend lease for workflow run");
     }
 
@@ -309,6 +357,10 @@ export class BackendPostgres implements Backend {
   }
 
   async sleepWorkflowRun(params: SleepWorkflowRunParams): Promise<WorkflowRun> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     // 'succeeded' status is deprecated
     const [updated] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
@@ -326,6 +378,7 @@ export class BackendPostgres implements Backend {
       .returning("*");
 
     if (!updated) {
+      logger.error("Failed to sleep workflow run: {params}", { params });
       throw new Error("Failed to sleep workflow run");
     }
 
@@ -333,6 +386,10 @@ export class BackendPostgres implements Backend {
   }
 
   async completeWorkflowRun(params: CompleteWorkflowRunParams): Promise<WorkflowRun> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     const [updated] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("workflow_runs")
@@ -352,6 +409,7 @@ export class BackendPostgres implements Backend {
       .returning("*");
 
     if (!updated) {
+      logger.error("Failed to complete workflow run: {params}", { params });
       throw new Error("Failed to complete workflow run");
     }
 
@@ -359,6 +417,10 @@ export class BackendPostgres implements Backend {
   }
 
   async failWorkflowRun(params: FailWorkflowRunParams): Promise<WorkflowRun> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     const { workflowRunId, error } = params;
     const { initialIntervalMs, backoffCoefficient, maximumIntervalMs } = DEFAULT_RETRY_POLICY;
 
@@ -396,6 +458,7 @@ export class BackendPostgres implements Backend {
       .returning("*");
 
     if (!updated) {
+      logger.error("Failed to mark workflow run failed: {params}", { params });
       throw new Error("Failed to mark workflow run failed");
     }
 
@@ -403,6 +466,10 @@ export class BackendPostgres implements Backend {
   }
 
   async cancelWorkflowRun(params: CancelWorkflowRunParams): Promise<WorkflowRun> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     const [updated] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("workflow_runs")
@@ -435,11 +502,16 @@ export class BackendPostgres implements Backend {
       // throw error for completed/failed workflows
       // 'succeeded' status is deprecated
       if (["succeeded", "completed", "failed"].includes(existing.status)) {
+        logger.error("Cannot cancel workflow run: {params} with status {status}", {
+          params,
+          status: existing.status,
+        });
         throw new Error(
           `Cannot cancel workflow run ${params.workflowRunId} with status ${existing.status}`,
         );
       }
 
+      logger.error("Failed to cancel workflow run: {params}", { params });
       throw new Error("Failed to cancel workflow run");
     }
 
@@ -447,6 +519,10 @@ export class BackendPostgres implements Backend {
   }
 
   async createStepAttempt(params: CreateStepAttemptParams): Promise<StepAttempt> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     const [stepAttempt] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("step_attempts")
@@ -466,6 +542,7 @@ export class BackendPostgres implements Backend {
       .returning("*");
 
     if (!stepAttempt) {
+      logger.error("Failed to create step attempt: {params}", { params });
       throw new Error("Failed to create step attempt");
     }
 
@@ -473,6 +550,10 @@ export class BackendPostgres implements Backend {
   }
 
   async getStepAttempt(params: GetStepAttemptParams): Promise<StepAttempt | null> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     const stepAttempt = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("step_attempts")
@@ -484,6 +565,10 @@ export class BackendPostgres implements Backend {
   }
 
   async listStepAttempts(params: ListStepAttemptsParams): Promise<PaginatedResponse<StepAttempt>> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     const limit = params.limit ?? DEFAULT_PAGINATION_PAGE_SIZE;
     const { after, before } = params;
 
@@ -569,6 +654,10 @@ export class BackendPostgres implements Backend {
   }
 
   async completeStepAttempt(params: CompleteStepAttemptParams): Promise<StepAttempt> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     const [updated] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("step_attempts as sa")
@@ -591,6 +680,7 @@ export class BackendPostgres implements Backend {
       .returning("sa.*");
 
     if (!updated) {
+      logger.error("Failed to mark step attempt completed: {params}", { params });
       throw new Error("Failed to mark step attempt completed");
     }
 
@@ -598,6 +688,10 @@ export class BackendPostgres implements Backend {
   }
 
   async failStepAttempt(params: FailStepAttemptParams): Promise<StepAttempt> {
+    if (!this.initialized) {
+      throw new Error("Backend not initialized");
+    }
+
     const [updated] = await this.knex
       .withSchema(DEFAULT_SCHEMA)
       .table("step_attempts as sa")
@@ -620,6 +714,7 @@ export class BackendPostgres implements Backend {
       .returning("sa.*");
 
     if (!updated) {
+      logger.error("Failed to mark step attempt failed: {params}", { params });
       throw new Error("Failed to mark step attempt failed");
     }
 

package/src/database/base.ts

@@ -1,5 +1,5 @@
 import path from "node:path";
-import type { Knex } from "knex";
+import knex, { type Knex } from "knex";
 
 export const DEFAULT_SCHEMA = "sonamu_tasks";
 
@@ -7,12 +7,17 @@ export const DEFAULT_SCHEMA = "sonamu_tasks";
  * migrate applies pending migrations to the database. Does nothing if the
  * database is already up to date.
  */
-export async function migrate(knex: Knex, schema: string) {
-  await knex.schema.createSchemaIfNotExists(schema);
-  await knex.migrate.latest({
-    directory: path.join(import.meta.dirname, "migrations"),
-    schemaName: schema,
-  });
+export async function migrate(config: Knex.Config, schema: string) {
+  const instance = knex({ ...config, pool: { min: 1, max: 1 } });
+  try {
+    await instance.schema.createSchemaIfNotExists(schema);
+    await instance.migrate.latest({
+      directory: path.join(import.meta.dirname, "migrations"),
+      schemaName: schema,
+    });
+  } finally {
+    await instance.destroy();
+  }
 }
 
 /**

package/src/database/pubsub.ts

@@ -65,9 +65,15 @@ export class PostgresPubSub {
 
   // destroy the listener and close the connection, do not destroy the knex connection
   async destroy() {
-    this._destroyed = true;
-    this._connection.off("close", this._onClosed);
-    await this.knex.client.destroyRawConnection(this._connection);
+    if (this._destroyed) {
+      return;
+    }
+    try {
+      this._connection.off("close", this._onClosed);
+      await this.knex.client.destroyRawConnection(this._connection);
+    } finally {
+      this._destroyed = true;
+    }
   }
 
   // create a new listener and connect to the database

package/src/execution.test.ts

@@ -8,10 +8,11 @@ describe("StepExecutor", () => {
   let backend: BackendPostgres;
 
   beforeAll(async () => {
-    backend = await BackendPostgres.connect(KNEX_GLOBAL_CONFIG, {
+    backend = new BackendPostgres(KNEX_GLOBAL_CONFIG, {
       namespaceId: randomUUID(),
       runMigrations: false,
     });
+    await backend.initialize();
   });
 
   afterAll(async () => {
@@ -164,10 +165,11 @@ describe("executeWorkflow", () => {
   let backend: BackendPostgres;
 
   beforeAll(async () => {
-    backend = await BackendPostgres.connect(KNEX_GLOBAL_CONFIG, {
+    backend = new BackendPostgres(KNEX_GLOBAL_CONFIG, {
       namespaceId: randomUUID(),
       runMigrations: false,
     });
+    await backend.initialize();
   });
 
   afterAll(async () => {

package/src/practices/01-remote-workflow.ts

@@ -10,11 +10,12 @@ async function getBackend(): Promise<BackendPostgres> {
     return _backend;
   }
 
-  _backend = await BackendPostgres.connect(KNEX_GLOBAL_CONFIG, {
+  _backend = new BackendPostgres(KNEX_GLOBAL_CONFIG, {
     runMigrations: true,
     namespaceId: randomUUID(),
   });
 
+  await _backend.initialize();
   return _backend;
 }
 

package/src/testing/connection.ts

@@ -1,5 +1,5 @@
 import { randomUUID } from "node:crypto";
-import knex, { type Knex } from "knex";
+import type { Knex } from "knex";
 import { BackendPostgres } from "../database/backend";
 import { migrate as baseMigrate, DEFAULT_SCHEMA } from "../database/base";
 
@@ -20,7 +20,7 @@ export const KNEX_GLOBAL_CONFIG: Knex.Config = {
 } as const;
 
 export async function migrate(): Promise<void> {
-  await baseMigrate(knex(KNEX_GLOBAL_CONFIG), DEFAULT_SCHEMA);
+  await baseMigrate(KNEX_GLOBAL_CONFIG, DEFAULT_SCHEMA);
 }
 
 export async function createBackend(): Promise<BackendPostgres> {
@@ -28,9 +28,11 @@ export async function createBackend(): Promise<BackendPostgres> {
     return backend;
   }
 
-  backend = await BackendPostgres.connect(KNEX_GLOBAL_CONFIG, {
+  backend = new BackendPostgres(KNEX_GLOBAL_CONFIG, {
     namespaceId: randomUUID(),
+    runMigrations: false,
   });
+  await backend.initialize();
 
   return backend;
 }

package/src/worker.test.ts

@@ -8,10 +8,11 @@ describe("Worker", () => {
   let backend: BackendPostgres;
 
   beforeEach(async () => {
-    backend = await BackendPostgres.connect(KNEX_GLOBAL_CONFIG, {
+    backend = new BackendPostgres(KNEX_GLOBAL_CONFIG, {
       namespaceId: randomUUID(),
       runMigrations: false,
     });
+    await backend.initialize();
   });
 
   afterEach(async () => {

package/templates/openworkflow.config.ts

@@ -13,7 +13,7 @@ const config: Knex.Config = {
 } as const;
 
 // Use Postgres (configured with Knex config)
-const backend = await BackendPostgres.connect(config, {
+const backend = new BackendPostgres(config, {
   runMigrations: false,
 });