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

Drizzle ORM storage implementation for durable-execution

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

// src/common.ts
import {
  applyTaskExecutionStorageUpdate
} from "durable-execution";
import { omitUndefinedValues } from "@gpahal/std/objects";
function taskExecutionStorageValueToDBValue(value) {
  return {
    rootTaskId: value.root?.taskId,
    rootExecutionId: value.root?.executionId,
    parentTaskId: value.parent?.taskId,
    parentExecutionId: value.parent?.executionId,
    indexInParentChildren: value.parent?.indexInParentChildren,
    isOnlyChildOfParent: value.parent?.isOnlyChildOfParent,
    isFinalizeOfParent: value.parent?.isFinalizeOfParent,
    taskId: value.taskId,
    executionId: value.executionId,
    isSleepingTask: value.isSleepingTask,
    sleepingTaskUniqueId: value.sleepingTaskUniqueId,
    retryOptions: value.retryOptions,
    sleepMsBeforeRun: value.sleepMsBeforeRun,
    timeoutMs: value.timeoutMs,
    areChildrenSequential: value.areChildrenSequential,
    input: value.input,
    executorId: value.executorId,
    status: value.status,
    isFinished: value.isFinished,
    runOutput: value.runOutput,
    output: value.output,
    error: value.error,
    retryAttempts: value.retryAttempts,
    startAt: value.startAt,
    startedAt: value.startedAt,
    expiresAt: value.expiresAt,
    waitingForChildrenStartedAt: value.waitingForChildrenStartedAt,
    waitingForFinalizeStartedAt: value.waitingForFinalizeStartedAt,
    finishedAt: value.finishedAt,
    children: value.children,
    activeChildrenCount: value.activeChildrenCount,
    onChildrenFinishedProcessingStatus: value.onChildrenFinishedProcessingStatus,
    onChildrenFinishedProcessingExpiresAt: value.onChildrenFinishedProcessingExpiresAt,
    onChildrenFinishedProcessingFinishedAt: value.onChildrenFinishedProcessingFinishedAt,
    finalize: value.finalize,
    closeStatus: value.closeStatus,
    closeExpiresAt: value.closeExpiresAt,
    closedAt: value.closedAt,
    needsPromiseCancellation: value.needsPromiseCancellation,
    createdAt: value.createdAt,
    updatedAt: value.updatedAt
  };
}
function taskExecutionDBValueToStorageValue(dbValue, update, updateExpiresAtWithStartedAt) {
  const value = {
    taskId: dbValue.taskId,
    executionId: dbValue.executionId,
    isSleepingTask: dbValue.isSleepingTask,
    retryOptions: dbValue.retryOptions,
    sleepMsBeforeRun: dbValue.sleepMsBeforeRun,
    timeoutMs: dbValue.timeoutMs,
    areChildrenSequential: dbValue.areChildrenSequential,
    input: dbValue.input,
    status: dbValue.status,
    isFinished: dbValue.isFinished,
    retryAttempts: dbValue.retryAttempts,
    startAt: dbValue.startAt,
    activeChildrenCount: dbValue.activeChildrenCount,
    onChildrenFinishedProcessingStatus: dbValue.onChildrenFinishedProcessingStatus,
    closeStatus: dbValue.closeStatus,
    needsPromiseCancellation: dbValue.needsPromiseCancellation,
    createdAt: dbValue.createdAt,
    updatedAt: dbValue.updatedAt
  };
  if (dbValue.rootTaskId && dbValue.rootExecutionId) {
    value.root = {
      taskId: dbValue.rootTaskId,
      executionId: dbValue.rootExecutionId
    };
  }
  if (dbValue.parentTaskId && dbValue.parentExecutionId) {
    value.parent = {
      taskId: dbValue.parentTaskId,
      executionId: dbValue.parentExecutionId,
      indexInParentChildren: dbValue.indexInParentChildren ?? 0,
      isOnlyChildOfParent: dbValue.isOnlyChildOfParent ?? false,
      isFinalizeOfParent: dbValue.isFinalizeOfParent ?? false
    };
  }
  if (dbValue.executorId) {
    value.executorId = dbValue.executorId;
  }
  if (dbValue.sleepingTaskUniqueId != null) {
    value.sleepingTaskUniqueId = dbValue.sleepingTaskUniqueId;
  }
  if (dbValue.runOutput != null) {
    value.runOutput = dbValue.runOutput;
  }
  if (dbValue.output != null) {
    value.output = dbValue.output;
  }
  if (dbValue.error) {
    value.error = dbValue.error;
  }
  if (dbValue.startedAt) {
    value.startedAt = dbValue.startedAt;
  }
  if (dbValue.expiresAt) {
    value.expiresAt = dbValue.expiresAt;
  }
  if (dbValue.waitingForChildrenStartedAt) {
    value.waitingForChildrenStartedAt = dbValue.waitingForChildrenStartedAt;
  }
  if (dbValue.waitingForFinalizeStartedAt) {
    value.waitingForFinalizeStartedAt = dbValue.waitingForFinalizeStartedAt;
  }
  if (dbValue.expiresAt) {
    value.expiresAt = dbValue.expiresAt;
  }
  if (dbValue.finishedAt) {
    value.finishedAt = dbValue.finishedAt;
  }
  if (dbValue.children) {
    value.children = dbValue.children;
  }
  if (dbValue.onChildrenFinishedProcessingExpiresAt) {
    value.onChildrenFinishedProcessingExpiresAt = dbValue.onChildrenFinishedProcessingExpiresAt;
  }
  if (dbValue.onChildrenFinishedProcessingFinishedAt) {
    value.onChildrenFinishedProcessingFinishedAt = dbValue.onChildrenFinishedProcessingFinishedAt;
  }
  if (dbValue.finalize) {
    value.finalize = dbValue.finalize;
  }
  if (dbValue.closeExpiresAt) {
    value.closeExpiresAt = dbValue.closeExpiresAt;
  }
  if (dbValue.closedAt) {
    value.closedAt = dbValue.closedAt;
  }
  return update ? applyTaskExecutionStorageUpdate(value, update, updateExpiresAtWithStartedAt) : value;
}
function taskExecutionStorageUpdateToDBUpdate(update) {
  const dbUpdate = omitUndefinedValues({
    ...update,
    children: update.children,
    unset: void 0
  });
  const updateUnset = update.unset;
  if (updateUnset) {
    for (const key in updateUnset) {
      if (updateUnset[key]) {
        dbUpdate[key] = null;
      }
    }
  }
  return dbUpdate;
}

