[validation-only] verify loopback-bound PGlite gateway and per-run credentials

src/generate.ts

@@ -1,4 +1,8 @@
 import * as zg from "zapatos/generate"
+import { spawn } from "node:child_process"
+import { randomBytes } from "node:crypto"
+import { createRequire } from "node:module"
+import type { Socket } from "node:net"
 import {
   getConnectionStringFromEnv,
   getPgConnectionFromEnv,
@@ -8,6 +12,42 @@ import { dumpTree } from "pg-schema-dump"
 import path from "path"
 import { migrate } from "./migrate"
 
+const dumpTreeModulePath = createRequire(__filename).resolve("pg-schema-dump")
+
+const dumpTreeWithConnection = async (
+  options: Parameters<typeof dumpTree>[0],
+  connectionString: string,
+) => {
+  const script = `const dumpTree = require(process.argv[1]).dumpTree;
+dumpTree(JSON.parse(process.argv[2])).catch((error) => {
+  console.error(error)
+  process.exit(1)
+})`
+
+  await new Promise<void>((resolve, reject) => {
+    const child = spawn(
+      process.execPath,
+      ["-e", script, dumpTreeModulePath, JSON.stringify(options)],
+      {
+        cwd: process.cwd(),
+        env: {
+          ...process.env,
+          POSTGRES_URI: connectionString,
+          PG_URI: connectionString,
+          DATABASE_URL: connectionString,
+        },
+        stdio: "inherit",
+      },
+    )
+
+    child.once("error", reject)
+    child.once("exit", (code, signal) => {
+      if (code === 0) resolve()
+      else reject(new Error(`pg-schema-dump exited with ${signal ?? code}`))
+    })
+  })
+}
+
 export const generate = async ({
   schemas,
   defaultDatabase,
@@ -27,31 +67,33 @@ export const generate = async ({
     const net = await import("node:net")
 
     const db = new PGlite()
-
-    await migrate({
-      client: db as any,
-      migrationsDir,
-      defaultDatabase,
-      cwd: process.cwd(),
-      schemas,
-    })
+    const gatewayPassword = randomBytes(24).toString("hex")
+    const sockets = new Set<Socket>()
+    const errors: unknown[] = []
 
     const server = net.createServer(async (socket) => {
+      sockets.add(socket)
+      socket.once("close", () => sockets.delete(socket))
       const connection = await fromNodeSocket(socket, {
         serverVersion: "16.3 (PGlite)",
         auth: {
           method: "password",
           validateCredentials: ({ username, password }: any) =>
-            username === "postgres" && password === "postgres",
-          getClearTextPassword: () => "postgres",
+            username === "postgres" && password === gatewayPassword,
+          getClearTextPassword: () => gatewayPassword,
         },
         async onStartup() {
           await (db as any).waitReady
         },
         async onMessage(data: Uint8Array, { isAuthenticated }: any) {
           if (!isAuthenticated) return
           try {
-            const { data: responseData } = await (db as any).execProtocol(data)
+            const { data: responseData } = await (db as any).execProtocol(
+              data,
+              {
+                throwOnError: false,
+              },
+            )
             return responseData
           } catch {
             return undefined
@@ -60,32 +102,64 @@ export const generate = async ({
       })
     })
 
-    await new Promise<void>((resolve) => server.listen(0, resolve))
-    const port = (server.address() as any).port
-    const connectionString = `postgres://postgres:postgres@127.0.0.1:${port}/postgres`
-
-    const prevDbUrl = process.env.DATABASE_URL
-    process.env.DATABASE_URL = connectionString
+    try {
+      await migrate({
+        client: db as any,
+        migrationsDir,
+        defaultDatabase,
+        cwd: process.cwd(),
+        schemas,
+      })
 
-    await zg.generate({
-      db: {
-        connectionString,
-      },
-      schemas: Object.fromEntries(
-        schemas.map((s) => [s, { include: "*", exclude: [] }]),
-      ),
-      outDir: dbDir,
-    })
+      await new Promise<void>((resolve, reject) => {
+        server.once("error", reject)
+        server.listen(0, "127.0.0.1", () => {
+          server.removeListener("error", reject)
+          resolve()
+        })
+      })
+      const port = (server.address() as any).port
+      const connectionString = `postgres://postgres:${gatewayPassword}@127.0.0.1:${port}/postgres`
 
-    await dumpTree({
-      targetDir: path.join(dbDir, "structure"),
-      defaultDatabase: "postgres",
-      schemas,
-    })
+      await zg.generate({
+        db: {
+          connectionString,
+        },
+        schemas: Object.fromEntries(
+          schemas.map((s) => [s, { include: "*", exclude: [] }]),
+        ),
+        outDir: dbDir,
+      })
 
-    server.close()
-    if (prevDbUrl === undefined) delete process.env.DATABASE_URL
-    else process.env.DATABASE_URL = prevDbUrl
+      await dumpTreeWithConnection(
+        {
+          targetDir: path.join(dbDir, "structure"),
+          defaultDatabase: "postgres",
+          schemas,
+        },
+        connectionString,
+      )
+    } catch (error) {
+      errors.push(error)
+    } finally {
+      try {
+        await new Promise<void>((resolve) => {
+          server.close(() => resolve())
+          for (const socket of sockets) socket.destroy()
+        })
+      } catch (error) {
+        errors.push(error)
+      }
+      try {
+        await db.close()
+      } catch (error) {
+        errors.push(error)
+      }
+    }
+    if (errors.length === 1) throw errors[0]
+    if (errors.length > 1) {
+      throw new AggregateError(errors, "PGlite generation and cleanup failed")
+    }
     return
   }
 

src/get-project-context.ts

@@ -1,5 +1,6 @@
 import fs from "fs"
 import path from "path"
+import { pathToFileURL } from "node:url"
 import { PgstrapConfig } from "./define-config"
 
 export interface Context extends PgstrapConfig {
@@ -13,7 +14,9 @@ export const getProjectContext = async (): Promise<Context> => {
     )
   }
 
-  const config = await import(path.join(process.cwd(), "pgstrap.config.js"))
+  const config = await import(
+    pathToFileURL(path.join(process.cwd(), "pgstrap.config.js")).href
+  )
 
   return {
     cwd: process.cwd(),

src/init.ts

@@ -16,7 +16,7 @@ export const initPgstrap = async (ctx: Pick<Context, "cwd">) => {
 
   pkg.scripts["db:migrate"] = "pgstrap migrate"
   pkg.scripts["db:reset"] = "pgstrap reset"
-  pkg.scripts["db:generate"] = "pgstrap generate"
+  pkg.scripts["db:generate"] = "pgstrap generate --pglite"
   pkg.scripts["db:create-migration"] = "pgstrap create-migration"
 
   if (!pkg.devDependencies) pkg.devDependencies = {}
@@ -26,9 +26,11 @@ export const initPgstrap = async (ctx: Pick<Context, "cwd">) => {
   }
 
   if (!fs.existsSync(path.join(cwd, "pgstrap.config.js"))) {
+    const configExport =
+      pkg.type === "module" ? "export default " : "module.exports = "
     fs.writeFileSync(
       path.join(cwd, "pgstrap.config.js"),
-      `module.exports = ${JSON.stringify(
+      `${configExport}${JSON.stringify(
         {
           defaultDatabase: "my_service_name",
           schemas: ["public"],

tests/generate.cleanup.fixture.ts

@@ -0,0 +1,147 @@
+import { expect, mock, test } from "bun:test"
+import { PGlite } from "@electric-sql/pglite"
+import * as net from "node:net"
+import fs from "node:fs"
+import os from "node:os"
+import path from "node:path"
+
+const nativeNet = { ...net }
+const databases: PGlite[] = []
+const servers: net.Server[] = []
+const sockets: net.Socket[] = []
+let closeFailure: Error | undefined
+
+class TrackedPGlite extends PGlite {
+  constructor() {
+    super()
+    databases.push(this)
+  }
+
+  override async close() {
+    if (closeFailure) throw closeFailure
+    await super.close()
+  }
+}
+
+mock.module("@electric-sql/pglite", () => ({ PGlite: TrackedPGlite }))
+mock.module("node:net", () => ({
+  ...nativeNet,
+  createServer: (listener: (socket: net.Socket) => void) => {
+    const server = nativeNet.createServer(listener)
+    servers.push(server)
+    server.on("connection", (socket) => sockets.push(socket))
+    server.once("listening", () => {
+      const address = server.address() as net.AddressInfo
+      const idleClient = nativeNet.createConnection({
+        host: "127.0.0.1",
+        port: address.port,
+      })
+      idleClient.on("error", () => {})
+      sockets.push(idleClient)
+    })
+    return server
+  },
+}))
+
+const { generate } = await import("../src/generate")
+
+async function isolated(
+  migration: string,
+  run: (options: Parameters<typeof generate>[0]) => Promise<void>,
+) {
+  const tmp = fs.mkdtempSync(path.join(os.tmpdir(), "pgstrap-cleanup-"))
+  const migrationsDir = path.join(tmp, "migrations")
+  fs.mkdirSync(migrationsDir)
+  fs.writeFileSync(path.join(migrationsDir, "001_probe.js"), migration)
+  databases.length = 0
+  servers.length = 0
+  sockets.length = 0
+  try {
+    await run({
+      schemas: ["public"],
+      defaultDatabase: "postgres",
+      dbDir: path.join(tmp, "db"),
+      migrationsDir,
+      pglite: true,
+    })
+  } finally {
+    closeFailure = undefined
+    for (const socket of sockets) socket.destroy()
+    for (const server of servers) {
+      if (server.listening) {
+        await new Promise<void>((resolve) => server.close(() => resolve()))
+      }
+    }
+    for (const database of databases) {
+      if (!database.closed) await database.close()
+    }
+    fs.rmSync(tmp, { recursive: true, force: true })
+  }
+}
+
+test("generation closes its PGlite database and idle gateway connections", async () => {
+  await isolated(
+    "exports.up = (pgm) => pgm.createTable('probe', { id: 'id' })",
+    async (options) => {
+      await generate(options)
+      const remainingConnections = await new Promise<number>(
+        (resolve, reject) =>
+          servers[0].getConnections((error, count) =>
+            error ? reject(error) : resolve(count),
+          ),
+      )
+      expect({
+        databaseClosed: databases[0].closed,
+        gatewayListening: servers[0].listening,
+        remainingConnections,
+      }).toEqual({
+        databaseClosed: true,
+        gatewayListening: false,
+        remainingConnections: 0,
+      })
+    },
+  )
+})
+
+test("migration failures close the temporary PGlite database", async () => {
+  await isolated(
+    "exports.up = () => { throw new Error('expected migration failure') }",
+    async (options) => {
+      await expect(generate(options)).rejects.toThrow(
+        "expected migration failure",
+      )
+      expect(databases[0].closed).toBe(true)
+      expect(servers.every((server) => !server.listening)).toBe(true)
+    },
+  )
+})
+
+test("generation and cleanup failures both remain available to the caller", async () => {
+  await isolated(
+    "exports.up = () => { throw new Error('expected migration failure') }",
+    async (options) => {
+      closeFailure = new Error("expected cleanup failure")
+      let failure: unknown
+      try {
+        await generate(options)
+      } catch (error) {
+        failure = error
+      }
+      expect(failure).toBeInstanceOf(AggregateError)
+      const errors = (failure as AggregateError).errors
+      expect(errors).toHaveLength(2)
+      expect(errors[0].message).toBe("expected migration failure")
+      expect(errors[1]).toBe(closeFailure)
+    },
+  )
+})
+
+test("cleanup failure rejects an otherwise successful generation", async () => {
+  await isolated(
+    "exports.up = (pgm) => pgm.createTable('probe', { id: 'id' })",
+    async (options) => {
+      closeFailure = new Error("expected cleanup failure")
+      await expect(generate(options)).rejects.toBe(closeFailure)
+    },
+  )
+})

tests/generate.cleanup.test.ts

@@ -0,0 +1,17 @@
+import { spawnSync } from "node:child_process"
+import { fileURLToPath } from "node:url"
+import { expect, test } from "bun:test"
+
+test("PGlite resources are released after generation and migration errors", () => {
+  const fixture = fileURLToPath(
+    new URL("./generate.cleanup.fixture.ts", import.meta.url),
+  )
+  const result = spawnSync(process.execPath, ["test", fixture], {
+    encoding: "utf8",
+    timeout: 30000,
+  })
+
+  process.stdout.write(result.stdout)
+  process.stderr.write(result.stderr)
+  expect(result.status, `${result.stdout}\n${result.stderr}`).toBe(0)
+}, 30000)