durable-execution-storage-drizzle/build/index.js

Drizzle ORM storage implementation for durable-execution

// src/sqlite.ts
import { and, eq, inArray, lt } from "drizzle-orm";
import {
  index,
  integer,
  sqliteTable,
  text,
  uniqueIndex
} from "drizzle-orm/sqlite-core";
import {
  createTransactionMutex
} from "durable-execution";

// src/common.ts
function storageObjectToInsertValue(obj) {
  return {
    rootTaskId: obj.rootTask?.taskId,
    rootExecutionId: obj.rootTask?.executionId,
    parentTaskId: obj.parentTask?.taskId,
    parentExecutionId: obj.parentTask?.executionId,
    isFinalizeTask: obj.parentTask?.isFinalizeTask,
    taskId: obj.taskId,
    executionId: obj.executionId,
    retryOptions: obj.retryOptions,
    timeoutMs: obj.timeoutMs,
    sleepMsBeforeRun: obj.sleepMsBeforeRun,
    runInput: obj.runInput,
    runOutput: obj.runOutput,
    output: obj.output,
    childrenTasksCompletedCount: obj.childrenTasksCompletedCount,
    childrenTasks: obj.childrenTasks,
    childrenTasksErrors: obj.childrenTasksErrors,
    finalizeTask: obj.finalizeTask,
    finalizeTaskError: obj.finalizeTaskError,
    error: obj.error,
    status: obj.status,
    isClosed: obj.isClosed,
    needsPromiseCancellation: obj.needsPromiseCancellation,
    retryAttempts: obj.retryAttempts,
    startAt: obj.startAt,
    startedAt: obj.startedAt,
    finishedAt: obj.finishedAt,
    expiresAt: obj.expiresAt,
    createdAt: obj.createdAt,
    updatedAt: obj.updatedAt
  };
}
function selectValueToStorageObject(row) {
  const obj = {
    taskId: row.taskId,
    executionId: row.executionId,
    retryOptions: row.retryOptions,
    timeoutMs: row.timeoutMs,
    sleepMsBeforeRun: row.sleepMsBeforeRun,
    runInput: row.runInput,
    childrenTasksCompletedCount: row.childrenTasksCompletedCount,
    status: row.status,
    isClosed: row.isClosed,
    needsPromiseCancellation: row.needsPromiseCancellation,
    retryAttempts: row.retryAttempts,
    startAt: row.startAt,
    createdAt: row.createdAt,
    updatedAt: row.updatedAt
  };
  if (row.rootTaskId && row.rootExecutionId) {
    obj.rootTask = {
      taskId: row.rootTaskId,
      executionId: row.rootExecutionId
    };
  }
  if (row.parentTaskId && row.parentExecutionId) {
    obj.parentTask = {
      taskId: row.parentTaskId,
      executionId: row.parentExecutionId,
      isFinalizeTask: row.isFinalizeTask ?? false
    };
  }
  if (row.runOutput != null) {
    obj.runOutput = row.runOutput;
  }
  if (row.output != null) {
    obj.output = row.output;
  }
  if (row.childrenTasks) {
    obj.childrenTasks = row.childrenTasks;
  }
  if (row.childrenTasksErrors) {
    obj.childrenTasksErrors = row.childrenTasksErrors;
  }
  if (row.finalizeTask) {
    obj.finalizeTask = row.finalizeTask;
  }
  if (row.finalizeTaskError) {
    obj.finalizeTaskError = row.finalizeTaskError;
  }
  if (row.error) {
    obj.error = row.error;
  }
  if (row.startedAt) {
    obj.startedAt = row.startedAt;
  }
  if (row.finishedAt) {
    obj.finishedAt = row.finishedAt;
  }
  if (row.expiresAt) {
    obj.expiresAt = row.expiresAt;
  }
  return obj;
}
function storageUpdateToUpdateValue(update) {
  const row = {};
  if (update.runOutput !== void 0) {
    row.runOutput = update.runOutput;
  }
  if (update.output !== void 0) {
    row.output = update.output;
  }
  if (update.childrenTasksCompletedCount !== void 0) {
    row.childrenTasksCompletedCount = update.childrenTasksCompletedCount;
  }
  if (update.childrenTasks !== void 0) {
    row.childrenTasks = update.childrenTasks;
  }
  if (update.childrenTasksErrors !== void 0) {
    row.childrenTasksErrors = update.childrenTasksErrors;
  }
  if (update.finalizeTask !== void 0) {
    row.finalizeTask = update.finalizeTask;
  }
  if (update.finalizeTaskError !== void 0) {
    row.finalizeTaskError = update.finalizeTaskError;
  }
  if (update.error !== void 0) {
    row.error = update.error;
  }
  if (update.unsetError) {
    row.error = null;
  }
  if (update.status !== void 0) {
    row.status = update.status;
  }
  if (update.isClosed !== void 0) {
    row.isClosed = update.isClosed;
  }
  if (update.needsPromiseCancellation !== void 0) {
    row.needsPromiseCancellation = update.needsPromiseCancellation;
  }
  if (update.retryAttempts !== void 0) {
    row.retryAttempts = update.retryAttempts;
  }
  if (update.startAt !== void 0) {
    row.startAt = update.startAt;
  }
  if (update.startedAt !== void 0) {
    row.startedAt = update.startedAt;
  }
  if (update.finishedAt !== void 0) {
    row.finishedAt = update.finishedAt;
  }
  if (update.expiresAt !== void 0) {
    row.expiresAt = update.expiresAt;
  }
  if (update.unsetExpiresAt) {
    row.expiresAt = null;
  }
  if (update.updatedAt !== void 0) {
    row.updatedAt = update.updatedAt;
  }
  return row;
}