// src/sqlite.ts
function createSQLiteTaskExecutionsTable(tableName = "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"),
      indexInParentChildren: integer("index_in_parent_children"),
      isOnlyChildOfParent: integer("is_only_child_of_parent", {
        mode: "boolean"
      }),
      isFinalizeOfParent: integer("is_finalize_of_parent", {
        mode: "boolean"
      }),
      taskId: text("task_id").notNull(),
      executionId: text("execution_id").notNull(),
      retryOptions: text("retry_options", { mode: "json" }).$type().notNull(),
      sleepMsBeforeRun: integer("sleep_ms_before_run").notNull(),
      timeoutMs: integer("timeout_ms").notNull(),
      areChildrenSequential: integer("are_children_sequential", {
        mode: "boolean"
      }).notNull(),
      input: text("input").notNull(),
      executorId: text("executor_id"),
      isSleepingTask: integer("is_sleeping_task", { mode: "boolean" }).notNull(),
      sleepingTaskUniqueId: text("sleeping_task_unique_id"),
      status: text("status").$type().notNull(),
      isFinished: integer("is_finished", { mode: "boolean" }).notNull(),
      runOutput: text("run_output"),
      output: text("output"),
      error: text("error", { mode: "json" }).$type(),
      retryAttempts: integer("retry_attempts").notNull(),
      startAt: integer("start_at").notNull(),
      startedAt: integer("started_at"),
      expiresAt: integer("expires_at"),
      waitingForChildrenStartedAt: integer("waiting_for_children_started_at"),
      waitingForFinalizeStartedAt: integer("waiting_for_finalize_started_at"),
      finishedAt: integer("finished_at"),
      children: text("children", { mode: "json" }).$type(),
      activeChildrenCount: integer("active_children_count").notNull(),
      onChildrenFinishedProcessingStatus: text("on_children_finished_processing_status").$type().notNull(),
      onChildrenFinishedProcessingExpiresAt: integer("on_children_finished_processing_expires_at"),
      onChildrenFinishedProcessingFinishedAt: integer(
        "on_children_finished_processing_finished_at"
      ),
      finalize: text("finalize", { mode: "json" }).$type(),
      closeStatus: text("close_status").$type().notNull(),
      closeExpiresAt: integer("close_expires_at"),
      closedAt: integer("closed_at"),
      needsPromiseCancellation: integer("needs_promise_cancellation", {
        mode: "boolean"
      }).notNull(),
      createdAt: integer("created_at").notNull(),
      updatedAt: integer("updated_at").notNull()
    },
    (table) => [
      uniqueIndex(`ix_${tableName}_execution_id`).on(table.executionId),
      uniqueIndex(`ix_${tableName}_sleeping_task_unique_id`).on(table.sleepingTaskUniqueId),
      index(`ix_${tableName}_status_start_at`).on(table.status, table.startAt),
      index(`ix_${tableName}_status_ocfp_status_acc_updated_at`).on(
        table.status,
        table.onChildrenFinishedProcessingStatus,
        table.activeChildrenCount,
        table.updatedAt
      ),
      index(`ix_${tableName}_close_status_updated_at`).on(table.closeStatus, table.updatedAt),
      index(`ix_${tableName}_status_is_sleeping_task_expires_at`).on(
        table.status,
        table.isSleepingTask,
        table.expiresAt
      ),
      index(`ix_${tableName}_on_ocfp_expires_at`).on(table.onChildrenFinishedProcessingExpiresAt),
      index(`ix_${tableName}_close_expires_at`).on(table.closeExpiresAt),
      index(`ix_${tableName}_executor_id_npc_updated_at`).on(
        table.executorId,
        table.needsPromiseCancellation,
        table.updatedAt
      ),
      index(`ix_${tableName}_parent_execution_id_is_finished`).on(
        table.parentExecutionId,
        table.isFinished
      ),
      index(`ix_${tableName}_is_finished_close_status_updated_at`).on(
        table.isFinished,
        table.closeStatus,
        table.updatedAt
      )
    ]
  );
}
function createSQLiteTaskExecutionsStorage(db, taskExecutionsTable, options = {}) {
  return new TaskExecutionsStorageWithMutex(
    new TaskExecutionsStorageWithBatching(
      new SQLiteTaskExecutionsStorageNonAtomic(db, taskExecutionsTable, options),
      true
    )
  );
}
var SQLiteTaskExecutionsStorageNonAtomic = class {
  db;
  taskExecutionsTable;
  enableTestMode;
  constructor(db, taskExecutionsTable, {
    enableTestMode = false
  } = {}) {
    this.db = db;
    this.taskExecutionsTable = taskExecutionsTable;
    this.enableTestMode = enableTestMode;
  }
  async insertMany(executions) {
    if (executions.length === 0) {
      return;
    }
    const rows = executions.map((execution) => taskExecutionStorageValueToDBValue(execution));
    await this.db.insert(this.taskExecutionsTable).values(rows);
  }
  async getById({
    executionId,
    filters
  }) {
    const rows = await this.db.select().from(this.taskExecutionsTable).where(getByIdWhereCondition(this.taskExecutionsTable, executionId, filters)).limit(1);
    return rows.length > 0 ? taskExecutionDBValueToStorageValue(rows[0]) : void 0;
  }
  async getBySleepingTaskUniqueId({
    sleepingTaskUniqueId
  }) {
    const rows = await this.db.select().from(this.taskExecutionsTable).where(eq(this.taskExecutionsTable.sleepingTaskUniqueId, sleepingTaskUniqueId)).limit(1);
    return rows.length > 0 ? taskExecutionDBValueToStorageValue(rows[0]) : void 0;
  }
  async updateById({
    executionId,
    filters,
    update
  }) {
    const dbUpdate = taskExecutionStorageUpdateToDBUpdate(update);
    await this.db.update(this.taskExecutionsTable).set(dbUpdate).where(getByIdWhereCondition(this.taskExecutionsTable, executionId, filters));
  }
  async updateByIdAndInsertChildrenIfUpdated({
    executionId,
    filters,
    update,
    childrenTaskExecutionsToInsertIfAnyUpdated
  }) {
    const dbUpdate = taskExecutionStorageUpdateToDBUpdate(update);
    if (childrenTaskExecutionsToInsertIfAnyUpdated.length === 0) {
      return await this.updateById({ executionId, filters, update });
    }
    const rowsToInsert = childrenTaskExecutionsToInsertIfAnyUpdated.map(
      (execution) => taskExecutionStorageValueToDBValue(execution)
    );
    await this.db.transaction(async (tx) => {
      const updateResult = await tx.update(this.taskExecutionsTable).set(dbUpdate).where(getByIdWhereCondition(this.taskExecutionsTable, executionId, filters)).returning({ executionId: this.taskExecutionsTable.executionId });
      if (updateResult.length > 0) {
        await tx.insert(this.taskExecutionsTable).values(rowsToInsert);
      }
    });
  }
  async updateByStatusAndStartAtLessThanAndReturn({
    status,
    startAtLessThan,
    update,
    updateExpiresAtWithStartedAt,
    limit
  }) {
    const dbUpdate = taskExecutionStorageUpdateToDBUpdate(update);
    const updatedRows = await this.db.transaction(async (tx) => {
      const rows = await tx.select().from(this.taskExecutionsTable).where(
        and(
          eq(this.taskExecutionsTable.status, status),
          lt(this.taskExecutionsTable.startAt, startAtLessThan)
        )
      ).orderBy(asc(this.taskExecutionsTable.startAt)).limit(limit);
      if (rows.length > 0) {
        await tx.update(this.taskExecutionsTable).set({
          ...dbUpdate,
          expiresAt: sql`(${updateExpiresAtWithStartedAt} + timeout_ms)`
        }).where(
          inArray(
            this.taskExecutionsTable.executionId,
            rows.map((row) => row.executionId)
          )
        );
      }
      return rows;
    });
    return updatedRows.map(
      (row) => taskExecutionDBValueToStorageValue(row, update, updateExpiresAtWithStartedAt)
    );
  }
  async updateByStatusAndOnChildrenFinishedProcessingStatusAndActiveChildrenCountZeroAndReturn({
    status,
    onChildrenFinishedProcessingStatus,
    update,
    limit
  }) {
    const dbUpdate = taskExecutionStorageUpdateToDBUpdate(update);
    const updatedRows = await this.db.transaction(async (tx) => {
      const rows = await tx.select().from(this.taskExecutionsTable).where(
        and(
          eq(this.taskExecutionsTable.status, status),
          eq(
            this.taskExecutionsTable.onChildrenFinishedProcessingStatus,
            onChildrenFinishedProcessingStatus
          ),
          eq(this.taskExecutionsTable.activeChildrenCount, 0)
        )
      ).orderBy(asc(this.taskExecutionsTable.updatedAt)).limit(limit);
      if (rows.length > 0) {
        await tx.update(this.taskExecutionsTable).set(dbUpdate).where(
          inArray(
            this.taskExecutionsTable.executionId,
            rows.map((row) => row.executionId)
          )
        );
      }
      return rows;
    });
    return updatedRows.map((row) => taskExecutionDBValueToStorageValue(row, update));
  }
  async updateByCloseStatusAndReturn({
    closeStatus,
    update,
    limit
  }) {
    const dbUpdate = taskExecutionStorageUpdateToDBUpdate(update);
    const updatedRows = await this.db.transaction(async (tx) => {
      const rows = await tx.select().from(this.taskExecutionsTable).where(eq(this.taskExecutionsTable.closeStatus, closeStatus)).orderBy(asc(this.taskExecutionsTable.updatedAt)).limit(limit);
      if (rows.length > 0) {
        await tx.update(this.taskExecutionsTable).set(dbUpdate).where(
          inArray(
            this.taskExecutionsTable.executionId,
            rows.map((row) => row.executionId)
          )
        );
      }
      return rows;
    });
    return updatedRows.map((row) => taskExecutionDBValueToStorageValue(row, update));
  }
  async updateByStatusAndIsSleepingTaskAndExpiresAtLessThan({
    status,
    isSleepingTask,
    expiresAtLessThan,
    update,
    limit
  }) {
    const dbUpdate = taskExecutionStorageUpdateToDBUpdate(update);
    return await this.db.transaction(async (tx) => {
      const rows = await tx.select({ executionId: this.taskExecutionsTable.executionId }).from(this.taskExecutionsTable).where(
        and(
          eq(this.taskExecutionsTable.status, status),
          eq(this.taskExecutionsTable.isSleepingTask, isSleepingTask),
          lt(this.taskExecutionsTable.expiresAt, expiresAtLessThan)
        )
      ).orderBy(asc(this.taskExecutionsTable.expiresAt)).limit(limit);
      if (rows.length > 0) {
        await tx.update(this.taskExecutionsTable).set(dbUpdate).where(
          inArray(
            this.taskExecutionsTable.executionId,
            rows.map((row) => row.executionId)
          )
        );
      }
      return rows.length;
    });
  }
  async updateByOnChildrenFinishedProcessingExpiresAtLessThan({
    onChildrenFinishedProcessingExpiresAtLessThan,
    update,
    limit
  }) {
    const dbUpdate = taskExecutionStorageUpdateToDBUpdate(update);
    return await this.db.transaction(async (tx) => {
      const rows = await tx.select({ executionId: this.taskExecutionsTable.executionId }).from(this.taskExecutionsTable).where(
        lt(
          this.taskExecutionsTable.onChildrenFinishedProcessingExpiresAt,
          onChildrenFinishedProcessingExpiresAtLessThan
        )
      ).orderBy(asc(this.taskExecutionsTable.onChildrenFinishedProcessingExpiresAt)).limit(limit);
      if (rows.length > 0) {
        await tx.update(this.taskExecutionsTable).set(dbUpdate).where(
          inArray(
            this.taskExecutionsTable.executionId,
            rows.map((row) => row.executionId)
          )
        );
      }
      return rows.length;
    });
  }
  async updateByCloseExpiresAtLessThan({
    closeExpiresAtLessThan,
    update,
    limit
  }) {
    const dbUpdate = taskExecutionStorageUpdateToDBUpdate(update);
    return await this.db.transaction(async (tx) => {
      const rows = await tx.select({ executionId: this.taskExecutionsTable.executionId }).from(this.taskExecutionsTable).where(lt(this.taskExecutionsTable.closeExpiresAt, closeExpiresAtLessThan)).orderBy(asc(this.taskExecutionsTable.closeExpiresAt)).limit(limit);
      if (rows.length > 0) {
        await tx.update(this.taskExecutionsTable).set(dbUpdate).where(
          inArray(
            this.taskExecutionsTable.executionId,
            rows.map((row) => row.executionId)
          )
        );
      }
      return rows.length;
    });
  }
  async updateByExecutorIdAndNeedsPromiseCancellationAndReturn({
    executorId,
    needsPromiseCancellation,
    update,
    limit
  }) {
    const dbUpdate = taskExecutionStorageUpdateToDBUpdate(update);
    const updatedRows = await this.db.transaction(async (tx) => {
      const rows = await tx.select().from(this.taskExecutionsTable).where(
        and(
          eq(this.taskExecutionsTable.executorId, executorId),
          eq(this.taskExecutionsTable.needsPromiseCancellation, needsPromiseCancellation)
        )
      ).orderBy(asc(this.taskExecutionsTable.updatedAt)).limit(limit);
      if (rows.length > 0) {
        await tx.update(this.taskExecutionsTable).set(dbUpdate).where(
          inArray(
            this.taskExecutionsTable.executionId,
            rows.map((row) => row.executionId)
          )
        );
      }
      return rows;
    });
    return updatedRows.map((row) => taskExecutionDBValueToStorageValue(row, update));
  }
  async getByParentExecutionId({
    parentExecutionId
  }) {
    const rows = await this.db.select().from(this.taskExecutionsTable).where(eq(this.taskExecutionsTable.parentExecutionId, parentExecutionId));
    return rows.map((row) => taskExecutionDBValueToStorageValue(row));
  }
  async updateByParentExecutionIdAndIsFinished({
    parentExecutionId,
    isFinished,
    update
  }) {
    const dbUpdate = taskExecutionStorageUpdateToDBUpdate(update);
    await this.db.update(this.taskExecutionsTable).set(dbUpdate).where(
      and(
        eq(this.taskExecutionsTable.parentExecutionId, parentExecutionId),
        eq(this.taskExecutionsTable.isFinished, isFinished)
      )
    );
  }
  async updateAndDecrementParentActiveChildrenCountByIsFinishedAndCloseStatus({
    isFinished,
    closeStatus,
    update,
    limit
  }) {
    const dbUpdate = taskExecutionStorageUpdateToDBUpdate(update);
    return await this.db.transaction(async (tx) => {
      const rows = await tx.select().from(this.taskExecutionsTable).where(
        and(
          eq(this.taskExecutionsTable.isFinished, isFinished),
          eq(this.taskExecutionsTable.closeStatus, closeStatus)
        )
      ).orderBy(asc(this.taskExecutionsTable.updatedAt)).limit(limit);
      if (rows.length > 0) {
        await tx.update(this.taskExecutionsTable).set(dbUpdate).where(
          inArray(
            this.taskExecutionsTable.executionId,
            rows.map((row) => row.executionId)
          )
        );
        const parentExecutionIdToDecrementValueMap = /* @__PURE__ */ new Map();
        for (const row of rows) {
          if (row.parentExecutionId != null) {
            parentExecutionIdToDecrementValueMap.set(
              row.parentExecutionId,
              (parentExecutionIdToDecrementValueMap.get(row.parentExecutionId) ?? 0) + 1
            );
          }
        }
        if (parentExecutionIdToDecrementValueMap.size > 0) {
          const parentIds = [...parentExecutionIdToDecrementValueMap.keys()].sort();
          const caseStatements = parentIds.map(
            (parentId) => sql`WHEN ${parentId} THEN ${parentExecutionIdToDecrementValueMap.get(parentId)}`
          );
          await tx.run(sql`
            UPDATE ${this.taskExecutionsTable}
            SET
              active_children_count = active_children_count - (
                CASE execution_id
                  ${sql.join(caseStatements, sql` `)}
                  ELSE 0
                END
              ),
              updated_at = ${dbUpdate.updatedAt}
            WHERE execution_id IN (${sql.join(parentIds, sql`, `)})
          `);
        }
      }
      return rows.length;
    });
  }
  async deleteById({ executionId }) {
    if (!this.enableTestMode) {
      return;
    }
    await this.db.delete(this.taskExecutionsTable).where(eq(this.taskExecutionsTable.executionId, executionId));
  }
  async deleteAll() {
    if (!this.enableTestMode) {
      return;
    }
    await this.db.delete(this.taskExecutionsTable);
  }
};
function getByIdWhereCondition(table, executionId, filters) {
  const conditions = [eq(table.executionId, executionId)];
  if (filters?.isSleepingTask != null) {
    conditions.push(eq(table.isSleepingTask, filters.isSleepingTask));
  }
  if (filters?.status != null) {
    conditions.push(eq(table.status, filters.status));
  }
  if (filters?.isFinished != null) {
    conditions.push(eq(table.isFinished, filters.isFinished));
  }
  return and(...conditions);
}
export {
  createSQLiteTaskExecutionsStorage,
  createSQLiteTaskExecutionsTable
};
//# sourceMappingURL=sqlite.js.map