db:generate works without an external Postgres via auto PGlite (#2)

README.md

@@ -55,7 +55,7 @@ npm install pgstrap --save-dev
 
 - `npm run db:migrate` - Run pending migrations
 - `npm run db:reset` - Drop and recreate the database, then run all migrations
-- `npm run db:generate` - Generate types and structure dumps. Use `pgstrap generate --pglite` to run migrations against an in-memory PGlite instance.
+- `npm run db:generate` - Generate types and structure dumps. When no Postgres connection env is configured, pgstrap runs migrations against an in-memory PGlite instance automatically.
 - `npm run db:create-migration` - Create a new migration file
 
 ### Configuration

package.json

@@ -8,6 +8,7 @@
   "scripts": {
     "test": "bun test",
     "build": "tsup ./src/index.ts ./src/cli.ts --outDir ./dist --dts --sourcemap inline",
+    "db:generate": "bun ./src/cli.ts generate",
     "format": "biome format . --write",
     "format:check": "biome format ."
   },

src/cli.ts

@@ -29,7 +29,7 @@ import { getProjectContext } from "./get-project-context"
     process.exit(0)
   })
   .command("migrate", "migrates the database", {}, async () => {
-    migrate(await getProjectContext())
+    await migrate(await getProjectContext())
   })
   .command(
     "generate",
@@ -38,7 +38,7 @@ import { getProjectContext } from "./get-project-context"
       yargs.option("pglite", { type: "boolean", default: false })
     },
     async (argv) => {
-      generate({ ...(await getProjectContext()), pglite: !!argv.pglite })
+      await generate({ ...(await getProjectContext()), pglite: !!argv.pglite })
     },
   )
   .parse()

src/generate.ts

@@ -1,33 +1,82 @@
 import * as zg from "zapatos/generate"
-import {
-  getConnectionStringFromEnv,
-  getPgConnectionFromEnv,
-} from "pg-connection-from-env"
+import { getConnectionStringFromEnv } from "pg-connection-from-env"
 import { Context } from "./get-project-context"
 import { dumpTree } from "pg-schema-dump"
 import path from "path"
 import { migrate } from "./migrate"
 
