drizzle-orm icon indicating copy to clipboard operation
drizzle-orm copied to clipboard

[TUTORIAL]: Using in-memory Postgres when testing with vitest

Open wilfredjonathanjames opened this issue 9 months ago • 16 comments

Sharing for posterity as it may help someone. A vi.mock to switch pg databases to in-memory pglite, push your drizzle schema, and insert seed data. pglite runs-WASM compiled postgres in memory, which is ideal for testing as there's no docker containers, no delay, and it's real pg. This solution supports parallelism and watch mode. Most importantly - it uses push in place of migrate, as there's no need to create migration files when testing.

Note to Drizzle maintainers, please do not remove drizzle-kit/api's pushSchema method. While it's not well documented, it allows us to push schema directly to test dbs. Please don't make me push via child_process execSync again.

Stack:

  • Typescript
  • Vitest
  • Drizzle
  • NodePostgres (drizzle-orm/node-postgres)

Mock:

// vitest.setup.ts

vi.mock("src/db", async (importOriginal) => {
  const { default: _, ...rest } =
    await importOriginal<typeof import("src/db")>()

  const { PGlite } = await vi.importActual<
    typeof import("@electric-sql/pglite")
  >("@electric-sql/pglite")
  const { drizzle } =
    await vi.importActual<typeof import("drizzle-orm/pglite")>(
      "drizzle-orm/pglite",
    )

  // use require to defeat dynamic require error
  // (https://github.com/drizzle-team/drizzle-orm/issues/2853#issuecomment-2668459509)
  const { createRequire } =
    await vi.importActual<typeof import("node:module")>("node:module")
  const require = createRequire(import.meta.url)
  const { pushSchema } =
    require("drizzle-kit/api") as typeof import("drizzle-kit/api")

  const { schema, User } = rest

  const client = new PGlite()
  const db = drizzle(client, { schema })

  // apply schema to db
  const { apply } = await pushSchema(schema, db as any)
  await apply()

  // seed test data
  await db
    .insert(User)
    .values({ email: "[email protected]", externalId: "external" })

  return { default: db, ...rest }
})

Detail:

// vitest.config.mts

import { defineConfig } from "vitest/config"
import tsconfigPaths from "vite-tsconfig-paths"

const config = defineConfig({
  plugins: [tsconfigPaths({ root: "./" })],
  test: {
    globals: true,
    environment: "node",
    setupFiles: "./vitest.setup.ts",
  },
})

export default config
// src/db/index.ts (mocked, db is not exposed to tests)

import { Pool } from "pg"
import { drizzle } from "drizzle-orm/node-postgres"

import { User } from "./user.sql"

export const schema = {
  User,
  ...
}

const db = drizzle(
  new Pool({
    ...
  }),
  {
    schema,
  },
)
export default db

export * from "./user.sql"

wilfredjonathanjames avatar Mar 03 '25 07:03 wilfredjonathanjames

this is great! I'm pretty new to this, so do you have an example usage in a test that would allow parallelism? Especially if you are adding new data to db in different tests. are the mocks isolated with vitest?

speedystar777 avatar Mar 06 '25 15:03 speedystar777

Thanks for the tutorial. I studied your code and trying to integrate into my project. however i noticed pglite cannot be instantiated in vitest globalSetup, so one db instance cannot be shared across all test files am I correct?

the only way to do it seems to be init the db in setupFiles. which means vitest has to spin up a new pglite instance for every test files plus doing schema migration before every test file runs.

that will negates all performance benefits of pglite i think. do you see any performance improvement though using pglite?

kevinlaw91 avatar Mar 11 '25 05:03 kevinlaw91

Thanks @wilfredjonathanjames !

that will negates all performance benefits of pglite i think. do you see any performance improvement though using pglite?

Can't say that a small performance impact bothers me-- this frees up a fair amount of infra burden for my developers and for my CI.

extradosages avatar Mar 11 '25 20:03 extradosages

Thanks!

If you have your schema separated (as needed for some drizzle-kit stuff), vitest.setup.ts can be simplified since you don't need to import the original src/db to split off the schema.

import { vi } from 'vitest';
import * as schema from './src/lib/server/db/schema';
import { PGlite } from "@electric-sql/pglite";
import {drizzle} from "drizzle-orm/pglite"
import { users } from './src/lib/server/db/schema';

