feat: expression-based concurrency keys (#889)

* feat: expression-based concurrency keys

* fix: build

* fix: typos

* fix: gen

* fix: migration

* fix: remove print statements

* fix: reassignment bugs, retries on closed transport, pr review

sql/constraints.sql

@@ -10,3 +10,5 @@ ALTER TABLE "InternalQueueItem" ADD CONSTRAINT "InternalQueueItem_priority_check
 CREATE INDEX CONCURRENTLY IF NOT EXISTS "StepRun_jobRunId_status_tenantId_idx"
 ON "StepRun" ("jobRunId", "status", "tenantId")
 WHERE "status" = 'PENDING';
+
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "WorkflowRun_parentStepRunId" ON "WorkflowRun"("parentStepRunId" ASC);

sql/migrations/20240918162532_v0.45.0.sql

@@ -0,0 +1,16 @@
+-- atlas:txmode none
+
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "WorkflowRun_parentStepRunId" ON "WorkflowRun"("parentStepRunId" ASC);
+
+-- Add value to enum type: "StepRunEventReason"
+ALTER TYPE "StepRunEventReason" ADD VALUE 'WORKFLOW_RUN_GROUP_KEY_SUCCEEDED';
+-- Add value to enum type: "StepRunEventReason"
+ALTER TYPE "StepRunEventReason" ADD VALUE 'WORKFLOW_RUN_GROUP_KEY_FAILED';
+-- Add value to enum type: "InternalQueue"
+ALTER TYPE "InternalQueue" ADD VALUE 'WORKFLOW_RUN_UPDATE';
+-- Modify "StepRunEvent" table
+ALTER TABLE "StepRunEvent" DROP CONSTRAINT "StepRunEvent_stepRunId_fkey", ALTER COLUMN "stepRunId" DROP NOT NULL, ADD COLUMN "workflowRunId" uuid NULL;
+-- Create index "StepRunEvent_workflowRunId_idx" to table: "StepRunEvent"
+CREATE INDEX "StepRunEvent_workflowRunId_idx" ON "StepRunEvent" ("workflowRunId");
+-- Modify "WorkflowConcurrency" table
+ALTER TABLE "WorkflowConcurrency" ADD COLUMN "concurrencyGroupExpression" text NULL;

sql/migrations/atlas.sum

@@ -1,4 +1,4 @@
-h1:GONXnb4KX2gX139WgjSkVzSS2yGkT+MXqmCtgO3ulMg=
+h1:TwaCvaPgsVfs199JWuoIa/1ipf3RIgAGfH76Wzm/4dw=
 20240115180414_init.sql h1:Ef3ZyjAHkmJPdGF/dEWCahbwgcg6uGJKnDxW2JCRi2k=
 20240122014727_v0_6_0.sql h1:o/LdlteAeFgoHJ3e/M4Xnghqt9826IE/Y/h0q95Acuo=
 20240126235456_v0_7_0.sql h1:KiVzt/hXgQ6esbdC6OMJOOWuYEXmy1yeCpmsVAHTFKs=
@@ -56,3 +56,4 @@ h1:GONXnb4KX2gX139WgjSkVzSS2yGkT+MXqmCtgO3ulMg=
 20240911124017_v0.44.6.sql h1:7MlkRZpRUra7GUPSYfFzqDSZGM1SKlP1j81xMHyptqc=
 20240911201831_v0.44.7.sql h1:rBiOx3qNLvl9lCedK6dlf6yTX3y+xCcmUJNEXPG8YYc=
 20240916115647_v0.44.8.sql h1:1D5fmikdME61lwN1/EbBTmGcKzpqESZUmfoHvmOeZl4=
+20240918162532_v0.45.0.sql h1:pOhBg/58SnD8qBpDx9OdxTFfjRxEanJrJF7rLEvUFSY=

sql/schema/schema.sql

@@ -2,7 +2,7 @@
 CREATE TYPE "ConcurrencyLimitStrategy" AS ENUM ('CANCEL_IN_PROGRESS', 'DROP_NEWEST', 'QUEUE_NEWEST', 'GROUP_ROUND_ROBIN');
 
 -- CreateEnum
-CREATE TYPE "InternalQueue" AS ENUM ('WORKER_SEMAPHORE_COUNT', 'STEP_RUN_UPDATE');
+CREATE TYPE "InternalQueue" AS ENUM ('WORKER_SEMAPHORE_COUNT', 'STEP_RUN_UPDATE', 'WORKFLOW_RUN_UPDATE');
 
 -- CreateEnum
 CREATE TYPE "InviteLinkStatus" AS ENUM ('PENDING', 'ACCEPTED', 'REJECTED');
@@ -20,7 +20,7 @@ CREATE TYPE "LimitResource" AS ENUM ('WORKFLOW_RUN', 'EVENT', 'WORKER', 'CRON',
 CREATE TYPE "LogLineLevel" AS ENUM ('DEBUG', 'INFO', 'WARN', 'ERROR');
 
 -- CreateEnum
-CREATE TYPE "StepRunEventReason" AS ENUM ('REQUEUED_NO_WORKER', 'REQUEUED_RATE_LIMIT', 'SCHEDULING_TIMED_OUT', 'ASSIGNED', 'STARTED', 'FINISHED', 'FAILED', 'RETRYING', 'CANCELLED', 'TIMED_OUT', 'REASSIGNED', 'SLOT_RELEASED', 'TIMEOUT_REFRESHED', 'RETRIED_BY_USER', 'SENT_TO_WORKER');
+CREATE TYPE "StepRunEventReason" AS ENUM ('REQUEUED_NO_WORKER', 'REQUEUED_RATE_LIMIT', 'SCHEDULING_TIMED_OUT', 'ASSIGNED', 'STARTED', 'FINISHED', 'FAILED', 'RETRYING', 'CANCELLED', 'TIMED_OUT', 'REASSIGNED', 'SLOT_RELEASED', 'TIMEOUT_REFRESHED', 'RETRIED_BY_USER', 'SENT_TO_WORKER', 'WORKFLOW_RUN_GROUP_KEY_SUCCEEDED', 'WORKFLOW_RUN_GROUP_KEY_FAILED');
 
 -- CreateEnum
 CREATE TYPE "StepRunEventSeverity" AS ENUM ('INFO', 'WARNING', 'CRITICAL');
@@ -396,12 +396,13 @@ CREATE TABLE "StepRunEvent" (
     "id" BIGSERIAL NOT NULL,
     "timeFirstSeen" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
     "timeLastSeen" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
-    "stepRunId" UUID NOT NULL,
+    "stepRunId" UUID,
     "reason" "StepRunEventReason" NOT NULL,
     "severity" "StepRunEventSeverity" NOT NULL,
     "message" TEXT NOT NULL,
     "count" INTEGER NOT NULL,
-    "data" JSONB
+    "data" JSONB,
+    "workflowRunId" UUID
 );
 
 -- CreateTable
@@ -736,6 +737,7 @@ CREATE TABLE "WorkflowConcurrency" (
     "getConcurrencyGroupId" UUID,
     "maxRuns" INTEGER NOT NULL DEFAULT 1,
     "limitStrategy" "ConcurrencyLimitStrategy" NOT NULL DEFAULT 'CANCEL_IN_PROGRESS',
+    "concurrencyGroupExpression" TEXT,
 
     CONSTRAINT "WorkflowConcurrency_pkey" PRIMARY KEY ("id")
 );
@@ -1074,6 +1076,9 @@ CREATE UNIQUE INDEX "StepRunEvent_id_key" ON "StepRunEvent"("id" ASC);
 -- CreateIndex
 CREATE INDEX "StepRunEvent_stepRunId_idx" ON "StepRunEvent"("stepRunId" ASC);
 
+-- CreateIndex
+CREATE INDEX "StepRunEvent_workflowRunId_idx" ON "StepRunEvent"("workflowRunId" ASC);
+
 -- CreateIndex
 CREATE UNIQUE INDEX "StepRunResultArchive_id_key" ON "StepRunResultArchive"("id" ASC);
 
@@ -1422,9 +1427,6 @@ ALTER TABLE "StepRun" ADD CONSTRAINT "StepRun_tickerId_fkey" FOREIGN KEY ("ticke
 -- AddForeignKey
 ALTER TABLE "StepRun" ADD CONSTRAINT "StepRun_workerId_fkey" FOREIGN KEY ("workerId") REFERENCES "Worker"("id") ON DELETE SET NULL ON UPDATE CASCADE;
 
--- AddForeignKey
-ALTER TABLE "StepRunEvent" ADD CONSTRAINT "StepRunEvent_stepRunId_fkey" FOREIGN KEY ("stepRunId") REFERENCES "StepRun"("id") ON DELETE CASCADE ON UPDATE CASCADE;
-
 -- AddForeignKey
 ALTER TABLE "StepRunResultArchive" ADD CONSTRAINT "StepRunResultArchive_stepRunId_fkey" FOREIGN KEY ("stepRunId") REFERENCES "StepRun"("id") ON DELETE CASCADE ON UPDATE CASCADE;
 
@@ -1629,3 +1631,5 @@ ALTER TABLE "InternalQueueItem" ADD CONSTRAINT "InternalQueueItem_priority_check
 CREATE INDEX CONCURRENTLY IF NOT EXISTS "StepRun_jobRunId_status_tenantId_idx"
 ON "StepRun" ("jobRunId", "status", "tenantId")
 WHERE "status" = 'PENDING';
+
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "WorkflowRun_parentStepRunId" ON "WorkflowRun"("parentStepRunId" ASC);