diff --git a/package.json b/package.json
index e41d51f..45939e8 100644
--- a/package.json
+++ b/package.json
@@ -7,7 +7,8 @@
   "main": "dist/index.cjs",
   "scripts": {
     "test": "npm run build && ava",
-    "build": "tsup ./src/index.ts ./src/cli.ts --outDir ./dist --dts --sourcemap inline"
+    "build": "tsup ./src/index.ts ./src/cli.ts --outDir ./dist --dts --sourcemap inline",
+    "db:generate": "bun run ./src/generate.ts"
   },
   "bin": {
     "pgstrap": "./dist/cli.cjs"
diff --git a/src/generate.ts b/src/generate.ts
index 59fb29e..2f8491f 100644
--- a/src/generate.ts
+++ b/src/generate.ts
@@ -1,27 +1,35 @@
-import * as zg from "zapatos/generate"
-import {
-  getConnectionStringFromEnv,
-  getPgConnectionFromEnv,
-} from "pg-connection-from-env"
-import { Context } from "./get-project-context"
-import { dumpTree } from "pg-schema-dump"
-import path from "path"
+import * as zg from "zapatos/generate";
+import { Context } from "./get-project-context";
+import { dumpTree } from "pg-schema-dump";
+import path from "path";
+import { PGlite } from "@electric-sql/pglite";
+
+class MockClient {
+  constructor(private pglite: PGlite) {}
+
+  async query(sql: string, params?: any[]) {
+    return this.pglite.query(sql, params);
+  }
+
+  async connect() {} 
+
+  async end() {
+    await this.pglite.close();
+  }
+}
 
 export const generate = async ({
   schemas,
   defaultDatabase,
   dbDir,
 }: Pick<Context, "schemas" | "defaultDatabase" | "dbDir">) => {
-  dbDir = dbDir ?? "./src/db"
+  dbDir = dbDir ?? "./src/db";
+
+  const pglite = new PGlite({ database: ":memory:" });
+  const client = new MockClient(pglite);
 
   await zg.generate({
-    db: {
-      connectionString: getConnectionStringFromEnv({
-        fallbackDefaults: {
-          database: defaultDatabase,
-        },
-      }),
-    },
+    db: client as any,
     schemas: Object.fromEntries(
       schemas.map((s) => [
         s,
@@ -32,11 +40,13 @@ export const generate = async ({
       ])
     ),
     outDir: dbDir,
-  })
+  });
 
   await dumpTree({
     targetDir: path.join(dbDir, "structure"),
     defaultDatabase,
     schemas,
-  })
-}
+  });
+
+  await client.end();
+};