async function MOCK_DB() {
	// use require to defeat dynamic require error
	// (https://github.com/drizzle-team/drizzle-orm/issues/2853#issuecomment-2668459509)
	const { createRequire } =
		await vi.importActual<typeof import("node:module")>("node:module")
	const require = createRequire(import.meta.url)
	const { pushSchema } =
		require("drizzle-kit/api") as typeof import("drizzle-kit/api")

	const client = new PGlite()
	const db = drizzle(client, { schema })

	// apply schema to db
	const { apply } = await pushSchema(schema, db as any)
	await apply()

	// seed test data
	await seed(db)

	return { db, ...schema }
}

async function seed(db) {
	await db
		.insert(users) // my table is lowercase
		.values({ email: "[email protected]", id: "aaa" })
}

// // seems to work but actually does not, still loads original file
// vi.mock("$lib/server/db", MOCK_DB)

vi.mock("./src/lib/server/db", MOCK_DB)

TjenWellens avatar Mar 24 '25 11:03 TjenWellens

Hi @kevinlaw91, on my m3 macbook air these are the comparative boot times:

pglite: 228.3ms
docker: 21.3ms (docker run -it --rm -e POSTGRES_PASSWORD=password postgres:latest)

Though I should mention--all but 13.8ms of pglite's boot time is spent on importOriginal, so with some optimisation this would probably beat docker. Pglite itself boots in 0.5ms.

I'd hazard that even without optimisation most engineers would swallow 200ms in favour of hard test separation and greater simplicity.

Wouldn't recommend going down the globalSetup route.

@extradosages happy to hear this might help you guys!

@TjenWellens this is a great point, and you'll likely improve boot time by 100ms or so. I'm exporting all of my db tables and utils from shared/db, otherwise your way would make sense.

wilfredjonathanjames avatar Mar 28 '25 08:03 wilfredjonathanjames

Hi @speedystar777 no need to turn anything on. Just plug the mock above into a setup file, tweak to suit, and your test database will run in memory. Setup files run once per test file so you will get one db instance per. Up to you how you handle conflicts within each file but there are strategies. Vitest runs same-file tests in series by default. Yell if you need a hand.

wilfredjonathanjames avatar Mar 28 '25 08:03 wilfredjonathanjames

@wilfredjonathanjames Hey thanks for your detailed explanation,

I'm trying to recreate this experience but I'm failing due to some very odd Zod errors...

schema.ts

import { pgTable, serial, varchar } from "drizzle-orm/pg-core";
export const people = pgTable("people", {
  id: serial("id").primaryKey(),
  name: varchar("name", { length: 255 }).notNull(),
});

index.ts

import { PGlite } from "@electric-sql/pglite";
import { drizzle } from "drizzle-orm/node-postgres";
import * as dbSchema from "./db/simpleSchema";
import { pushSchema } from "drizzle-kit/api";

const pg = new PGlite();
const db = drizzle(pg, { schema: dbSchema });
const { apply } = await pushSchema(dbSchema, db as any); // <----- Errors here
await apply();

Error:

