From 0f3168d7b5bd03e0e06afa7ecb349fc7c6993411 Mon Sep 17 00:00:00 2001
From: Benjie Gillam <benjie@jemjie.com>
Date: Mon, 4 Dec 2023 11:43:33 +0000
Subject: [PATCH] Use PgContextPlugin in example schema runner

---
 .../dataplan-pg/scripts/runExampleSchema.mjs  | 39 +++++++++++--------
 grafast/dataplan-pg/src/adaptors/pg.ts        | 24 ++++++++----
 grafast/dataplan-pg/src/interfaces.ts         |  6 +++
 grafast/grafast/src/interfaces.ts             |  1 +
 4 files changed, 47 insertions(+), 23 deletions(-)

diff --git a/grafast/dataplan-pg/scripts/runExampleSchema.mjs b/grafast/dataplan-pg/scripts/runExampleSchema.mjs
index 2569f5221b..b6a6820ff0 100755
--- a/grafast/dataplan-pg/scripts/runExampleSchema.mjs
+++ b/grafast/dataplan-pg/scripts/runExampleSchema.mjs
@@ -5,32 +5,39 @@ import { grafast } from "grafast";
 import { isAsyncIterable } from "iterall";
 import JSON5 from "json5";
 import { strict as assert } from "node:assert";
-import pg from "pg";
 
-import { createWithPgClient } from "../dist/adaptors/pg.js";
 import { schema } from "./exampleSchemaExport.mjs";
+import { resolvePresets } from "graphile-config";
+import { PgContextPlugin } from "@dataplan/pg";
+import { makePgService } from "@dataplan/pg/adaptors/pg";
 
-const pool = new pg.Pool({
-  connectionString: process.env.TEST_DATABASE_URL || "graphile_crystal",
-});
+const connectionString =
+  process.env.TEST_DATABASE_URL || "postgres:///graphile_crystal";
+const pgService = makePgService({ connectionString });
+const preset = {
+  plugins: [PgContextPlugin],
+  pgServices: [pgService],
+  grafast: {
+    context() {
+      return {
+        pgSettings: {
+          timezone: "UTC",
+        },
+      };
+    },
+  },
+};
+const resolvedPreset = resolvePresets([preset]);
 
 async function runTestQuery(basePath) {
   const source = await readFile(`${basePath}.test.graphql`, "utf8");
   const expectedData = JSON5.parse(await readFile(`${basePath}.json5`, "utf8"));
 
-  const withPgClient = createWithPgClient({
-    pool,
-  });
-
   const result = await grafast({
     schema,
     source,
-    contextValue: {
-      pgSettings: {
-        timezone: "UTC",
-      },
-      withPgClient,
-    },
+    resolvedPreset,
+    requestContext: {},
   });
   const operationType = "query";
 
@@ -158,5 +165,5 @@ try {
     await runTestQuery(basePath);
   }
 } finally {
-  pool.end();
+  await pgService.release?.();
 }
diff --git a/grafast/dataplan-pg/src/adaptors/pg.ts b/grafast/dataplan-pg/src/adaptors/pg.ts
index 309f59f3a5..8eb92a783c 100644
--- a/grafast/dataplan-pg/src/adaptors/pg.ts
+++ b/grafast/dataplan-pg/src/adaptors/pg.ts
@@ -730,11 +730,12 @@ export function makePgService(
     );
   }
   const Pool = pg.Pool ?? (pg as any).default?.Pool;
-  const pool =
-    options.pool ??
-    new Pool({
-      connectionString,
-    });
+  const releasers: (() => void | PromiseLike<void>)[] = [];
+  let pool = options.pool;
+  if (!pool) {
+    pool = new Pool({ connectionString });
+    releasers.push(() => pool!.end());
+  }
   if (!options.pool) {
     // If you pass your own pool, you're responsible for doing this yourself
     pool.on("connect", (client) => {
@@ -746,8 +747,11 @@ export function makePgService(
       console.error("Client error (in pool)", e);
     });
   }
-  const pgSubscriber =
-    options.pgSubscriber ?? (pubsub ? new PgSubscriber(pool) : null);
+  let pgSubscriber = options.pgSubscriber ?? null;
+  if (!pgSubscriber && pubsub) {
+    pgSubscriber = new PgSubscriber(pool);
+    releasers.push(() => pgSubscriber!.release?.());
+  }
   const service: GraphileConfig.PgServiceConfiguration = {
     name,
     schemas: Array.isArray(schemas) ? schemas : [schemas ?? "public"],
@@ -762,6 +766,12 @@ export function makePgService(
       pool,
       superuserConnectionString,
     },
+    async release() {
+      // Release in reverse order
+      for (const releaser of [...releasers].reverse()) {
+        await releaser();
+      }
+    },
   };
   return service;
 }
diff --git a/grafast/dataplan-pg/src/interfaces.ts b/grafast/dataplan-pg/src/interfaces.ts
index c5bcb2132f..15df0457a9 100644
--- a/grafast/dataplan-pg/src/interfaces.ts
+++ b/grafast/dataplan-pg/src/interfaces.ts
@@ -453,6 +453,12 @@ declare global {
         Grafast.Context & object,
         GrafastSubscriber<any> | null | undefined
       >;
+
+      /**
+       * Call this when you no longer need this service configuration any more;
+       * releases any created resources (e.g. connection pools).
+       */
+      release?: () => void | PromiseLike<void>;
     }
 
     interface Preset {
diff --git a/grafast/grafast/src/interfaces.ts b/grafast/grafast/src/interfaces.ts
index c9aab07a63..cf5afa5f30 100644
--- a/grafast/grafast/src/interfaces.ts
+++ b/grafast/grafast/src/interfaces.ts
@@ -848,6 +848,7 @@ export type GrafastSubscriber<
   subscribe<TTopic extends keyof TTopics = keyof TTopics>(
     topic: TTopic,
   ): PromiseOrDirect<AsyncIterableIterator<TTopics[TTopic]>>;
+  release?(): PromiseOrDirect<void>;
 };
 
 /**