durable-execution-storage-drizzle
Version:
Drizzle ORM storage implementation for durable-execution
315 lines (313 loc) • 9.82 kB
JavaScript
// src/mysql.ts
import {
and,
eq,
inArray,
lt
} from "drizzle-orm";
import {
bigint,
boolean,
index,
int,
json,
mysqlTable,
text,
timestamp,
uniqueIndex,
varchar
} from "drizzle-orm/mysql-core";
// 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/mysql.ts
function createDurableTaskExecutionsMySQLTable(tableName = "durable_task_executions") {
return mysqlTable(
tableName,
{
id: bigint("id", { mode: "number" }).primaryKey().autoincrement(),
rootTaskId: text("root_task_id"),
rootExecutionId: text("root_execution_id"),
parentTaskId: text("parent_task_id"),
parentExecutionId: text("parent_execution_id"),
isFinalizeTask: boolean("is_finalize_task"),
taskId: text("task_id").notNull(),
executionId: varchar("execution_id", { length: 255 }).notNull(),
retryOptions: json("retry_options").$type().notNull(),
timeoutMs: int("timeout_ms").notNull(),
sleepMsBeforeRun: int("sleep_ms_before_run").notNull(),
runInput: text("run_input").notNull(),
runOutput: text("run_output"),
output: text("output"),
childrenTasksCompletedCount: int("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: varchar("status", { length: 32 }).$type().notNull(),
isClosed: boolean("is_closed").notNull(),
needsPromiseCancellation: boolean("needs_promise_cancellation").notNull(),
retryAttempts: int("retry_attempts").notNull(),
startAt: timestamp("start_at").notNull(),
startedAt: timestamp("started_at"),
finishedAt: timestamp("finished_at"),
expiresAt: timestamp("expires_at"),
createdAt: timestamp("created_at").notNull(),
updatedAt: timestamp("updated_at").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 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(buildWhereCondition(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(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 rowsToUpdate = await this.tx.select({ executionId: this.table.executionId }).from(this.table).where(buildWhereCondition(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(inArray(this.table.executionId, executionIds));
return executionIds;
}
};
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);
}
export {
createDurableTaskExecutionsMySQLTable,
createMySQLDurableStorage
};
//# sourceMappingURL=mysql.js.map