Added KV store feature

This commit is contained in:
Eric Allam
2023-03-12 20:12:12 +00:00
parent 5e762a9cbf
commit 1673d4527a
22 changed files with 1146 additions and 39 deletions
+5
View File
@@ -0,0 +1,5 @@
---
"@trigger.dev/sdk": patch
---
Added kv storage to persist data in between runs and between workflows
@@ -1,4 +1,9 @@
import type { SecureString } from "@trigger.dev/common-schemas"; import {
KVDeleteSchema,
KVGetSchema,
KVSetSchema,
SecureString,
} from "@trigger.dev/common-schemas";
import { import {
CustomEventSchema, CustomEventSchema,
ErrorSchema, ErrorSchema,
@@ -105,6 +110,25 @@ async function parseStep(
type: "CUSTOM_EVENT" as const, type: "CUSTOM_EVENT" as const,
input: await CustomEventSchema.parseAsync(original.input), input: await CustomEventSchema.parseAsync(original.input),
}; };
case "KV_GET":
return {
...base,
type: "KV_GET" as const,
output: original.output,
input: await KVGetSchema.parseAsync(original.input),
};
case "KV_SET":
return {
...base,
type: "KV_SET" as const,
input: await KVSetSchema.parseAsync(original.input),
};
case "KV_DELETE":
return {
...base,
type: "KV_DELETE" as const,
input: await KVDeleteSchema.parseAsync(original.input),
};
case "OUTPUT": case "OUTPUT":
return { return {
...base, ...base,
@@ -13,6 +13,7 @@ import {
ArrowPathRoundedSquareIcon, ArrowPathRoundedSquareIcon,
ChatBubbleOvalLeftEllipsisIcon, ChatBubbleOvalLeftEllipsisIcon,
CheckCircleIcon, CheckCircleIcon,
CircleStackIcon,
ExclamationCircleIcon, ExclamationCircleIcon,
ExclamationTriangleIcon, ExclamationTriangleIcon,
} from "@heroicons/react/24/solid"; } from "@heroicons/react/24/solid";
@@ -517,6 +518,12 @@ function StepBody({ step }: { step: Step }) {
return <FetchRequestStep request={step} />; return <FetchRequestStep request={step} />;
case "DURABLE_DELAY": case "DURABLE_DELAY":
return <DelayStep step={step} />; return <DelayStep step={step} />;
case "KV_GET":
return <KVGetStep step={step} />;
case "KV_SET":
return <KVSetStep step={step} />;
case "KV_DELETE":
return <KVDeleteStep step={step} />;
} }
return <></>; return <></>;
} }
@@ -708,6 +715,104 @@ function CustomEventStep({ event }: { event: StepType<Step, "CUSTOM_EVENT"> }) {
); );
} }
function KVGetStep({ step }: { step: StepType<Step, "KV_GET"> }) {
const scope = step.input.namespace.split(":")[0];
return (
<>
<div className="flex gap-16">
<div>
<Body size="extra-small" className={workflowNodeUppercaseClasses}>
Key
</Body>
<Header2 size="small" className="mb-2 text-slate-300">
{step.input.key}
</Header2>
</div>
<div>
<Body size="extra-small" className={workflowNodeUppercaseClasses}>
Scope
</Body>
<Header2 size="small" className="mb-2 text-slate-300">
{scope}
</Header2>
</div>
</div>
{step.output && (
<>
<Header4>Output</Header4>
<CodeBlock code={stringifyCode(step.output)} align="top" />
</>
)}
</>
);
}
function KVSetStep({ step }: { step: StepType<Step, "KV_SET"> }) {
const scope = step.input.namespace.split(":")[0];
return (
<>
<div className="flex gap-16">
<div>
<Body size="extra-small" className={workflowNodeUppercaseClasses}>
Key
</Body>
<Header2 size="small" className="mb-2 text-slate-300">
{step.input.key}
</Header2>
</div>
<div>
<Body size="extra-small" className={workflowNodeUppercaseClasses}>
Scope
</Body>
<Header2 size="small" className="mb-2 text-slate-300">
{scope}
</Header2>
</div>
</div>
{step.input.value && (
<>
<Header4>Value</Header4>
<CodeBlock code={stringifyCode(step.input.value)} align="top" />
</>
)}
</>
);
}
function KVDeleteStep({ step }: { step: StepType<Step, "KV_DELETE"> }) {
const scope = step.input.namespace.split(":")[0];
return (
<>
<div className="flex gap-16">
<div>
<Body size="extra-small" className={workflowNodeUppercaseClasses}>
Key
</Body>
<Header2 size="small" className="mb-2 text-slate-300">
{step.input.key}
</Header2>
</div>
<div>
<Body size="extra-small" className={workflowNodeUppercaseClasses}>
Scope
</Body>
<Header2 size="small" className="mb-2 text-slate-300">
{scope}
</Header2>
</div>
</div>
</>
);
}
function RunOnceStep({ event }: { event: StepType<Step, "RUN_ONCE"> }) { function RunOnceStep({ event }: { event: StepType<Step, "RUN_ONCE"> }) {
return ( return (
<> <>
@@ -964,6 +1069,18 @@ const stepInfo: Record<Step["type"], { label: string; icon: ReactNode }> = {
label: "Run once", label: "Run once",
icon: <KeyIcon className={styleClass} />, icon: <KeyIcon className={styleClass} />,
}, },
KV_GET: {
label: "Get Key Value",
icon: <CircleStackIcon className={styleClass} />,
},
KV_SET: {
label: "Set Key Value",
icon: <CircleStackIcon className={styleClass} />,
},
KV_DELETE: {
label: "Delete Key Value",
icon: <CircleStackIcon className={styleClass} />,
},
} as const; } as const;
type LogLevel = StepType<Step, "LOG_MESSAGE">["input"]["level"]; type LogLevel = StepType<Step, "LOG_MESSAGE">["input"]["level"];
@@ -15,7 +15,7 @@ export class ResolveDelay {
step: { step: {
include: { include: {
run: { run: {
include: { environment: true }, include: { environment: true, workflow: true },
}, },
}, },
}, },
@@ -23,13 +23,20 @@ export class ResolveDelay {
}); });
if (!existingDelay) { if (!existingDelay) {
throw new Error("Delay not found"); return;
} }
if (existingDelay.resolvedAt) { if (existingDelay.resolvedAt) {
return existingDelay; return existingDelay;
} }
if (
existingDelay.step.run.workflow.disabledAt ||
existingDelay.step.run.workflow.archivedAt
) {
return;
}
const delay = await this.#prismaClient.durableDelay.update({ const delay = await this.#prismaClient.durableDelay.update({
where: { id }, where: { id },
data: { resolvedAt: new Date() }, data: { resolvedAt: new Date() },
@@ -0,0 +1,253 @@
import type {
KVDeleteOperation,
KVGetOperation,
KVSetOperation,
} from "@trigger.dev/common-schemas";
import type { PrismaClient } from "~/db.server";
import { prisma } from "~/db.server";
import type { WorkflowRunStep } from "~/models/workflowRun.server";
import { createStepOnce } from "~/models/workflowRunStep.server";
import { taskQueue } from "../messageBroker.server";
export class KVGetService {
#prismaClient: PrismaClient;
constructor(prismaClient: PrismaClient = prisma) {
this.#prismaClient = prismaClient;
}
async call(
key: string,
runId: string,
apiKey: string,
timestamp: string,
data: KVGetOperation
) {
const environment = await this.#prismaClient.runtimeEnvironment.findUnique({
where: {
apiKey,
},
include: {
organization: true,
},
});
if (!environment) {
throw new Error("Invalid API key");
}
const workflowRun = await this.#prismaClient.workflowRun.findUnique({
where: {
id: runId,
},
include: {
workflow: true,
},
});
if (!workflowRun) {
throw new Error("Invalid workflow run ID");
}
if (workflowRun.workflow.organizationId !== environment.organizationId) {
throw new Error("Invalid workflow run ID");
}
const fullKey = `${data.namespace}:${data.key}`;
const kvItem = await this.#prismaClient.keyValueItem.findUnique({
where: {
environmentId_key: {
environmentId: environment.id,
key: fullKey,
},
},
});
const idempotentStep = await createStepOnce(workflowRun.id, key, {
type: "KV_GET",
input: data,
output: kvItem?.value ? kvItem.value : undefined,
context: {},
status: "SUCCESS",
ts: timestamp,
});
return {
output: idempotentStep.step.output,
environment,
};
}
}
export class KVSetService {
#prismaClient: PrismaClient;
constructor(prismaClient: PrismaClient = prisma) {
this.#prismaClient = prismaClient;
}
async call(
key: string,
runId: string,
apiKey: string,
timestamp: string,
data: KVSetOperation
) {
const environment = await this.#prismaClient.runtimeEnvironment.findUnique({
where: {
apiKey,
},
include: {
organization: true,
},
});
if (!environment) {
throw new Error("Invalid API key");
}
const workflowRun = await this.#prismaClient.workflowRun.findUnique({
where: {
id: runId,
},
include: {
workflow: true,
},
});
if (!workflowRun) {
throw new Error("Invalid workflow run ID");
}
if (workflowRun.workflow.organizationId !== environment.organizationId) {
throw new Error("Invalid workflow run ID");
}
const fullKey = `${data.namespace}:${data.key}`;
const value = JSON.parse(JSON.stringify(data.value));
const idempotentStep = await createStepOnce(workflowRun.id, key, {
type: "KV_SET",
input: {
...data,
value,
},
context: {},
status: "PENDING",
ts: timestamp,
});
if (idempotentStep.status === "EXISTING") {
return environment;
}
await this.#prismaClient.keyValueItem.upsert({
where: {
environmentId_key: {
environmentId: environment.id,
key: fullKey,
},
},
update: {
value,
},
create: {
environmentId: environment.id,
key: fullKey,
value,
},
});
await this.#prismaClient.workflowRunStep.update({
where: {
id: idempotentStep.step.id,
},
data: {
status: "SUCCESS",
},
});
return environment;
}
}
export class KVDeleteService {
#prismaClient: PrismaClient;
constructor(prismaClient: PrismaClient = prisma) {
this.#prismaClient = prismaClient;
}
async call(
key: string,
runId: string,
apiKey: string,
timestamp: string,
data: KVDeleteOperation
) {
const environment = await this.#prismaClient.runtimeEnvironment.findUnique({
where: {
apiKey,
},
include: {
organization: true,
},
});
if (!environment) {
throw new Error("Invalid API key");
}
const workflowRun = await this.#prismaClient.workflowRun.findUnique({
where: {
id: runId,
},
include: {
workflow: true,
},
});
if (!workflowRun) {
throw new Error("Invalid workflow run ID");
}
if (workflowRun.workflow.organizationId !== environment.organizationId) {
throw new Error("Invalid workflow run ID");
}
const fullKey = `${data.namespace}:${data.key}`;
const idempotentStep = await createStepOnce(workflowRun.id, key, {
type: "KV_DELETE",
input: data,
context: {},
status: "PENDING",
ts: timestamp,
});
if (idempotentStep.status === "EXISTING") {
return environment;
}
await this.#prismaClient.keyValueItem.delete({
where: {
environmentId_key: {
environmentId: environment.id,
key: fullKey,
},
},
});
await this.#prismaClient.workflowRunStep.update({
where: {
id: idempotentStep.step.id,
},
data: {
status: "SUCCESS",
},
});
return environment;
}
}
+114 -12
View File
@@ -59,6 +59,11 @@ import { RegisterExternalSource } from "./externalSources/registerExternalSource
import { CreateFetchRequest } from "./fetches/createFetchRequest.server"; import { CreateFetchRequest } from "./fetches/createFetchRequest.server";
import { PerformFetchRequest } from "./fetches/performFetchRequest.server"; import { PerformFetchRequest } from "./fetches/performFetchRequest.server";
import { StartFetchRequest } from "./fetches/startFetchRequest.server"; import { StartFetchRequest } from "./fetches/startFetchRequest.server";
import {
KVDeleteService,
KVGetService,
KVSetService,
} from "./kv/services.server";
import type { PulsarClient } from "./pulsarClient.server"; import type { PulsarClient } from "./pulsarClient.server";
import { createPulsarClient } from "./pulsarClient.server"; import { createPulsarClient } from "./pulsarClient.server";
import { CreateIntegrationRequest } from "./requests/createIntegrationRequest.server"; import { CreateIntegrationRequest } from "./requests/createIntegrationRequest.server";
@@ -280,6 +285,100 @@ function createCommandSubscriber() {
return true; return true;
}, },
SEND_KV_GET: async (id, data, properties) => {
const service = new KVGetService();
const { output, environment } = await service.call(
data.key,
properties["x-workflow-run-id"],
properties["x-api-key"],
properties["x-timestamp"],
data.get
);
await commandResponsePublisher.publish(
"RESOLVE_KV_GET",
{
key: data.key,
operation: {
key: data.get.key,
namespace: data.get.namespace,
output,
},
},
{
"x-workflow-run-id": properties["x-workflow-run-id"],
"x-api-key": properties["x-api-key"],
"x-org-id": environment.organization.id,
"x-workflow-id": properties["x-workflow-id"],
"x-env": environment.slug,
}
);
return true;
},
SEND_KV_SET: async (id, data, properties) => {
const service = new KVSetService();
const environment = await service.call(
data.key,
properties["x-workflow-run-id"],
properties["x-api-key"],
properties["x-timestamp"],
data.set
);
await commandResponsePublisher.publish(
"RESOLVE_KV_SET",
{
key: data.key,
operation: {
key: data.set.key,
namespace: data.set.namespace,
},
},
{
"x-workflow-run-id": properties["x-workflow-run-id"],
"x-api-key": properties["x-api-key"],
"x-org-id": environment.organization.id,
"x-workflow-id": properties["x-workflow-id"],
"x-env": environment.slug,
}
);
return true;
},
SEND_KV_DELETE: async (id, data, properties) => {
const service = new KVDeleteService();
const environment = await service.call(
data.key,
properties["x-workflow-run-id"],
properties["x-api-key"],
properties["x-timestamp"],
data.delete
);
await commandResponsePublisher.publish(
"RESOLVE_KV_DELETE",
{
key: data.key,
operation: {
key: data.delete.key,
namespace: data.delete.namespace,
},
},
{
"x-workflow-run-id": properties["x-workflow-run-id"],
"x-api-key": properties["x-api-key"],
"x-org-id": environment.organization.id,
"x-workflow-id": properties["x-workflow-id"],
"x-env": environment.slug,
}
);
return true;
},
TRIGGER_CUSTOM_EVENT: async (id, data, properties) => { TRIGGER_CUSTOM_EVENT: async (id, data, properties) => {
await triggerEventInRun( await triggerEventInRun(
data.key, data.key,
@@ -571,19 +670,22 @@ function createTaskQueue() {
RESOLVE_DELAY: async (id, data, properties) => { RESOLVE_DELAY: async (id, data, properties) => {
const service = new ResolveDelay(); const service = new ResolveDelay();
const { step } = await service.call(data.id); const delay = await service.call(data.id);
if (delay) {
commandResponsePublisher.publish(
"RESOLVE_DELAY",
{ id: data.id, key: delay.step.idempotencyKey },
{
"x-workflow-run-id": delay.step.run.id,
"x-api-key": delay.step.run.environment.apiKey,
"x-org-id": delay.step.run.environment.organizationId,
"x-workflow-id": delay.step.run.workflowId,
"x-env": delay.step.run.environment.slug,
}
);
}
commandResponsePublisher.publish(
"RESOLVE_DELAY",
{ id: data.id, key: step.idempotencyKey },
{
"x-workflow-run-id": step.run.id,
"x-api-key": step.run.environment.apiKey,
"x-org-id": step.run.environment.organizationId,
"x-workflow-id": step.run.workflowId,
"x-env": step.run.environment.slug,
}
);
return true; return true;
}, },
INTEGRATION_REQUEST_CREATED: async (id, data, properties) => { INTEGRATION_REQUEST_CREATED: async (id, data, properties) => {
@@ -0,0 +1,20 @@
-- AlterEnum
ALTER TYPE "WorkflowRunStepType" ADD VALUE 'KEY_VALUE';
-- CreateTable
CREATE TABLE "KeyValueItem" (
"id" TEXT NOT NULL,
"key" TEXT NOT NULL,
"value" JSONB NOT NULL,
"environmentId" TEXT NOT NULL,
"createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updatedAt" TIMESTAMP(3) NOT NULL,
CONSTRAINT "KeyValueItem_pkey" PRIMARY KEY ("id")
);
-- CreateIndex
CREATE UNIQUE INDEX "KeyValueItem_environmentId_key_key" ON "KeyValueItem"("environmentId", "key");
-- AddForeignKey
ALTER TABLE "KeyValueItem" ADD CONSTRAINT "KeyValueItem_environmentId_fkey" FOREIGN KEY ("environmentId") REFERENCES "RuntimeEnvironment"("id") ON DELETE CASCADE ON UPDATE CASCADE;
@@ -0,0 +1,14 @@
/*
Warnings:
- The values [KEY_VALUE] on the enum `WorkflowRunStepType` will be removed. If these variants are still used in the database, this will fail.
*/
-- AlterEnum
BEGIN;
CREATE TYPE "WorkflowRunStepType_new" AS ENUM ('OUTPUT', 'LOG_MESSAGE', 'DURABLE_DELAY', 'CUSTOM_EVENT', 'INTEGRATION_REQUEST', 'DISCONNECTION', 'FETCH_REQUEST', 'RUN_ONCE', 'KV_GET', 'KV_SET', 'KV_DELETE');
ALTER TABLE "WorkflowRunStep" ALTER COLUMN "type" TYPE "WorkflowRunStepType_new" USING ("type"::text::"WorkflowRunStepType_new");
ALTER TYPE "WorkflowRunStepType" RENAME TO "WorkflowRunStepType_old";
ALTER TYPE "WorkflowRunStepType_new" RENAME TO "WorkflowRunStepType";
DROP TYPE "WorkflowRunStepType_old";
COMMIT;
+19
View File
@@ -114,6 +114,7 @@ model RuntimeEnvironment {
schedulerSources SchedulerSource[] schedulerSources SchedulerSource[]
internalSources InternalSource[] internalSources InternalSource[]
deployments ProjectDeployment[] deployments ProjectDeployment[]
keyValueItems KeyValueItem[]
@@unique([organizationId, slug]) @@unique([organizationId, slug])
} }
@@ -526,6 +527,24 @@ enum WorkflowRunStepType {
DISCONNECTION DISCONNECTION
FETCH_REQUEST FETCH_REQUEST
RUN_ONCE RUN_ONCE
KV_GET
KV_SET
KV_DELETE
}
model KeyValueItem {
id String @id @default(cuid())
key String
value Json
environment RuntimeEnvironment @relation(fields: [environmentId], references: [id], onDelete: Cascade, onUpdate: Cascade)
environmentId String
createdAt DateTime @default(now())
updatedAt DateTime @updatedAt
@@unique([environmentId, key])
} }
model FetchRequest { model FetchRequest {
+67 -24
View File
@@ -63,12 +63,11 @@ export class WorkflowRunController {
subscriptionType: "Exclusive", subscriptionType: "Exclusive",
subscriptionInitialPosition: "Latest", subscriptionInitialPosition: "Latest",
}, },
filter: {
"x-workflow-run-id": this.#runId,
},
handlers: { handlers: {
RESOLVE_DELAY: async (id, data, properties) => { RESOLVE_DELAY: async (id, data, properties) => {
if (properties["x-workflow-run-id"] !== this.#runId) {
return true;
}
this.#logger.debug( this.#logger.debug(
"Received resolve delay request", "Received resolve delay request",
id, id,
@@ -91,10 +90,6 @@ export class WorkflowRunController {
return success; return success;
}, },
RESOLVE_INTEGRATION_REQUEST: async (id, data, properties) => { RESOLVE_INTEGRATION_REQUEST: async (id, data, properties) => {
if (properties["x-workflow-run-id"] !== this.#runId) {
return true;
}
this.#logger.debug( this.#logger.debug(
"Received resolve integration request", "Received resolve integration request",
id, id,
@@ -118,10 +113,6 @@ export class WorkflowRunController {
return success; return success;
}, },
RESOLVE_RUN_ONCE: async (id, data, properties) => { RESOLVE_RUN_ONCE: async (id, data, properties) => {
if (properties["x-workflow-run-id"] !== this.#runId) {
return true;
}
this.#logger.debug("Received resolve runOnce", id, data, properties); this.#logger.debug("Received resolve runOnce", id, data, properties);
const success = await this.#hostRPC.send("RESOLVE_RUN_ONCE", { const success = await this.#hostRPC.send("RESOLVE_RUN_ONCE", {
@@ -140,10 +131,6 @@ export class WorkflowRunController {
return success; return success;
}, },
REJECT_INTEGRATION_REQUEST: async (id, data, properties) => { REJECT_INTEGRATION_REQUEST: async (id, data, properties) => {
if (properties["x-workflow-run-id"] !== this.#runId) {
return true;
}
this.#logger.debug( this.#logger.debug(
"Received reject integration request", "Received reject integration request",
id, id,
@@ -167,10 +154,6 @@ export class WorkflowRunController {
return success; return success;
}, },
RESOLVE_FETCH_REQUEST: async (id, data, properties) => { RESOLVE_FETCH_REQUEST: async (id, data, properties) => {
if (properties["x-workflow-run-id"] !== this.#runId) {
return true;
}
this.#logger.debug( this.#logger.debug(
"Received resolve fetch request", "Received resolve fetch request",
id, id,
@@ -194,10 +177,6 @@ export class WorkflowRunController {
return success; return success;
}, },
REJECT_FETCH_REQUEST: async (id, data, properties) => { REJECT_FETCH_REQUEST: async (id, data, properties) => {
if (properties["x-workflow-run-id"] !== this.#runId) {
return true;
}
this.#logger.debug( this.#logger.debug(
"Received reject fetch request", "Received reject fetch request",
id, id,
@@ -218,6 +197,70 @@ export class WorkflowRunController {
}, },
}); });
return success;
},
RESOLVE_KV_GET: async (id, data, properties) => {
this.#logger.debug(
"Received RESOLVE_KV_GET request",
id,
data,
properties
);
const success = await this.#hostRPC.send("RESOLVE_KV_GET", {
output: data.operation.output,
key: data.key,
meta: {
workflowId: properties["x-workflow-id"],
organizationId: properties["x-org-id"],
environment: properties["x-env"],
apiKey: properties["x-api-key"],
runId: properties["x-workflow-run-id"],
},
});
return success;
},
RESOLVE_KV_SET: async (id, data, properties) => {
this.#logger.debug(
"Received RESOLVE_KV_SET request",
id,
data,
properties
);
const success = await this.#hostRPC.send("RESOLVE_KV_SET", {
key: data.key,
meta: {
workflowId: properties["x-workflow-id"],
organizationId: properties["x-org-id"],
environment: properties["x-env"],
apiKey: properties["x-api-key"],
runId: properties["x-workflow-run-id"],
},
});
return success;
},
RESOLVE_KV_DELETE: async (id, data, properties) => {
this.#logger.debug(
"Received RESOLVE_KV_DELETE request",
id,
data,
properties
);
const success = await this.#hostRPC.send("RESOLVE_KV_DELETE", {
key: data.key,
meta: {
workflowId: properties["x-workflow-id"],
organizationId: properties["x-org-id"],
environment: properties["x-env"],
apiKey: properties["x-api-key"],
runId: properties["x-workflow-run-id"],
},
});
return success; return success;
}, },
}, },
+45
View File
@@ -169,6 +169,51 @@ export class TriggerServer {
return !!response; return !!response;
}, },
SEND_KV_GET: async (request) => {
const runController = this.#runControllers.get(request.runId);
if (!runController) {
// TODO: need to recover from this issue by trying to reconnect
return false;
}
const response = await runController.publish("SEND_KV_GET", {
key: request.key,
get: request.get,
});
return !!response;
},
SEND_KV_SET: async (request) => {
const runController = this.#runControllers.get(request.runId);
if (!runController) {
// TODO: need to recover from this issue by trying to reconnect
return false;
}
const response = await runController.publish("SEND_KV_SET", {
key: request.key,
set: request.set,
});
return !!response;
},
SEND_KV_DELETE: async (request) => {
const runController = this.#runControllers.get(request.runId);
if (!runController) {
// TODO: need to recover from this issue by trying to reconnect
return false;
}
const response = await runController.publish("SEND_KV_DELETE", {
key: request.key,
delete: request.delete,
});
return !!response;
},
INITIALIZE_RUN_ONCE: async (request) => { INITIALIZE_RUN_ONCE: async (request) => {
const runController = this.#runControllers.get(request.runId); const runController = this.#runControllers.get(request.runId);
+1
View File
@@ -6,3 +6,4 @@ export * from "./events";
export * from "./triggers"; export * from "./triggers";
export * from "./fetch"; export * from "./fetch";
export * from "./runOnce"; export * from "./runOnce";
export * from "./kv";
+24
View File
@@ -0,0 +1,24 @@
import { z } from "zod";
import { SerializableJsonSchema } from "./json";
export const KVGetSchema = z.object({
key: z.string(),
namespace: z.string(),
});
export type KVGetOperation = z.infer<typeof KVGetSchema>;
export const KVDeleteSchema = z.object({
key: z.string(),
namespace: z.string(),
});
export type KVDeleteOperation = z.infer<typeof KVDeleteSchema>;
export const KVSetSchema = z.object({
key: z.string(),
namespace: z.string(),
value: SerializableJsonSchema,
});
export type KVSetOperation = z.infer<typeof KVSetSchema>;
@@ -3,6 +3,7 @@ import {
FetchOutputSchema, FetchOutputSchema,
JsonSchema, JsonSchema,
ResolveRunOnceOuputSchema, ResolveRunOnceOuputSchema,
SerializableJsonSchema,
} from "@trigger.dev/common-schemas"; } from "@trigger.dev/common-schemas";
export const HostRPCSchema = { export const HostRPCSchema = {
@@ -114,6 +115,46 @@ export const HostRPCSchema = {
}), }),
response: z.boolean(), response: z.boolean(),
}, },
RESOLVE_KV_GET: {
request: z.object({
key: z.string(),
output: SerializableJsonSchema,
meta: z.object({
environment: z.string(),
workflowId: z.string(),
organizationId: z.string(),
apiKey: z.string(),
runId: z.string(),
}),
}),
response: z.boolean(),
},
RESOLVE_KV_SET: {
request: z.object({
key: z.string(),
meta: z.object({
environment: z.string(),
workflowId: z.string(),
organizationId: z.string(),
apiKey: z.string(),
runId: z.string(),
}),
}),
response: z.boolean(),
},
RESOLVE_KV_DELETE: {
request: z.object({
key: z.string(),
meta: z.object({
environment: z.string(),
workflowId: z.string(),
organizationId: z.string(),
apiKey: z.string(),
runId: z.string(),
}),
}),
response: z.boolean(),
},
}; };
export type HostRPC = typeof HostRPCSchema; export type HostRPC = typeof HostRPCSchema;
@@ -171,6 +171,43 @@ export const ServerRPCSchema = {
}), }),
response: z.boolean(), response: z.boolean(),
}, },
SEND_KV_GET: {
request: z.object({
runId: z.string(),
key: z.string(),
timestamp: z.string(),
get: z.object({
namespace: z.string(),
key: z.string(),
}),
}),
response: z.boolean(),
},
SEND_KV_SET: {
request: z.object({
runId: z.string(),
key: z.string(),
timestamp: z.string(),
set: z.object({
namespace: z.string(),
key: z.string(),
value: SerializableJsonSchema,
}),
}),
response: z.boolean(),
},
SEND_KV_DELETE: {
request: z.object({
runId: z.string(),
key: z.string(),
timestamp: z.string(),
delete: z.object({
namespace: z.string(),
key: z.string(),
}),
}),
response: z.boolean(),
},
}; };
export type ServerRPC = typeof ServerRPCSchema; export type ServerRPC = typeof ServerRPCSchema;
@@ -2,12 +2,14 @@ import { commandResponses as integrationRequests } from "../schemas/integrationR
import { commandResponses as delays } from "../schemas/delays"; import { commandResponses as delays } from "../schemas/delays";
import { commandResponses as fetchRequests } from "../schemas/fetchRequests"; import { commandResponses as fetchRequests } from "../schemas/fetchRequests";
import { commandResponses as runOnce } from "../schemas/runOnce"; import { commandResponses as runOnce } from "../schemas/runOnce";
import { commandResponses as kvStorage } from "../schemas/kvStorage";
const Catalog = { const Catalog = {
...integrationRequests, ...integrationRequests,
...delays, ...delays,
...fetchRequests, ...fetchRequests,
...runOnce, ...runOnce,
...kvStorage,
}; };
export default Catalog; export default Catalog;
@@ -5,6 +5,7 @@ import { commands as customEvents } from "../schemas/customEvents";
import { commands as delays } from "../schemas/delays"; import { commands as delays } from "../schemas/delays";
import { commands as fetchRequests } from "../schemas/fetchRequests"; import { commands as fetchRequests } from "../schemas/fetchRequests";
import { commands as runOnce } from "../schemas/runOnce"; import { commands as runOnce } from "../schemas/runOnce";
import { commands as kvStorage } from "../schemas/kvStorage";
const Catalog = { const Catalog = {
...integrationRequests, ...integrationRequests,
@@ -14,6 +15,7 @@ const Catalog = {
...delays, ...delays,
...fetchRequests, ...fetchRequests,
...runOnce, ...runOnce,
...kvStorage,
}; };
export default Catalog; export default Catalog;
@@ -0,0 +1,68 @@
import {
KVDeleteSchema,
KVGetSchema,
KVSetSchema,
} from "@trigger.dev/common-schemas";
import { z } from "zod";
import {
WorkflowRunEventPropertiesSchema,
WorkflowSendRunEventPropertiesSchema,
} from "../sharedSchemas";
export const commandResponses = {
RESOLVE_KV_GET: {
data: z.object({
key: z.string(),
operation: z.object({
key: z.string(),
namespace: z.string(),
output: z.any(),
}),
}),
properties: WorkflowRunEventPropertiesSchema,
},
RESOLVE_KV_SET: {
data: z.object({
key: z.string(),
operation: z.object({
key: z.string(),
namespace: z.string(),
}),
}),
properties: WorkflowRunEventPropertiesSchema,
},
RESOLVE_KV_DELETE: {
data: z.object({
key: z.string(),
operation: z.object({
key: z.string(),
namespace: z.string(),
}),
}),
properties: WorkflowRunEventPropertiesSchema,
},
};
export const commands = {
SEND_KV_GET: {
data: z.object({
key: z.string(),
get: KVGetSchema,
}),
properties: WorkflowSendRunEventPropertiesSchema,
},
SEND_KV_SET: {
data: z.object({
key: z.string(),
set: KVSetSchema,
}),
properties: WorkflowSendRunEventPropertiesSchema,
},
SEND_KV_DELETE: {
data: z.object({
key: z.string(),
delete: KVDeleteSchema,
}),
properties: WorkflowSendRunEventPropertiesSchema,
},
};
@@ -40,12 +40,14 @@ export type ZodSubscriberOptions<
config: Omit<PulsarConsumerConfig, "listener">; config: Omit<PulsarConsumerConfig, "listener">;
schema: SubscriberSchema; schema: SubscriberSchema;
handlers: ZodSubscriberHandlers<SubscriberSchema>; handlers: ZodSubscriberHandlers<SubscriberSchema>;
filter?: Record<string, string>;
}; };
export class ZodSubscriber<SubscriberSchema extends MessageCatalogSchema> { export class ZodSubscriber<SubscriberSchema extends MessageCatalogSchema> {
#config: Omit<PulsarConsumerConfig, "listener">; #config: Omit<PulsarConsumerConfig, "listener">;
#schema: SubscriberSchema; #schema: SubscriberSchema;
#handlers: ZodSubscriberHandlers<SubscriberSchema>; #handlers: ZodSubscriberHandlers<SubscriberSchema>;
#filter?: Record<string, string>;
#subscriber?: PulsarConsumer; #subscriber?: PulsarConsumer;
#client: PulsarClient; #client: PulsarClient;
@@ -58,6 +60,7 @@ export class ZodSubscriber<SubscriberSchema extends MessageCatalogSchema> {
this.#schema = options.schema; this.#schema = options.schema;
this.#handlers = options.handlers; this.#handlers = options.handlers;
this.#client = options.client; this.#client = options.client;
this.#filter = options.filter;
this.#logger = new Logger("trigger.dev subscriber"); this.#logger = new Logger("trigger.dev subscriber");
} }
@@ -123,6 +126,23 @@ export class ZodSubscriber<SubscriberSchema extends MessageCatalogSchema> {
const eventTimestamp = msg.getEventTimestamp(); const eventTimestamp = msg.getEventTimestamp();
const redeliveryCount = msg.getRedeliveryCount(); const redeliveryCount = msg.getRedeliveryCount();
const filter = this.#filter;
// Return if the filter exists and doesn't match
if (filter) {
const filterKeys = Object.keys(filter);
if (
filterKeys.some((key) => {
return properties[key] !== filter[key];
})
) {
await consumer.acknowledge(msg);
return;
}
}
this.#logger.debug("#onMessage", { this.#logger.debug("#onMessage", {
messageId, messageId,
publishedTimestamp, publishedTimestamp,
+193
View File
@@ -20,6 +20,12 @@ import chalk from "chalk";
import getRepoInfo from "git-repo-info"; import getRepoInfo from "git-repo-info";
import gitRemoteOriginUrl from "git-remote-origin-url"; import gitRemoteOriginUrl from "git-remote-origin-url";
import { readFile } from "node:fs/promises"; import { readFile } from "node:fs/promises";
import {
ContextKeyValueStorage,
KvDeleteFunction,
KvGetFunction,
KvSetFunction,
} from "./keyValueStorage";
const zodErrorMessageOptions: ErrorMessageOptions = { const zodErrorMessageOptions: ErrorMessageOptions = {
delimiter: { delimiter: {
@@ -93,6 +99,30 @@ export class TriggerClient<TSchema extends z.ZodTypeAny> {
} }
>(); >();
#kvGetCallbacks = new Map<
string,
{
resolve: (output: any) => void;
reject: (err?: any) => void;
}
>();
#kvSetCallbacks = new Map<
string,
{
resolve: () => void;
reject: (err?: any) => void;
}
>();
#kvDeleteCallbacks = new Map<
string,
{
resolve: () => void;
reject: (err?: any) => void;
}
>();
constructor(trigger: Trigger<TSchema>, options: TriggerOptions<TSchema>) { constructor(trigger: Trigger<TSchema>, options: TriggerOptions<TSchema>) {
this.#trigger = trigger; this.#trigger = trigger;
this.#options = options; this.#options = options;
@@ -384,6 +414,78 @@ export class TriggerClient<TSchema extends z.ZodTypeAny> {
return true; return true;
}, },
RESOLVE_KV_GET: async (data) => {
this.#logger.debug("Handling RESOLVE_KV_GET", data);
const getCallbacks = this.#kvGetCallbacks.get(
messageKey(data.meta.runId, data.key)
);
if (!getCallbacks) {
this.#logger.debug(
`Could not find kvGet callbacks for request ID ${messageKey(
data.meta.runId,
data.key
)}. This can happen when a workflow run is resumed`
);
return true;
}
const { resolve } = getCallbacks;
resolve(data.output);
return true;
},
RESOLVE_KV_SET: async (data) => {
this.#logger.debug("Handling RESOLVE_KV_SET", data);
const setCallbacks = this.#kvSetCallbacks.get(
messageKey(data.meta.runId, data.key)
);
if (!setCallbacks) {
this.#logger.debug(
`Could not find kvSet callbacks for request ID ${messageKey(
data.meta.runId,
data.key
)}. This can happen when a workflow run is resumed`
);
return true;
}
const { resolve } = setCallbacks;
resolve();
return true;
},
RESOLVE_KV_DELETE: async (data) => {
this.#logger.debug("Handling RESOLVE_KV_DELETE", data);
const deleteCallbacks = this.#kvDeleteCallbacks.get(
messageKey(data.meta.runId, data.key)
);
if (!deleteCallbacks) {
this.#logger.debug(
`Could not find kvDelete callbacks for request ID ${messageKey(
data.meta.runId,
data.key
)}. This can happen when a workflow run is resumed`
);
return true;
}
const { resolve } = deleteCallbacks;
resolve();
return true;
},
TRIGGER_WORKFLOW: async (data) => { TRIGGER_WORKFLOW: async (data) => {
this.#logger.debug("Handling TRIGGER_WORKFLOW", data); this.#logger.debug("Handling TRIGGER_WORKFLOW", data);
@@ -440,12 +542,103 @@ export class TriggerClient<TSchema extends z.ZodTypeAny> {
}; };
}; };
const kvGetFunction: KvGetFunction = async (op) => {
const result = new Promise<any>((resolve, reject) => {
this.#kvGetCallbacks.set(messageKey(data.id, op.idempotencyKey), {
resolve,
reject,
});
});
await serverRPC.send("SEND_KV_GET", {
runId: data.id,
key: op.idempotencyKey,
get: {
namespace: op.namespace,
key: op.key,
},
timestamp: String(highPrecisionTimestamp()),
});
const output = await result;
return output;
};
const kvSetFunction: KvSetFunction = async (op) => {
const result = new Promise<void>((resolve, reject) => {
this.#kvSetCallbacks.set(messageKey(data.id, op.idempotencyKey), {
resolve,
reject,
});
});
await serverRPC.send("SEND_KV_SET", {
runId: data.id,
key: op.idempotencyKey,
set: {
namespace: op.namespace,
key: op.key,
value: op.value,
},
timestamp: String(highPrecisionTimestamp()),
});
await result;
return;
};
const kvDeleteFunction: KvDeleteFunction = async (op) => {
const result = new Promise<void>((resolve, reject) => {
this.#kvDeleteCallbacks.set(
messageKey(data.id, op.idempotencyKey),
{
resolve,
reject,
}
);
});
await serverRPC.send("SEND_KV_DELETE", {
runId: data.id,
key: op.idempotencyKey,
delete: {
namespace: op.namespace,
key: op.key,
},
timestamp: String(highPrecisionTimestamp()),
});
await result;
return;
};
const ctx: TriggerContext = { const ctx: TriggerContext = {
id: data.id, id: data.id,
environment: data.meta.environment, environment: data.meta.environment,
apiKey: data.meta.apiKey, apiKey: data.meta.apiKey,
organizationId: data.meta.organizationId, organizationId: data.meta.organizationId,
isTest: data.meta.isTest, isTest: data.meta.isTest,
kv: new ContextKeyValueStorage(
`workflow:${data.meta.workflowId}`,
kvGetFunction,
kvSetFunction,
kvDeleteFunction
),
globalKv: new ContextKeyValueStorage(
`org:${data.meta.organizationId}`,
kvGetFunction,
kvSetFunction,
kvDeleteFunction
),
runKv: new ContextKeyValueStorage(
`run:${data.id}`,
kvGetFunction,
kvSetFunction,
kvDeleteFunction
),
logger: new ContextLogger(async (level, message, properties) => { logger: new ContextLogger(async (level, message, properties) => {
await serverRPC.send("SEND_LOG", { await serverRPC.send("SEND_LOG", {
runId: data.id, runId: data.id,
@@ -0,0 +1,61 @@
import { TriggerKeyValueStorage } from "./types";
export type KvSetFunction = (operation: {
key: string;
namespace: string;
idempotencyKey: string;
value: any;
}) => Promise<void>;
export type KvGetFunction = (operation: {
key: string;
namespace: string;
idempotencyKey: string;
}) => Promise<any>;
export type KvDeleteFunction = (operation: {
key: string;
namespace: string;
idempotencyKey: string;
}) => Promise<any>;
export class ContextKeyValueStorage implements TriggerKeyValueStorage {
getCount: number = 0;
setCount: number = 0;
deleteCount: number = 0;
constructor(
private namespace: string,
private onGet: KvGetFunction,
private onSet: KvSetFunction,
private onDelete: KvDeleteFunction
) {}
get<T>(key: string): Promise<T | undefined> {
const operation = {
key,
namespace: this.namespace,
idempotencyKey: `get:${this.namespace}:${key}:${this.getCount++}`,
};
return this.onGet(operation);
}
set<T>(key: string, value: T): Promise<void> {
const operation = {
key,
namespace: this.namespace,
idempotencyKey: `set:${this.namespace}:${key}:${this.setCount++}`,
value,
};
return this.onSet(operation);
}
delete(key: string): Promise<void> {
const operation = {
key,
namespace: this.namespace,
idempotencyKey: `delete:${this.namespace}:${key}:${this.deleteCount++}`,
};
return this.onDelete(operation);
}
}
+9
View File
@@ -58,6 +58,12 @@ export type TriggerFetch = <TBodySchema extends z.ZodTypeAny = z.ZodTypeAny>(
export type TriggerRunOnceCallback = (idempotencyKey: string) => Promise<any>; export type TriggerRunOnceCallback = (idempotencyKey: string) => Promise<any>;
export interface TriggerKeyValueStorage {
get<T>(key: string): Promise<T | undefined>;
set<T>(key: string, value: T): Promise<void>;
delete(key: string): Promise<void>;
}
export interface TriggerContext { export interface TriggerContext {
id: string; id: string;
environment: string; environment: string;
@@ -77,6 +83,9 @@ export interface TriggerContext {
callback: T callback: T
): Promise<Awaited<ReturnType<T>>>; ): Promise<Awaited<ReturnType<T>>>;
fetch: TriggerFetch; fetch: TriggerFetch;
kv: TriggerKeyValueStorage;
globalKv: TriggerKeyValueStorage;
runKv: TriggerKeyValueStorage;
} }
export interface TriggerLogger { export interface TriggerLogger {