-export const generate = async ({
+const postgresEnvKeys = [
+  "DATABASE_URL",
+  "PGHOST",
+  "PGPORT",
+  "PGDATABASE",
+  "PGUSER",
+  "PGPASSWORD",
+]
+
+export const hasExternalDatabaseConfig = (
+  env: Record<string, string | undefined> = process.env,
+) => postgresEnvKeys.some((key) => Boolean(env[key]))
+
+let pgliteEnvLock = Promise.resolve()
+
+const withPgliteEnvLock = async <T>(fn: () => Promise<T>) => {
+  const previous = pgliteEnvLock
+  let release!: () => void
+  pgliteEnvLock = new Promise<void>((resolve) => {
+    release = resolve
+  })
+
+  await previous
+  try {
+    return await fn()
+  } finally {
+    release()
+  }
+}
+
+const generateWithPglite = async ({
   schemas,
   defaultDatabase,
   dbDir,
-  pglite = false,
   migrationsDir,
-}: Pick<Context, "schemas" | "defaultDatabase" | "dbDir"> & {
-  pglite?: boolean
-  migrationsDir?: string
+}: Pick<Context, "schemas" | "defaultDatabase"> & {
+  dbDir: string
+  migrationsDir: string
 }) => {
-  dbDir = dbDir ?? "./src/db"
-  migrationsDir = migrationsDir ?? path.join(dbDir, "migrations")
+  const { PGlite } = await import("@electric-sql/pglite")
+  const { fromNodeSocket } = await import("pg-gateway/node")
+  const net = await import("node:net")
 
-  if (pglite) {
-    const { PGlite } = await import("@electric-sql/pglite")
-    const { fromNodeSocket } = await import("pg-gateway/node")
-    const net = await import("node:net")
-
-    const db = new PGlite()
+  const db = new PGlite()
+  const server = net.createServer(async (socket) => {
+    await fromNodeSocket(socket, {
+      serverVersion: "16.3 (PGlite)",
+      auth: {
+        method: "password",
+        validateCredentials: ({ username, password }: any) =>
+          username === "postgres" && password === "postgres",
+        getClearTextPassword: () => "postgres",
+      },
+      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)
+          return responseData
+        } catch (error) {
+          if (process.env.DEBUG?.includes("pgstrap")) {
+            console.error("PGlite protocol execution failed", error)
+          }
+          return undefined
+        }
+      },
+    })
+  })
 
+  try {
     await migrate({
       client: db as any,
       migrationsDir,
@@ -36,56 +85,73 @@ export const generate = async ({
       schemas,
     })
 
-    const server = net.createServer(async (socket) => {
-      const connection = await fromNodeSocket(socket, {
-        serverVersion: "16.3 (PGlite)",
-        auth: {
-          method: "password",
-          validateCredentials: ({ username, password }: any) =>
-            username === "postgres" && password === "postgres",
-          getClearTextPassword: () => "postgres",
-        },
-        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)
-            return responseData
-          } catch {
-            return undefined
-          }
-        },
-      })
-    })
-
-    await new Promise<void>((resolve) => server.listen(0, resolve))
+    await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", 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
+    await withPgliteEnvLock(async () => {
+      const previousEnv = Object.fromEntries(
+        postgresEnvKeys.map((key) => [key, process.env[key]]),
+      ) as Record<string, string | undefined>
 
-    await zg.generate({
-      db: {
-        connectionString,
-      },
-      schemas: Object.fromEntries(
-        schemas.map((s) => [s, { include: "*", exclude: [] }]),
-      ),
-      outDir: dbDir,
+      try {
+        for (const key of postgresEnvKeys) delete process.env[key]
+        process.env.DATABASE_URL = connectionString
+
+        await zg.generate({
+          db: {
+            connectionString,
+          },
+          schemas: Object.fromEntries(
+            schemas.map((s) => [s, { include: "*", exclude: [] }]),
+          ),
+          outDir: dbDir,
+        })
+
+        await dumpTree({
+          targetDir: path.join(dbDir, "structure"),
+          defaultDatabase: "postgres",
+          schemas,
+        })
+      } finally {
+        for (const [key, value] of Object.entries(previousEnv)) {
+          if (value === undefined) delete process.env[key]
+          else process.env[key] = value
+        }
+      }
     })
+  } finally {
+    await new Promise<void>((resolve, reject) => {
+      server.close((err) => {
+        if (err) reject(err)
+        else resolve()
+      })
+    }).catch(() => null)
+
+    await (db as any).close?.()
+  }
+}
+
+export const generate = async ({
+  schemas,
+  defaultDatabase,
+  dbDir,
+  pglite = false,
+  migrationsDir,
+}: Pick<Context, "schemas" | "defaultDatabase" | "dbDir"> & {
+  pglite?: boolean
+  migrationsDir?: string
+}) => {
+  dbDir = dbDir ?? "./src/db"
+  migrationsDir = migrationsDir ?? path.join(dbDir, "migrations")
 
-    await dumpTree({
-      targetDir: path.join(dbDir, "structure"),
-      defaultDatabase: "postgres",
+  if (pglite || !hasExternalDatabaseConfig()) {
+    await generateWithPglite({
       schemas,
+      defaultDatabase,
+      dbDir,
+      migrationsDir,
     })
-
-    server.close()
-    if (prevDbUrl === undefined) delete process.env.DATABASE_URL
-    else process.env.DATABASE_URL = prevDbUrl
     return
   }
 

tests/generate.pglite.test.ts

@@ -2,7 +2,7 @@ import { test, expect } from "bun:test"
 import fs from "fs"
 import os from "os"
 import path from "path"
-import { generate } from "../src/generate"
+import { generate, hasExternalDatabaseConfig } from "../src/generate"
 
 const migrationFile = `
 exports.up = async (pgm) => {
@@ -13,23 +13,38 @@ exports.down = async (pgm) => {
 }
 `
 
-test("generate with pglite runs migrations and dumps structure", async () => {
-  const tmp = fs.mkdtempSync(path.join(os.tmpdir(), "pgstrap-generate-"))
-  const migrationsDir = path.join(tmp, "migrations")
+const postgresEnvKeys = [
+  "DATABASE_URL",
+  "PGHOST",
+  "PGPORT",
+  "PGDATABASE",
+  "PGUSER",
+  "PGPASSWORD",
+]
+
+const snapshotPostgresEnv = () =>
+  Object.fromEntries(
+    postgresEnvKeys.map((key) => [key, process.env[key]]),
+  ) as Record<string, string | undefined>
+
+const restorePostgresEnv = (
+  previousEnv: Record<string, string | undefined>,
+) => {
+  for (const [key, value] of Object.entries(previousEnv)) {
+    if (value === undefined) delete process.env[key]
+    else process.env[key] = value
+  }
+}
+
+const writeMigration = (migrationsDir: string) => {
   fs.mkdirSync(migrationsDir, { recursive: true })
   fs.writeFileSync(
     path.join(migrationsDir, "001_create_table.js"),
     migrationFile,
   )
+}
 
-  await generate({
-    schemas: ["public"],
-    defaultDatabase: "postgres",
-    dbDir: path.join(tmp, "db"),
-    migrationsDir,
-    pglite: true,
-  })
-
+const expectGeneratedArtifacts = (tmp: string) => {
   const zapatosFile = path.join(tmp, "db", "zapatos", "schema.d.ts")
   const structureDir = path.join(
     tmp,
@@ -42,6 +57,68 @@ test("generate with pglite runs migrations and dumps structure", async () => {
 
   expect(fs.existsSync(zapatosFile)).toBe(true)
   expect(fs.existsSync(path.join(structureDir, "table.sql"))).toBe(true)
+}
+
+test("hasExternalDatabaseConfig detects postgres env settings", () => {
+  expect(hasExternalDatabaseConfig({})).toBe(false)
+  expect(
+    hasExternalDatabaseConfig({ DATABASE_URL: "postgres://example" }),
+  ).toBe(true)
+  expect(hasExternalDatabaseConfig({ PGHOST: "127.0.0.1" })).toBe(true)
+})
+
+test("generate without postgres env uses pglite for migrations and dumps structure", async () => {
+  const tmp = fs.mkdtempSync(path.join(os.tmpdir(), "pgstrap-generate-"))
+  const migrationsDir = path.join(tmp, "migrations")
+  const previousEnv = snapshotPostgresEnv()
+
+  try {
+    writeMigration(migrationsDir)
+
+    for (const key of postgresEnvKeys) {
+      delete process.env[key]
+    }
+
+    await generate({
+      schemas: ["public"],
+      defaultDatabase: "postgres",
+      dbDir: path.join(tmp, "db"),
+      migrationsDir,
+    })
+
+    expectGeneratedArtifacts(tmp)
+  } finally {
+    restorePostgresEnv(previousEnv)
+    fs.rmSync(tmp, { recursive: true, force: true })
+  }
+})
+
+test("generate with pglite true ignores external postgres env settings", async () => {
+  const tmp = fs.mkdtempSync(path.join(os.tmpdir(), "pgstrap-generate-"))
+  const migrationsDir = path.join(tmp, "migrations")
+  const previousEnv = snapshotPostgresEnv()
+
+  try {
+    writeMigration(migrationsDir)
+
+    process.env.DATABASE_URL = "postgres://external:external@example.invalid/db"
+    process.env.PGHOST = "example.invalid"
+    process.env.PGPORT = "5432"
+    process.env.PGDATABASE = "external"
+    process.env.PGUSER = "external"
+    process.env.PGPASSWORD = "external"
+
+    await generate({
+      schemas: ["public"],
+      defaultDatabase: "postgres",
+      dbDir: path.join(tmp, "db"),
+      migrationsDir,
+      pglite: true,
+    })
 
-  fs.rmSync(tmp, { recursive: true, force: true })
+    expectGeneratedArtifacts(tmp)
+  } finally {
+    restorePostgresEnv(previousEnv)
+    fs.rmSync(tmp, { recursive: true, force: true })
+  }
 })