bun run ./src/test.ts
[✓] Pulling schema from database...
3791 |     ZodError = class _ZodError extends Error {
3792 |       get errors() {
3793 |         return this.issues;
3794 |       }
3795 |       constructor(issues) {
3796 |         super();
               ^
ZodError: [
  {
    "code": "invalid_type",
    "expected": "string",
    "received": "undefined",
    "path": [
      "enums",
      "public.undefined",
      "name"
    ],
    "message": "Required"
  },
  {
    "code": "invalid_type",
    "expected": "string",
    "received": "undefined",
    "path": [
      "enums",
      "public.undefined",
      "values",
      0
    ],
    "message": "Required"
  },
  {
    "code": "invalid_type",
    "expected": "string",
    "received": "undefined",
    "path": [
      "enums",
      "public.undefined",
      "values",
      1
    ],
    "message": "Required"
  },
  {
    "code": "invalid_type",
    "expected": "string",
    "received": "undefined",
    "path": [
      "enums",
      "public.undefined",
      "values",
      2
    ],
    "message": "Required"
  },
  {
    "code": "invalid_type",
    "expected": "string",
    "received": "undefined",
    "path": [
      "enums",
      "public.undefined",
      "values",
      3
    ],
    "message": "Required"
  },

(and it goes on incrementing path for a couple hundred values).

Have you run into this issue? I know you're running vi, perhaps my issue is related to bun run instead.

Aside: If, instead of running pushSchema, I create the schema myself with raw sql, it works fine. So it definitely seems related to either the function itself or my use of it.

Anyway, if you have any help I would greatly appreciate it :)

Yours, Levi

leviFrosty avatar May 09 '25 15:05 leviFrosty

@wilfredjonathanjames Hey thanks for your detailed explanation,

I'm trying to recreate this experience but I'm failing due to some very odd Zod errors...

schema.ts

import { pgTable, serial, varchar } from "drizzle-orm/pg-core"; export const people = pgTable("people", { id: serial("id").primaryKey(), name: varchar("name", { length: 255 }).notNull(), }); index.ts

import { PGlite } from "@electric-sql/pglite"; import { drizzle } from "drizzle-orm/node-postgres"; import * as dbSchema from "./db/simpleSchema"; import { pushSchema } from "drizzle-kit/api";

const pg = new PGlite(); const db = drizzle(pg, { schema: dbSchema }); const { apply } = await pushSchema(dbSchema, db as any); // <----- Errors here await apply(); Error:

bun run ./src/test.ts
[✓] Pulling schema from database...
3791 |     ZodError = class _ZodError extends Error {
3792 |       get errors() {
3793 |         return this.issues;
3794 |       }
3795 |       constructor(issues) {
3796 |         super();
               ^
ZodError: [
  {
    "code": "invalid_type",
    "expected": "string",
    "received": "undefined",
    "path": [
      "enums",
      "public.undefined",
      "name"
    ],
    "message": "Required"
  },
  {
    "code": "invalid_type",
    "expected": "string",
    "received": "undefined",
    "path": [
      "enums",
      "public.undefined",
      "values",
      0
    ],
    "message": "Required"
  },
  {
    "code": "invalid_type",
    "expected": "string",
    "received": "undefined",
    "path": [
      "enums",
      "public.undefined",
      "values",
      1
    ],
    "message": "Required"
  },
  {
    "code": "invalid_type",
    "expected": "string",
    "received": "undefined",
    "path": [
      "enums",
      "public.undefined",
      "values",
      2
    ],
    "message": "Required"
  },
  {
    "code": "invalid_type",
    "expected": "string",
    "received": "undefined",
    "path": [
      "enums",
      "public.undefined",
      "values",
      3
    ],
    "message": "Required"
  },

(and it goes on incrementing path for a couple hundred values).

Have you run into this issue? I know you're running vi, perhaps my issue is related to bun run instead.

Aside: If, instead of running pushSchema, I create the schema myself with raw sql, it works fine. So it definitely seems related to either the function itself or my use of it.

Anyway, if you have any help I would greatly appreciate it :)

Yours, Levi

Turns out I just messed up the import when transitioning from node-postgres to PGlite I changed import { drizzle } from "drizzle-orm/node-postgres"; -> import { drizzle } from "drizzle-orm/pglite"; and everything works great :)

leviFrosty avatar May 09 '25 16:05 leviFrosty

Hello! I am trying to reproduce this in a trpc api but I am stuck in doing so. Do you have an example of the full folder structure? Thanks!

AdolfoRoquero avatar Jun 03 '25 21:06 AdolfoRoquero

Hi @AdolfoRoquero, you should have everything you need here. A little hard to help if you don't share the issue you're having

wilfredjonathanjames avatar Jul 02 '25 00:07 wilfredjonathanjames

Hey everyone!

I've created this message to send in a batch to all opened issues we have, just because there are a lot of them and I want to update all of you with our current work, why issues are not responded to, and the amount of work that has been done by our team over ~8 months.

I saw a lot of issues with suggestions on how to fix something while we were not responding – so thanks everyone. Also, thanks to everyone patiently waiting for a response from us and continuing to use Drizzle!

We currently have 4 major branches with a lot of work done. Each branch was handled by different devs and teams to make sure we could make all the changes in parallel.


First branch is drizzle-kit rewrite

All of the work can be found on the alternation-engine branch. Here is a PR with the work done: https://github.com/drizzle-team/drizzle-orm/pull/4439

As you can see, it has 167k added lines of code and 67k removed, which means we've completely rewritten the drizzle-kit alternation engine, the way we handle diffs for each dialect, together with expanding our test suite from 600 tests to ~9k test units for all different types of actions you can do with kit. More importantly, we changed the migration folder structure and made commutative migrations, so you won't face complex conflicts on migrations when working in a team.

What's left here:

  • We are finishing handling defaults for Postgres, the last being geometry (yes, we fixed the srid issue here as well).
  • We are finishing commutative migrations for all dialects.
  • We are finishing up the command, so the migration flow will be as simple as drizzle-kit up for you.

Where it brings us:

  • We are getting drizzle-kit into a new good shape where we can call it [email protected]!

Timeline:

  • We need ~2 weeks to finish all of the above and send this branch to beta for testing.

Second big branch is a complex one with several HUGE updates

  • Bringing Relational Queries v2 finally live. We've done a lot of work here to actually make it faster than RQBv1 and much better from a DX point of view. But in implementing it, we had to make another big rewrite, so we completely rewrote the drizzle-orm type system, which made it much simpler and improved type performance by ~21.4x:
(types instantiations for 3300 lines production drizzle schema + 990 lines relations)

TS v5.8.3: 728.8k -> 34.1k
TS v5.9.2: 553.7k -> 25.4k

You can read more about it here.

What's left here:

Where it brings us:

  • We are getting drizzle-orm into a new good shape where we can call it [email protected]!

Breaking changes:

  • We will have them, but we will have open channels for everyone building on top of drizzle types, so we can guide you through all the changes.

Third branch is adding support for CockroachDB and MSSQL dialects

Support for them is already in the alternation-engine branch and will be available together with the drizzle-kit rewrite.

Summary

All of the work we are doing is crucial and should be done sooner rather than later. We've received a lot of feedback and worked really hard to find the best strategies and decisions for API, DX, architecture, etc., so we can confidently mark it as v1 and be sure we can improve it and remain flexible for all the features you are asking for, while becoming even better for everyone building on top of the drizzle API as well.

We didn't want to stay with some legacy decisions and solutions we had, and instead wanted to shape Drizzle in a way that will be best looking ahead to 2025–2026 trends (v1 will get proper effect support, etc.).

We believe that all of the effort we've put in will boost Drizzle and benefit everyone using it.

Thanks everyone, as we said, we are here to stay for a long time to build a great tool together!

Timelines

We are hoping to get v1 for drizzle in beta this fall and same timeline for latest. Right after that we can go through all of the issues and PRs and resond everyone. v1 for drizzle should close ~70% of all the bug tickets we have, so on beta release we will start marking them as closed!

AndriiSherman avatar Aug 30 '25 18:08 AndriiSherman

I'm unable to make multiple connections at the same time using this setup.

apparently to achieve that, i'd need to setup pglite and drizzle using the worker approach?

i'm unable to find anything that helps me understand how to do that

airtonix avatar Sep 09 '25 08:09 airtonix

Is there anyone facing a connection stuck when using this approach with some transactions queries ?

aecsar avatar Nov 03 '25 18:11 aecsar

Is there anyone facing a connection stuck when using this approach with some transactions queries ?

I fixed the issue, which was related to using base database instance instead of transaction inside of a transaction. Thanks for this.

aecsar avatar Nov 03 '25 23:11 aecsar

I haven't looked into it, but it's possible that https://github.com/drizzle-team/drizzle-orm/pull/4999 may have some material impact on the best way to approach this (for the better, afaict!).

extradosages avatar Nov 04 '25 14:11 extradosages

AMAZING SOLUTION! THANK YOU!

https://github.com/drizzle-team/drizzle-orm/issues/2853#issuecomment-2668459509 is now resolved and we can use import agian!

FWIW, here's a slightly less verbose version of the setup file:

import { PGlite } from '@electric-sql/pglite';
import * as originalDb from './path/to/db';
import { pushSchema } from 'drizzle-kit/api';
import { drizzle } from 'drizzle-orm/pglite';
import { vi } from 'vitest';

vi.doMock('./path/to/db', async () => {
  const db = drizzle(new PGlite(), { schema: originalDb.schema });

  await (await pushSchema(originalDb.schema, db as any)).apply();

  return { ...originalDb, db };
});

You can then do something like this in your test to setup db:

import { db } from './path/to/db';

beforeEach(async () => {
  await db.insert(...).values({ ... })
})

ycmjason avatar Nov 13 '25 07:11 ycmjason