From a3e1e16804446ddf5579fc995228884fe2aef0a4 Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Thu, 1 Aug 2024 11:09:05 +0100
Subject: [PATCH] Fix close and add support for new stage 3 ECMAScript explicit
 resource management (#133)

* Fix db.close #132

* Add support for Stage 3 ECMAScript Explicit Resource Management #101
---
 packages/pglite/src/pglite.ts         | 37 ++++++++++++++++++---------
 packages/pglite/src/worker/index.ts   | 11 +++++++-
 packages/pglite/tests/basic.test.js   | 20 +++++++++++++++
 packages/pglite/tests/targets/base.js | 13 ++++++++++
 4 files changed, 68 insertions(+), 13 deletions(-)

diff --git a/packages/pglite/src/pglite.ts b/packages/pglite/src/pglite.ts
index 46372b90..31384b07 100644
--- a/packages/pglite/src/pglite.ts
+++ b/packages/pglite/src/pglite.ts
@@ -31,7 +31,7 @@ import {
   NotificationResponseMessage,
 } from "pg-protocol/dist/messages.js";
 
-export class PGlite implements PGliteInterface {
+export class PGlite implements PGliteInterface, AsyncDisposable {
   fs?: Filesystem;
   protected mod?: PostgresMod;
 
@@ -344,22 +344,35 @@ export class PGlite implements PGliteInterface {
     }
 
     // Close the database
-    await new Promise<void>(async (resolve, reject) => {
-      try {
-        await this.execProtocol(serialize.end());
-      } catch (e) {
-        const err = e as { name: string; status: number };
-        if (err.name === "ExitStatus" && err.status === 0) {
-          resolve();
-        } else {
-          reject(e);
-        }
+    try {
+      await this.execProtocol(serialize.end());
+    } catch (e) {
+      const err = e as { name: string; status: number };
+      if (err.name === "ExitStatus" && err.status === 0) {
+        // Database closed successfully
+        // An earlier build of PGlite would throw an error here when closing
+        // leaving this here for now. I believe it was a bug in Emscripten.
+      } else {
+        throw e;
       }
-    });
+    }
+
+    // Close the filesystem
+    await this.fs!.close();
+
     this.#closed = true;
     this.#closing = false;
   }
 
+  /**
+   * Close the database when the object exits scope
+   * Stage 3 ECMAScript Explicit Resource Management
+   * https://www.typescriptlang.org/docs/handbook/release-notes/typescript-5-2.html#using-declarations-and-explicit-resource-management
+   */
+  async [Symbol.asyncDispose]() {
+    await this.close();
+  }
+
   /**
    * Execute a single SQL statement
    * This uses the "Extended Query" postgres wire protocol message.
diff --git a/packages/pglite/src/worker/index.ts b/packages/pglite/src/worker/index.ts
index 15f205b2..fd724121 100644
--- a/packages/pglite/src/worker/index.ts
+++ b/packages/pglite/src/worker/index.ts
@@ -16,7 +16,7 @@ export type PGliteWorkerOptions = PGliteOptions & {
   id?: string;
 };
 
-export class PGliteWorker implements PGliteInterface {
+export class PGliteWorker implements PGliteInterface, AsyncDisposable {
   #initPromise: Promise<void>;
   #debug: DebugLevel = 0;
 
@@ -300,6 +300,15 @@ export class PGliteWorker implements PGliteInterface {
     this.#workerProcess.terminate();
   }
 
+  /**
+   * Close the database when the object exits scope
+   * Stage 3 ECMAScript Explicit Resource Management
+   * https://www.typescriptlang.org/docs/handbook/release-notes/typescript-5-2.html#using-declarations-and-explicit-resource-management
+   */
+  async [Symbol.asyncDispose]() {
+    await this.close();
+  }
+
   /**
    * Execute a single SQL statement
    * This uses the "Extended Query" postgres wire protocol message.
diff --git a/packages/pglite/tests/basic.test.js b/packages/pglite/tests/basic.test.js
index e46fe0db..01feccf3 100644
--- a/packages/pglite/tests/basic.test.js
+++ b/packages/pglite/tests/basic.test.js
@@ -371,3 +371,23 @@ test("basic copy to/from blob", async (t) => {
     affectedRows: 0,
   });
 });
+
+test("basic close", async (t) => {
+  const db = new PGlite();
+  await db.query(`
+    CREATE TABLE IF NOT EXISTS test (
+      id SERIAL PRIMARY KEY,
+      name TEXT
+    );
+  `);
+  await db.query("INSERT INTO test (name) VALUES ('test');");
+  await db.close();
+  await t.throwsAsync(
+    async () => {
+      await db.query("SELECT * FROM test;");
+    },
+    {
+      message: "PGlite is closed",
+    }
+  );
+});
diff --git a/packages/pglite/tests/targets/base.js b/packages/pglite/tests/targets/base.js
index 7239c005..df3ad5ea 100644
--- a/packages/pglite/tests/targets/base.js
+++ b/packages/pglite/tests/targets/base.js
@@ -132,6 +132,19 @@ export function tests(env, dbFilename, target) {
     });
   });
 
+  test.serial(`targets ${target} close`, async (t) => {
+    const err = await evaluate(async () => {
+      try {
+        await db.close();
+      } catch (e) {
+        console.error(e);
+        return e.message;
+      }
+      return null;
+    });
+    t.is(err, null);
+  });
+
   if (dbFilename === "memory://") {
     // Skip the rest of the tests for memory:// as it's not persisted
     return;