fix: fall back to pglite for type generation

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. Uses your configured Postgres database when it is available, and falls back to an in-memory PGlite database when Postgres is not running. Use `pgstrap generate --pglite` to force PGlite, or `pgstrap generate --no-pglite` to require Postgres.
 - `npm run db:create-migration` - Create a new migration file
 
 ### Configuration

src/cli.ts

@@ -35,10 +35,17 @@ import { getProjectContext } from "./get-project-context"
     "generate",
     "generate types and sql documentation from database",
     (yargs) => {
-      yargs.option("pglite", { type: "boolean", default: false })
+      yargs.option("pglite", {
+        type: "boolean",
+        describe:
+          "Force PGlite generation. By default pgstrap uses Postgres when available and falls back to PGlite.",
+      })
     },
     async (argv) => {
-      generate({ ...(await getProjectContext()), pglite: !!argv.pglite })
+      await generate({
+        ...(await getProjectContext()),
+        pglite: typeof argv.pglite === "boolean" ? argv.pglite : "auto",
+      })
     },
   )
   .parse()

src/generate.ts

@@ -1,117 +1,241 @@
 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"
+import { Client } from "pg"
 
-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")
+type GenerateMode = "auto" | "pglite" | "postgres"
+type GenerateRuntimeContext = Pick<Context, "schemas" | "defaultDatabase"> & {
+  dbDir: string
+}
 
-  if (pglite) {
-    const { PGlite } = await import("@electric-sql/pglite")
-    const { fromNodeSocket } = await import("pg-gateway/node")
-    const net = await import("node:net")
+const getSchemasConfig = (schemas: string[]) =>
+  Object.fromEntries(
+    schemas.map((s) => [s, { include: "*" as const, exclude: [] as string[] }]),
+  )
 
-    const db = new PGlite()
+const getPostgresConnectionString = (defaultDatabase: string) =>
+  getConnectionStringFromEnv({
+    fallbackDefaults: {
+      database: defaultDatabase,
+    },
+  })
 
-    await migrate({
-      client: db as any,
-      migrationsDir,
-      defaultDatabase,
-      cwd: process.cwd(),
-      schemas,
-    })
+const isPostgresUnavailableError = (error: unknown) => {
+  const err = error as { code?: string; message?: string; cause?: unknown }
+  const cause = err.cause as { code?: string } | undefined
+  const code = err.code ?? cause?.code
+  if (
+    code &&
+    [
+      "ECONNREFUSED",
+      "ECONNRESET",
+      "ENOTFOUND",
+      "ETIMEDOUT",
+      "EHOSTUNREACH",
+      "ENETUNREACH",
+    ].includes(code)
+  ) {
+    return true
+  }
 
-    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
-          }
-        },
-      })
-    })
+  const message = String(err.message ?? error)
+  return /timeout expired|Connection terminated unexpectedly|connect .* refused/i.test(
+    message,
+  )
+}
 
-    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 canConnectToPostgres = async (connectionString: string) => {
+  const client = new Client({
+    connectionString,
+    connectionTimeoutMillis: 1500,
+  })
 
-    const prevDbUrl = process.env.DATABASE_URL
-    process.env.DATABASE_URL = connectionString
+  try {
+    await client.connect()
+    return true
+  } catch (error) {
+    if (isPostgresUnavailableError(error)) return false
+    throw error
+  } finally {
+    await client.end().catch(() => undefined)
+  }
+}
+
+const getGenerateMode = (pglite?: boolean | "auto"): GenerateMode => {
+  if (pglite === true) return "pglite"
+  if (pglite === false) return "postgres"
+  return "auto"
+}
+
+const withDatabaseUrl = async <T>(
+  connectionString: string,
+  fn: () => Promise<T>,
+) => {
+  const prevDbUrl = process.env.DATABASE_URL
+  process.env.DATABASE_URL = connectionString
+
+  try {
+    return await fn()
+  } finally {
+    if (prevDbUrl === undefined) delete process.env.DATABASE_URL
+    else process.env.DATABASE_URL = prevDbUrl
+  }
+}
 
+const generateUsingConnectionString = async ({
+  connectionString,
+  schemas,
+  defaultDatabase,
+  dbDir,
+  setDatabaseUrl = false,
+}: GenerateRuntimeContext & {
+  connectionString: string
+  setDatabaseUrl?: boolean
+}) => {
+  const run = async () => {
     await zg.generate({
       db: {
         connectionString,
       },
-      schemas: Object.fromEntries(
-        schemas.map((s) => [s, { include: "*", exclude: [] }]),
-      ),
+      schemas: getSchemasConfig(schemas),
       outDir: dbDir,
     })
 
     await dumpTree({
       targetDir: path.join(dbDir, "structure"),
-      defaultDatabase: "postgres",
+      defaultDatabase,
       schemas,
     })
+  }
 
-    server.close()
-    if (prevDbUrl === undefined) delete process.env.DATABASE_URL
-    else process.env.DATABASE_URL = prevDbUrl
-    return
+  if (setDatabaseUrl) {
+    return withDatabaseUrl(connectionString, run)
   }
 
-  await zg.generate({
-    db: {
-      connectionString: getConnectionStringFromEnv({
-        fallbackDefaults: {
-          database: defaultDatabase,
-        },
-      }),
-    },
-    schemas: Object.fromEntries(
-      schemas.map((s) => [
-        s,
-        {
-          include: "*",
-          exclude: [],
-        },
-      ]),
-    ),
-    outDir: dbDir,
-  })
+  return run()
+}
 
