From 0f5ae1a73090eb1fe4bb7a076100c59daa1b0082 Mon Sep 17 00:00:00 2001
From: Sam Willis <sam.willis@gmail.com>
Date: Tue, 18 Jun 2024 22:00:34 +0100
Subject: [PATCH] More extension api

---
 packages/pglite/src/interface.ts | 27 ++++++++++++++-----
 packages/pglite/src/pglite.ts    | 45 +++++++++++++++++++++++++++++---
 2 files changed, 61 insertions(+), 11 deletions(-)

diff --git a/packages/pglite/src/interface.ts b/packages/pglite/src/interface.ts
index 1bb6227d..fda850b5 100644
--- a/packages/pglite/src/interface.ts
+++ b/packages/pglite/src/interface.ts
@@ -20,10 +20,19 @@ export interface ExecProtocolOptions {
   syncToFs?: boolean;
 }
 
-export interface Extension<T = any> {
-  name: string;
-  setup: (mod: any, pg: PGliteInterface) => T;
-  init: () => void;
+export interface ExtensionSetupResult {
+  emscriptenOpts?: any;
+  namespaceObj?: any;
+  init?: () => Promise<void>;
+  close?: () => Promise<void>;
+}
+
+export interface Extension {
+  name?: string;
+  setup: (
+    pg: PGliteInterface,
+    emscriptenOpts: any
+  ) => Promise<ExtensionSetupResult>;
 }
 
 export type Extensions = {
@@ -62,9 +71,13 @@ export type PGliteInterface = {
 
 export type PGliteInterfaceExtensions<E> = E extends Extensions
   ? {
-      [K in keyof E]: ReturnType<E[K]["setup"]> extends undefined | null | void
-        ? never
-        : ReturnType<E[K]["setup"]>;
+      [K in keyof E]: Awaited<
+        ReturnType<E[K]["setup"]>
+      >["namespaceObj"] extends infer N
+        ? N extends undefined | null | void
+          ? never
+          : N
+        : never;
     }
   : {};
 
diff --git a/packages/pglite/src/pglite.ts b/packages/pglite/src/pglite.ts
index e568f142..6a21eb41 100644
--- a/packages/pglite/src/pglite.ts
+++ b/packages/pglite/src/pglite.ts
@@ -14,6 +14,8 @@ import type {
   QueryOptions,
   ExecProtocolOptions,
   PGliteInterfaceExtensions,
+  Extensions,
+  Extension,
 } from "./interface.js";
 
 // Importing the source as the built version is not ESM compatible
@@ -30,6 +32,7 @@ export class PGlite implements PGliteInterface {
   fs?: Filesystem;
   protected emp?: any;
 
+  #extensions: Extensions;
   #initStarted = false;
   #ready = false;
   #eventTarget: EventTarget;
@@ -37,6 +40,7 @@ export class PGlite implements PGliteInterface {
   #closed = false;
   #inTransaction = false;
   #relaxedDurability = false;
+  #extensionsClose: Array<() => Promise<void>> = [];
 
   #resultAccumulator: Uint8Array[] = [];
 
@@ -98,6 +102,9 @@ export class PGlite implements PGliteInterface {
       this.#resultAccumulator.push(e.detail);
     });
 
+    // Save the extensions for later use
+    this.#extensions = options.extensions ?? {};
+
     // Initialize the database, and store the promise so we can wait for it to be ready
     this.waitReady = this.#init(options ?? {});
   }
@@ -114,6 +121,7 @@ export class PGlite implements PGliteInterface {
       this.fs = await loadFs(dataDir, fsType);
     }
 
+    const extensionInitFns: Array<() => Promise<void>> = [];
     let firstRun = false;
     await new Promise<void>(async (resolve, reject) => {
       if (this.#initStarted) {
@@ -157,6 +165,23 @@ export class PGlite implements PGliteInterface {
         Event: PGEvent,
       };
 
+      // Setup extensions
+      for (const [extName, ext] of Object.entries(this.#extensions)) {
+        const extRet = await ext.setup(this, emscriptenOpts);
+        if (extRet.emscriptenOpts) {
+          emscriptenOpts = extRet.emscriptenOpts;
+        }
+        if (extRet.namespaceObj) {
+          (this as any)[extName] = extRet.namespaceObj;
+        }
+        if (extRet.init) {
+          extensionInitFns.push(extRet.init);
+        }
+        if (extRet.close) {
+          this.#extensionsClose.push(extRet.close);
+        }
+      }
+
       emscriptenOpts = await this.fs!.emscriptenOpts(emscriptenOpts);
       const emp = await EmPostgresFactory(emscriptenOpts);
       this.emp = emp;
@@ -168,6 +193,11 @@ export class PGlite implements PGliteInterface {
     await this.#runExec(`
       SET search_path TO public;
     `);
+
+    // Init extensions
+    for (const initFn of extensionInitFns) {
+      await initFn();
+    }
   }
 
   /**
@@ -220,6 +250,13 @@ export class PGlite implements PGliteInterface {
   async close() {
     await this.#checkReady();
     this.#closing = true;
+
+    // Close all extensions
+    for (const closeFn of this.#extensionsClose) {
+      await closeFn();
+    }
+
+    // Close the database
     await new Promise<void>(async (resolve, reject) => {
       try {
         await this.execProtocol(serialize.end());
@@ -519,8 +556,9 @@ export class PGlite implements PGliteInterface {
   }
 
   /**
-   * Create a new PGlite instance with extensions on the interface
-   * (The main constructor does enable extensions, but does not add them to the interface)
+   * Create a new PGlite instance with extensions on the Typescript interface
+   * (The main constructor does enable extensions, however due to the limitations
+   * of Typescript, the extensions are not available on the instance interface)
    * @param dataDir The directory to store the database files
    *                Prefix with idb:// to use indexeddb filesystem in the browser
    *                Use memory:// to use in-memory filesystem
@@ -528,9 +566,8 @@ export class PGlite implements PGliteInterface {
    * @returns A new PGlite instance with extensions
    */
   static withExtensions<O extends PGliteOptions>(
-    dataDir?: string,
     options?: O
   ): PGlite & PGliteInterfaceExtensions<O["extensions"]> {
-    return new PGlite(dataDir, options) as any;
+    return new PGlite(options) as any;
   }
 }