Add pglite option to pgstrap generate

README.md

@@ -57,6 +57,7 @@ npm install pgstrap --save-dev
 - `npm run db:reset` - Drop and recreate the database, then run all migrations
 - `npm run db:generate` - Generate types and structure dumps
 - `npm run db:create-migration` - Create a new migration file
+- `pgstrap generate --pglite` - Run migrations and generate schema using an in-memory PGlite instance
 
 ### Configuration
 

package.json

@@ -23,6 +23,7 @@
   "license": "MIT",
   "dependencies": {
     "@electric-sql/pglite": "^0.2.10",
+    "pg-gateway": "^0.3.0-beta.4",
     "pg-schema-dump": "^2.0.2",
     "yargs": "^17.7.2"
   },

src/cli.ts

@@ -1,6 +1,13 @@
 #!/usr/bin/env node
 import yargs from "yargs"
-import { migrate, reset, generate, createMigration, initPgstrap } from "./"
+import {
+  migrate,
+  reset,
+  generate,
+  createMigration,
+  initPgstrap,
+  generateWithPglite,
+} from "./"
 import { getProjectContext } from "./get-project-context"
 ;(yargs as any)
   .command("init", "initialize pgstrap", {}, async () => {
@@ -34,9 +41,20 @@ import { getProjectContext } from "./get-project-context"
   .command(
     "generate",
     "generate types and sql documentation from database",
-    {},
-    async () => {
-      generate(await getProjectContext())
+    (yargs) => {
+      yargs.option("pglite", {
+        type: "boolean",
+        describe: "use pglite instead of connecting to Postgres",
+        default: false,
+      })
+    },
+    async (argv) => {
+      const ctx = await getProjectContext()
+      if (argv.pglite) {
+        await generateWithPglite(ctx)
+      } else {
+        await generate(ctx)
+      }
     },
   )
   .parse()

src/generate.ts

@@ -1,11 +1,12 @@
 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 { PGlite } from "@electric-sql/pglite"
+import net from "net"
+import { fromNodeSocket } from "pg-gateway/node"
+import { migrate } from "./migrate"
 
 export const generate = async ({
   schemas,
@@ -40,3 +41,61 @@ export const generate = async ({
     schemas,
   })
 }
+
+export const generateWithPglite = async (
+  ctx: Context & { migrationsDir?: string },
+) => {
+  const db = new PGlite()
+
+  const migrationsDir =
+    ctx.migrationsDir ?? path.join(ctx.dbDir ?? "./src/db", "migrations")
+
+  await migrate({
+    ...ctx,
+    migrationsDir,
+    client: db as any,
+  })
+
+  const server = net.createServer(async (socket) => {
+    const connection = await fromNodeSocket(socket as any, {
+      auth: { method: "trust" },
+      async onStartup() {
+        await db.waitReady
+        return false
+      },
+      async onMessage(data: any, { isAuthenticated }: any) {
+        if (!isAuthenticated) return false
+
+        try {
+          const [[_, responseData]] = (await db.execProtocol(
+            data as any,
+          )) as any
+          connection.sendData(responseData)
+        } catch (err) {
+          connection.sendError(err as Error)
+          connection.sendReadyForQuery()
+        }
+        return true
+      },
+    })
+
+    socket.on("end", () => {
+      // connection cleanup if needed
+    })
+  })
+
+  await new Promise<void>((resolve) => server.listen(0, () => resolve()))
+  server.unref()
+  const port = (server.address() as net.AddressInfo).port
+  const prevUrl = process.env.DATABASE_URL
+  process.env.DATABASE_URL = `postgres://localhost:${port}`
+
+  try {
+    await generate(ctx)
+  } finally {
+    if (prevUrl === undefined) delete process.env.DATABASE_URL
+    else process.env.DATABASE_URL = prevUrl
+    await new Promise<void>((resolve) => server.close(() => resolve()))
+    await db.close()
+  }
+}

tests/generate-pglite.test.ts

@@ -0,0 +1,38 @@
+import test from "ava"
+import fs from "fs"
+import path from "path"
+import { generateWithPglite } from "../src/generate"
+
+const migrationContent = `
+  exports.up = async (pgm) => {
+    pgm.createTable('foo', {
+      id: 'id'
+    })
+  }
+  exports.down = async (pgm) => {
+    pgm.dropTable('foo')
+  }
+`
+
+test("generateWithPglite creates structure", async (t) => {
+  const cwd = fs.mkdtempSync(path.join(__dirname, "temp_proj"))
+  const migrationsDir = path.join(cwd, "migrations")
+  fs.mkdirSync(migrationsDir, { recursive: true })
+  fs.writeFileSync(
+    path.join(migrationsDir, "001_create_foo.cjs"),
+    migrationContent,
+  )
+
+  await generateWithPglite({
+    schemas: ["public"],
+    defaultDatabase: "test_db",
+    dbDir: path.join(cwd, "db"),
+    migrationsDir,
+    cwd,
+  })
+
+  const expected = path.join(cwd, "db/structure/public/tables/foo/table.sql")
+  t.true(fs.existsSync(expected))
+
+  fs.rmSync(cwd, { recursive: true, force: true })
+})

tests/pglite.test.ts

@@ -28,7 +28,7 @@ test("migration of a pglite db works", async (t) => {
     }
   `
   fs.writeFileSync(
-    path.join(migrationsDir, "001_create_test_table.js"),
+    path.join(migrationsDir, "001_create_test_table.cjs"),
     migrationContent,
   )