|
| 1 | +import { fileURLToPath } from "node:url" |
| 2 | +import { PGlite } from "@electric-sql/pglite" |
| 3 | +import { type PostgresClient, type PostgresDb, postgresSchema } from "@platform/db-postgres" |
| 4 | +import { sql } from "drizzle-orm" |
| 5 | +import { drizzle } from "drizzle-orm/pglite" |
| 6 | +import { migrate } from "drizzle-orm/pglite/migrator" |
| 7 | + |
| 8 | +const MIGRATIONS_FOLDER = fileURLToPath(new URL("../../../db-postgres/drizzle", import.meta.url)) |
| 9 | + |
| 10 | +const unsafeCast = <T>(value: unknown): T => value as T |
| 11 | + |
| 12 | +export interface InMemoryPostgres { |
| 13 | + readonly client: PGlite |
| 14 | + readonly db: ReturnType<typeof drizzle> |
| 15 | + readonly postgresDb: PostgresDb |
| 16 | + /** Admin client — runs as the table owner (superuser). RLS is bypassed. */ |
| 17 | + readonly adminPostgresClient: PostgresClient |
| 18 | + /** |
| 19 | + * App-role client whose transactions run under the `latitude_app` role so |
| 20 | + * that Postgres RLS policies are enforced. The role switch is injected via |
| 21 | + * drizzle's `tx.execute()` (transaction-scoped) before each query batch. |
| 22 | + */ |
| 23 | + readonly appPostgresClient: PostgresClient |
| 24 | +} |
| 25 | + |
| 26 | +const createPostgresClientFromDb = (postgresDb: PostgresDb): PostgresClient => { |
| 27 | + const transaction = <T>(fn: (txDb: PostgresDb) => Promise<T>): Promise<T> => |
| 28 | + (postgresDb as unknown as { transaction: (fn: (tx: unknown) => Promise<T>) => Promise<T> }).transaction( |
| 29 | + async (tx) => fn(tx as PostgresDb), |
| 30 | + ) |
| 31 | + return unsafeCast<PostgresClient>({ db: postgresDb, transaction }) |
| 32 | +} |
| 33 | + |
| 34 | +/** |
| 35 | + * Create an app-role client that switches to `latitude_app` for each |
| 36 | + * transaction so that RLS policies are enforced. |
| 37 | + * |
| 38 | + * The role switch uses `SET LOCAL ROLE` executed through drizzle's transaction |
| 39 | + * `tx` handle (not via a raw `pglite.exec` call) to avoid interleaving raw |
| 40 | + * PGlite I/O with drizzle's serialized connection queue, which would deadlock |
| 41 | + * on the single-connection PGlite instance. |
| 42 | + */ |
| 43 | +const createAppRoleClient = (postgresDb: PostgresDb): PostgresClient => { |
| 44 | + const transaction = <T>(fn: (txDb: PostgresDb) => Promise<T>): Promise<T> => |
| 45 | + (postgresDb as unknown as { transaction: (fn: (tx: unknown) => Promise<T>) => Promise<T> }).transaction( |
| 46 | + async (tx) => { |
| 47 | + // Switch to the runtime role so RLS policies apply. |
| 48 | + // SET LOCAL is transaction-scoped and reverts automatically on |
| 49 | + // commit/rollback — no cleanup needed. |
| 50 | + await (tx as PostgresDb).execute(sql`SET LOCAL ROLE latitude_app`) |
| 51 | + return fn(tx as PostgresDb) |
| 52 | + }, |
| 53 | + ) |
| 54 | + return unsafeCast<PostgresClient>({ db: postgresDb, transaction }) |
| 55 | +} |
| 56 | + |
| 57 | +export const createInMemoryPostgres = async (): Promise<InMemoryPostgres> => { |
| 58 | + const client = new PGlite() |
| 59 | + |
| 60 | + // Create the runtime role before migrations so the grant migration finds it. |
| 61 | + await client.exec("CREATE ROLE latitude_app NOLOGIN") |
| 62 | + |
| 63 | + const db = drizzle({ client, schema: postgresSchema }) |
| 64 | + await migrate(db, { migrationsFolder: MIGRATIONS_FOLDER }) |
| 65 | + |
| 66 | + const postgresDb = unsafeCast<PostgresDb>(db) |
| 67 | + |
| 68 | + return { |
| 69 | + client, |
| 70 | + db, |
| 71 | + postgresDb, |
| 72 | + adminPostgresClient: createPostgresClientFromDb(postgresDb), |
| 73 | + appPostgresClient: createAppRoleClient(postgresDb), |
| 74 | + } |
| 75 | +} |
| 76 | + |
| 77 | +export const closeInMemoryPostgres = async (database: InMemoryPostgres): Promise<void> => { |
| 78 | + await database.client.close() |
| 79 | +} |
0 commit comments