From 59e16a97afffd17412e38780741293d627dad7a7 Mon Sep 17 00:00:00 2001
From: nicktrn <55853254+nicktrn@users.noreply.github.com>
Date: Wed, 25 Oct 2023 21:40:57 +0000
Subject: [PATCH] Upsert batch job function

---
 apps/webapp/app/services/worker.server.ts | 74 ++++++++++++++++++++---
 1 file changed, 65 insertions(+), 9 deletions(-)

diff --git a/apps/webapp/app/services/worker.server.ts b/apps/webapp/app/services/worker.server.ts
index d07a90e59..ffd718026 100644
--- a/apps/webapp/app/services/worker.server.ts
+++ b/apps/webapp/app/services/worker.server.ts
@@ -22,6 +22,7 @@ import { PerformTaskOperationService } from "./tasks/performTaskOperation.server
 import { ProcessCallbackTimeoutService } from "./tasks/processCallbackTimeout";
 import { ProbeEndpointService } from "./endpoints/probeEndpoint.server";
 import { PgNotifyService } from "./db/pgNotify.server";
+import { runMigrations } from "graphile-worker";
 
 const workerCatalog = {
   indexEndpoint: z.object({
@@ -122,7 +123,14 @@ if (env.NODE_ENV === "production") {
 }
 
 export async function init() {
-  await addMigrationDelayAndNotify();
+  await detectAndPrepareForMigrations();
+
+  await runMigrations({
+    connectionString: env.DATABASE_URL,
+    schema: env.WORKER_SCHEMA,
+  });
+
+  await upsertBatchJobFunction();
 
   if (env.WORKER_ENABLED === "true") {
     await workerQueue.initialize();
@@ -133,6 +141,45 @@ export async function init() {
   }
 }
 
+async function upsertBatchJobFunction() {
+  const prismaSchema = new URL(env.DATABASE_URL).searchParams.get("schema") ?? "public";
+
+  await prisma.$executeRawUnsafe(`
+    CREATE OR REPLACE FUNCTION ${prismaSchema}.add_batch_job(
+      job_key text,
+      job_key_new int,
+      task_identifier text,
+      payload json,
+      maximum_payloads int,
+      run_at timestamptz
+    ) RETURNS ${env.WORKER_SCHEMA}.jobs AS $$
+    DECLARE
+      v_job ${env.WORKER_SCHEMA}.jobs;
+    BEGIN
+      IF json_typeof(payload) IS DISTINCT FROM 'array' THEN
+        RAISE EXCEPTION 'Must only call add_batch_job with an array payload';
+      END IF;
+
+      v_job := ${env.WORKER_SCHEMA}.add_job(
+        identifier := task_identifier,
+        payload := payload,
+        run_at := run_at,
+        job_key := job_key,
+        job_key_mode := 'preserve_run_at'
+      );
+      
+      IF json_array_length(v_job.payload) >= maximum_payloads THEN
+        UPDATE jobs SET run_at = NOW() WHERE jobs.id = v_job.id RETURNING * INTO v_job;
+        -- lie that this job was just inserted so a worker picks it up ASAP
+        PERFORM pg_notify('jobs:insert', '');
+      END IF;
+
+      RETURN v_job;
+    END
+    $$ LANGUAGE plpgsql VOLATILE;
+  `);
+}
+
 async function graphileSchemaExists() {
   const schemaCount = await prisma.$executeRaw`
     SELECT schema_name FROM information_schema.schemata
@@ -142,13 +189,7 @@ async function graphileSchemaExists() {
   return schemaCount === 1;
 }
 
-/** Helper for graphile-worker v0.14.0 migration. No-op if already migrated. */
-async function addMigrationDelayAndNotify() {
-  if (!(await graphileSchemaExists())) {
-    // no schema yet, likely first start
-    return;
-  }
-
+async function getLatestMigration() {
   const migrationQueryResult = await prisma.$queryRawUnsafe(`
     SELECT id FROM ${env.WORKER_SCHEMA}.migrations
     ORDER BY id DESC LIMIT 1
@@ -160,10 +201,25 @@ async function addMigrationDelayAndNotify() {
 
   if (!migrationResults.length) {
     // no migrations applied yet
+    return -1;
+  }
+
+  return migrationResults[0].id;
+}
+
+/** Helper for graphile-worker v0.14.0 migration. No-op if already migrated. */
+async function detectAndPrepareForMigrations() {
+  if (!(await graphileSchemaExists())) {
+    // no schema yet, likely first start
     return;
   }
 
-  const latestMigration = migrationResults[0].id;
+  const latestMigration = await getLatestMigration();
+
+  if (latestMigration < 0) {
+    // no migrations found
+    return;
+  }
 
   // the first v0.14.0 migration has ID 11
   if (latestMigration > 10) {