// src/sqlite.ts
function createDurableTaskExecutionsSQLiteTable(tableName = "durable_task_executions") {
  return sqliteTable(
    tableName,
    {
      id: integer("id").primaryKey({ autoIncrement: true }),
      rootTaskId: text("root_task_id"),
      rootExecutionId: text("root_execution_id"),
      parentTaskId: text("parent_task_id"),
      parentExecutionId: text("parent_execution_id"),
      isFinalizeTask: integer("is_finalize_task", {
        mode: "boolean"
      }),
      taskId: text("task_id").notNull(),
      executionId: text("execution_id").notNull(),
      retryOptions: text("retry_options", { mode: "json" }).$type().notNull(),
      timeoutMs: integer("timeout_ms").notNull(),
      sleepMsBeforeRun: integer("sleep_ms_before_run").notNull(),
      runInput: text("run_input").notNull(),
      runOutput: text("run_output"),
      output: text("output"),
      childrenTasksCompletedCount: integer("children_tasks_completed_count").notNull(),
      childrenTasks: text("children_tasks", { mode: "json" }).$type(),
      childrenTasksErrors: text("children_tasks_errors", { mode: "json" }).$type(),
      finalizeTask: text("finalize_task", {
        mode: "json"
      }).$type(),
      finalizeTaskError: text("finalize_task_error", {
        mode: "json"
      }).$type(),
      error: text("error", { mode: "json" }).$type(),
      status: text("status").$type().notNull(),
      isClosed: integer("is_closed", { mode: "boolean" }).notNull(),
      needsPromiseCancellation: integer("needs_promise_cancellation", {
        mode: "boolean"
      }).notNull(),
      retryAttempts: integer("retry_attempts").notNull(),
      startAt: integer("start_at", { mode: "timestamp" }).notNull(),
      startedAt: integer("started_at", { mode: "timestamp" }),
      finishedAt: integer("finished_at", { mode: "timestamp" }),
      expiresAt: integer("expires_at", { mode: "timestamp" }),
      createdAt: integer("created_at", { mode: "timestamp" }).notNull(),
      updatedAt: integer("updated_at", { mode: "timestamp" }).notNull()
    },
    (table) => [
      uniqueIndex(`ix_${tableName}_execution_id`).on(table.executionId),
      index(`ix_${tableName}_status_is_closed_expires_at`).on(
        table.status,
        table.isClosed,
        table.expiresAt
      ),
      index(`ix_${tableName}_status_start_at`).on(table.status, table.startAt)
    ]
  );
}
function createSQLiteDurableStorage(db, table) {
  return new SQLiteDurableStorage(db, table);
}
var SQLiteDurableStorage = class {
  db;
  table;
  transactionMutex;
  constructor(db, table) {
    this.db = db;
    this.table = table;
    this.transactionMutex = createTransactionMutex();
  }
  async withTransaction(fn) {
    await this.transactionMutex.acquire();
    try {
      return await this.db.transaction(async (tx) => {
        const durableTx = new SQLiteDurableStorageTx(tx, this.table);
        return await fn(durableTx);
      });
    } finally {
      this.transactionMutex.release();
    }
  }
};
var SQLiteDurableStorageTx = class {
  tx;
  table;
  constructor(tx, table) {
    this.tx = tx;
    this.table = table;
  }
  async insertTaskExecutions(executions) {
    if (executions.length === 0) {
      return;
    }
    const rows = executions.map((execution) => storageObjectToInsertValue(execution));
    await this.tx.insert(this.table).values(rows);
  }
  async getTaskExecutionIds(where, limit) {
    const query = this.tx.select({ executionId: this.table.executionId }).from(this.table).where(buildWhereCondition(this.table, where));
    const rows = await (limit != null && limit > 0 ? query.limit(limit) : query);
    return rows.map((row) => row.executionId);
  }
  async getTaskExecutions(where, limit) {
    const query = this.tx.select().from(this.table).where(buildWhereCondition(this.table, where));
    const rows = await (limit != null && limit > 0 ? query.limit(limit) : query);
    return rows.map((row) => selectValueToStorageObject(row));
  }
  async updateTaskExecutions(where, update) {
    const rows = await this.tx.update(this.table).set(storageUpdateToUpdateValue(update)).where(buildWhereCondition(this.table, where)).returning({ executionId: this.table.executionId });
    return rows.map((row) => row.executionId);
  }
};
function buildWhereCondition(table, where) {
  const conditions = [];
  switch (where.type) {
    case "by_execution_ids": {
      conditions.push(inArray(table.executionId, where.executionIds));
      if (where.statuses) {
        conditions.push(inArray(table.status, where.statuses));
      }
      if (where.needsPromiseCancellation !== void 0) {
        conditions.push(eq(table.needsPromiseCancellation, where.needsPromiseCancellation));
      }
      break;
    }
    case "by_statuses": {
      conditions.push(inArray(table.status, where.statuses));
      if (where.isClosed !== void 0) {
        conditions.push(eq(table.isClosed, where.isClosed));
      }
      if (where.expiresAtLessThan) {
        conditions.push(lt(table.expiresAt, where.expiresAtLessThan));
      }
      break;
    }
    case "by_start_at_less_than": {
      conditions.push(lt(table.startAt, where.startAtLessThan));
      if (where.statuses) {
        conditions.push(inArray(table.status, where.statuses));
      }
      break;
    }
  }
  return and(...conditions);
}