-  await dumpTree({
-    targetDir: path.join(dbDir, "structure"),
+const generateWithPglite = async ({
+  schemas,
+  defaultDatabase,
+  dbDir,
+  migrationsDir,
+}: GenerateRuntimeContext & {
+  migrationsDir: string
+}) => {
+  const { PGlite } = await import("@electric-sql/pglite")
+  const { fromNodeSocket } = await import("pg-gateway/node")
+  const net = await import("node:net")
+
+  const db = new PGlite()
+  await db.waitReady
+
+  await migrate({
+    client: db as any,
+    migrationsDir,
     defaultDatabase,
+    cwd: process.cwd(),
+    schemas,
+  })
+
+  const sockets = new Set<import("node:net").Socket>()
+  const server = net.createServer((socket) => {
+    sockets.add(socket)
+    socket.on("close", () => sockets.delete(socket))
+
+    fromNodeSocket(socket, {
+      serverVersion: "16.3 (PGlite)",
+      auth: {
+        method: "password",
+        validateCredentials: ({ username, password }: any) =>
+          username === "postgres" && password === "postgres",
+        getClearTextPassword: () => "postgres",
+      },
+      async onStartup() {
+        await db.waitReady
+      },
+      async onMessage(data: Uint8Array, { isAuthenticated }: any) {
+        if (!isAuthenticated) return
+        try {
+          const { data: responseData } = await db.execProtocol(data)
+          return responseData
+        } catch {
+          return undefined
+        }
+      },
+    }).catch((error) => socket.destroy(error))
+  })
+
+  try {
+    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`
+
+    await generateUsingConnectionString({
+      connectionString,
+      schemas,
+      defaultDatabase: "postgres",
+      dbDir,
+      setDatabaseUrl: true,
+    })
+  } finally {
+    for (const socket of sockets) socket.destroy()
+    if (server.listening) {
+      await new Promise<void>((resolve, reject) =>
+        server.close((error) => (error ? reject(error) : resolve())),
+      )
+    }
+    await db.close()
+  }
+}
+
+export const generate = async ({
+  schemas,
+  defaultDatabase,
+  dbDir,
+  pglite = "auto",
+  migrationsDir,
+}: Pick<Context, "schemas" | "defaultDatabase" | "dbDir"> & {
+  pglite?: boolean | "auto"
+  migrationsDir?: string
+}) => {
+  dbDir = dbDir ?? "./src/db"
+  migrationsDir = migrationsDir ?? path.join(dbDir, "migrations")
+  const mode = getGenerateMode(pglite)
+
+  if (mode === "pglite") {
+    await generateWithPglite({
+      schemas,
+      defaultDatabase,
+      dbDir,
+      migrationsDir,
+    })
+    return
+  }
+
+  const connectionString = getPostgresConnectionString(defaultDatabase)
+
+  if (mode === "auto") {
+    const postgresAvailable = await canConnectToPostgres(connectionString)
+    if (!postgresAvailable) {
+      console.info("Postgres unavailable; generating with PGlite.")
+      await generateWithPglite({
+        schemas,
+        defaultDatabase,
+        dbDir,
+        migrationsDir,
+      })
+      return
+    }
+  }
+
+  await generateUsingConnectionString({
+    connectionString,
     schemas,
+    defaultDatabase,
+    dbDir,
   })
 }

tests/generate.pglite.test.ts

@@ -45,3 +45,42 @@ test("generate with pglite runs migrations and dumps structure", async () => {
 
   fs.rmSync(tmp, { recursive: true, force: true })
 })
+
+test("generate falls back to pglite when postgres is unavailable", async () => {
+  const tmp = fs.mkdtempSync(path.join(os.tmpdir(), "pgstrap-generate-"))
+  const migrationsDir = path.join(tmp, "migrations")
+  fs.mkdirSync(migrationsDir, { recursive: true })
+  fs.writeFileSync(
+    path.join(migrationsDir, "001_create_table.js"),
+    migrationFile,
+  )
+
+  const prevDbUrl = process.env.DATABASE_URL
+  process.env.DATABASE_URL = "postgres://postgres:postgres@127.0.0.1:1/postgres"
+
+  try {
+    await generate({
+      schemas: ["public"],
+      defaultDatabase: "postgres",
+      dbDir: path.join(tmp, "db"),
+      migrationsDir,
+    })
+
+    const zapatosFile = path.join(tmp, "db", "zapatos", "schema.d.ts")
+    const structureDir = path.join(
+      tmp,
+      "db",
+      "structure",
+      "public",
+      "tables",
+      "foo",
+    )
+
+    expect(fs.existsSync(zapatosFile)).toBe(true)
+    expect(fs.existsSync(path.join(structureDir, "table.sql"))).toBe(true)
+  } finally {
+    if (prevDbUrl === undefined) delete process.env.DATABASE_URL
+    else process.env.DATABASE_URL = prevDbUrl
+    fs.rmSync(tmp, { recursive: true, force: true })
+  }
+})