diff --git a/README.md b/README.md index 20e3127..e5a49e3 100644 --- a/README.md +++ b/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 diff --git a/src/cli.ts b/src/cli.ts index 9a9bdec..41db988 100644 --- a/src/cli.ts +++ b/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() diff --git a/src/generate.ts b/src/generate.ts index f337094..2c5b11c 100644 --- a/src/generate.ts +++ b/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 & { - pglite?: boolean - migrationsDir?: string -}) => { - dbDir = dbDir ?? "./src/db" - migrationsDir = migrationsDir ?? path.join(dbDir, "migrations") +type GenerateMode = "auto" | "pglite" | "postgres" +type GenerateRuntimeContext = Pick & { + 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((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 ( + connectionString: string, + fn: () => Promise, +) => { + 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() + 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((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((resolve, reject) => + server.close((error) => (error ? reject(error) : resolve())), + ) + } + await db.close() + } +} + +export const generate = async ({ + schemas, + defaultDatabase, + dbDir, + pglite = "auto", + migrationsDir, +}: Pick & { + 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, }) } diff --git a/tests/generate.pglite.test.ts b/tests/generate.pglite.test.ts index 56dcd53..57ef6b4 100644 --- a/tests/generate.pglite.test.ts +++ b/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 }) + } +})