// src/pg.ts
import { and as and2, eq as eq2, inArray as inArray2, lt as lt2 } from "drizzle-orm";
import {
  bigint,
  boolean,
  index as index2,
  integer as integer2,
  json,
  pgTable,
  text as text2,
  timestamp,
  uniqueIndex as uniqueIndex2
} from "drizzle-orm/pg-core";
function createDurableTaskExecutionsPgTable(tableName = "durable_task_executions") {
  return pgTable(
    tableName,
    {
      id: bigint("id", { mode: "number" }).primaryKey().generatedAlwaysAsIdentity(),
      rootTaskId: text2("root_task_id"),
      rootExecutionId: text2("root_execution_id"),
      parentTaskId: text2("parent_task_id"),
      parentExecutionId: text2("parent_execution_id"),
      isFinalizeTask: boolean("is_finalize_task"),
      taskId: text2("task_id").notNull(),
      executionId: text2("execution_id").notNull(),
      retryOptions: json("retry_options").$type().notNull(),
      timeoutMs: integer2("timeout_ms").notNull(),
      sleepMsBeforeRun: integer2("sleep_ms_before_run").notNull(),
      runInput: text2("run_input").notNull(),
      runOutput: text2("run_output"),
      output: text2("output"),
      childrenTasksCompletedCount: integer2("children_tasks_completed_count").notNull(),
      childrenTasks: json("children_tasks").$type(),
      childrenTasksErrors: json("children_tasks_errors").$type(),
      finalizeTask: json("finalize_task").$type(),
      finalizeTaskError: json("finalize_task_error").$type(),
      error: json("error").$type(),
      status: text2("status").$type().notNull(),
      isClosed: boolean("is_closed").notNull(),
      needsPromiseCancellation: boolean("needs_promise_cancellation").notNull(),
      retryAttempts: integer2("retry_attempts").notNull(),
      startAt: timestamp("start_at", { withTimezone: true }).notNull(),
      startedAt: timestamp("started_at", { withTimezone: true }),
      finishedAt: timestamp("finished_at", { withTimezone: true }),
      expiresAt: timestamp("expires_at", { withTimezone: true }),
      createdAt: timestamp("created_at", { withTimezone: true }).notNull(),
      updatedAt: timestamp("updated_at", { withTimezone: true }).notNull()
    },
    (table) => [
      uniqueIndex2(`ix_${tableName}_execution_id`).on(table.executionId),
      index2(`ix_${tableName}_status_is_closed_expires_at`).on(
        table.status,
        table.isClosed,
        table.expiresAt
      ),
      index2(`ix_${tableName}_status_start_at`).on(table.status, table.startAt)
    ]
  );
}
function createPgDurableStorage(db, table) {
  return new PgDurableStorage(db, table);
}
var PgDurableStorage = class {
  db;
  table;
  constructor(db, table) {
    this.db = db;
    this.table = table;
  }
  async withTransaction(fn) {
    return await this.db.transaction(async (tx) => {
      const durableTx = new PgDurableStorageTx(tx, this.table);
      return await fn(durableTx);
    });
  }
};
var PgDurableStorageTx = class {
  tx;
  table;
  constructor(tx, table) {
    this.tx = tx;
    this.table = table;
  }
  async insertTaskExecutions(executions) {
    if (executions.length === 0) {
      return;
    }
    const rows = executions.map((execution) => storageObjectToInsertValue(execution));
    await this.tx.insert(this.table).values(rows);
  }
  async getTaskExecutionIds(where, limit) {
    let rows = [];
    const query = this.tx.select({ executionId: this.table.executionId }).from(this.table).where(buildWhereCondition2(this.table, where));
    rows = await (limit != null && limit > 0 ? query.limit(limit) : query);
    return rows.map((row) => row.executionId);
  }
  async getTaskExecutions(where, limit) {
    const query = this.tx.select().from(this.table).where(buildWhereCondition2(this.table, where));
    const rows = await (limit != null && limit > 0 ? query.limit(limit) : query);
    return rows.map((row) => selectValueToStorageObject(row));
  }
  async updateTaskExecutions(where, update) {
    const rows = await this.tx.update(this.table).set(storageUpdateToUpdateValue(update)).where(buildWhereCondition2(this.table, where)).returning({ executionId: this.table.executionId });
    return rows.map((row) => row.executionId);
  }
};
function buildWhereCondition2(table, where) {
  const conditions = [];
  switch (where.type) {
    case "by_execution_ids": {
      conditions.push(inArray2(table.executionId, where.executionIds));
      if (where.statuses) {
        conditions.push(inArray2(table.status, where.statuses));
      }
      if (where.needsPromiseCancellation !== void 0) {
        conditions.push(eq2(table.needsPromiseCancellation, where.needsPromiseCancellation));
      }
      break;
    }
    case "by_statuses": {
      conditions.push(inArray2(table.status, where.statuses));
      if (where.isClosed !== void 0) {
        conditions.push(eq2(table.isClosed, where.isClosed));
      }
      if (where.expiresAtLessThan) {
        conditions.push(lt2(table.expiresAt, where.expiresAtLessThan));
      }
      break;
    }
    case "by_start_at_less_than": {
      conditions.push(lt2(table.startAt, where.startAtLessThan));
      if (where.statuses) {
        conditions.push(inArray2(table.status, where.statuses));
      }
      break;
    }
  }
  return and2(...conditions);
}

// src/mysql.ts
import {
  and as and3,
  eq as eq3,
  inArray as inArray3,
  lt as lt3
} from "drizzle-orm";
import {
  bigint as bigint2,
  boolean as boolean2,
  index as index3,
  int,
  json as json2,
  mysqlTable,
  text as text3,
  timestamp as timestamp2,
  uniqueIndex as uniqueIndex3,
  varchar
} from "drizzle-orm/mysql-core";
function createDurableTaskExecutionsMySQLTable(tableName = "durable_task_executions") {
  return mysqlTable(
    tableName,
    {
      id: bigint2("id", { mode: "number" }).primaryKey().autoincrement(),
      rootTaskId: text3("root_task_id"),
      rootExecutionId: text3("root_execution_id"),
      parentTaskId: text3("parent_task_id"),
      parentExecutionId: text3("parent_execution_id"),
      isFinalizeTask: boolean2("is_finalize_task"),
      taskId: text3("task_id").notNull(),
      executionId: varchar("execution_id", { length: 255 }).notNull(),
      retryOptions: json2("retry_options").$type().notNull(),
      timeoutMs: int("timeout_ms").notNull(),
      sleepMsBeforeRun: int("sleep_ms_before_run").notNull(),
      runInput: text3("run_input").notNull(),
      runOutput: text3("run_output"),
      output: text3("output"),
      childrenTasksCompletedCount: int("children_tasks_completed_count").notNull(),
      childrenTasks: json2("children_tasks").$type(),
      childrenTasksErrors: json2("children_tasks_errors").$type(),
      finalizeTask: json2("finalize_task").$type(),
      finalizeTaskError: json2("finalize_task_error").$type(),
      error: json2("error").$type(),
      status: varchar("status", { length: 32 }).$type().notNull(),
      isClosed: boolean2("is_closed").notNull(),
      needsPromiseCancellation: boolean2("needs_promise_cancellation").notNull(),
      retryAttempts: int("retry_attempts").notNull(),
      startAt: timestamp2("start_at").notNull(),
      startedAt: timestamp2("started_at"),
      finishedAt: timestamp2("finished_at"),
      expiresAt: timestamp2("expires_at"),
      createdAt: timestamp2("created_at").notNull(),
      updatedAt: timestamp2("updated_at").notNull()
    },
    (table) => [
      uniqueIndex3(`ix_${tableName}_execution_id`).on(table.executionId),
      index3(`ix_${tableName}_status_is_closed_expires_at`).on(
        table.status,
        table.isClosed,
        table.expiresAt
      ),
      index3(`ix_${tableName}_status_start_at`).on(table.status, table.startAt)
    ]
  );
}
function createMySQLDurableStorage(db, table) {
  return new MySQLDurableStorage(db, table);
}
var MySQLDurableStorage = class {
  db;
  table;
  constructor(db, table) {
    this.db = db;
    this.table = table;
  }
  async withTransaction(fn) {
    return await this.db.transaction(async (tx) => {
      const durableTx = new MySQLDurableStorageTx(tx, this.table);
      return await fn(durableTx);
    });
  }
};
var MySQLDurableStorageTx = class {
  tx;
  table;
  constructor(tx, table) {
    this.tx = tx;
    this.table = table;
  }
  async insertTaskExecutions(executions) {
    if (executions.length === 0) {
      return;
    }
    const rows = executions.map((execution) => storageObjectToInsertValue(execution));
    await this.tx.insert(this.table).values(rows);
  }
  async getTaskExecutionIds(where, limit) {
    let rows = [];
    const query = this.tx.select({ executionId: this.table.executionId }).from(this.table).where(buildWhereCondition3(this.table, where));
    rows = await (limit != null && limit > 0 ? query.limit(limit) : query);
    return rows.map((row) => row.executionId);
  }
  async getTaskExecutions(where, limit) {
    const query = this.tx.select().from(this.table).where(buildWhereCondition3(this.table, where));
    const rows = await (limit != null && limit > 0 ? query.limit(limit) : query);
    return rows.map((row) => selectValueToStorageObject(row));
  }
  async updateTaskExecutions(where, update) {
    const rowsToUpdate = await this.tx.select({ executionId: this.table.executionId }).from(this.table).where(buildWhereCondition3(this.table, where)).for("update");
    if (rowsToUpdate.length === 0) {
      return [];
    }
    const executionIds = rowsToUpdate.map((row) => row.executionId);
    await this.tx.update(this.table).set(storageUpdateToUpdateValue(update)).where(inArray3(this.table.executionId, executionIds));
    return executionIds;
  }
};
function buildWhereCondition3(table, where) {
  const conditions = [];
  switch (where.type) {
    case "by_execution_ids": {
      conditions.push(inArray3(table.executionId, where.executionIds));
      if (where.statuses) {
        conditions.push(inArray3(table.status, where.statuses));
      }
      if (where.needsPromiseCancellation !== void 0) {
        conditions.push(eq3(table.needsPromiseCancellation, where.needsPromiseCancellation));
      }
      break;
    }
    case "by_statuses": {
      conditions.push(inArray3(table.status, where.statuses));
      if (where.isClosed !== void 0) {
        conditions.push(eq3(table.isClosed, where.isClosed));
      }
      if (where.expiresAtLessThan) {
        conditions.push(lt3(table.expiresAt, where.expiresAtLessThan));
      }
      break;
    }
    case "by_start_at_less_than": {
      conditions.push(lt3(table.startAt, where.startAtLessThan));
      if (where.statuses) {
        conditions.push(inArray3(table.status, where.statuses));
      }
      break;
    }
  }
  return and3(...conditions);
}
export {
  createDurableTaskExecutionsMySQLTable,
  createDurableTaskExecutionsPgTable,
  createDurableTaskExecutionsSQLiteTable,
  createMySQLDurableStorage,
  createPgDurableStorage,
  createSQLiteDurableStorage
};
//# sourceMappingURL=index.js.map