diff --git a/packages/db/package.json b/packages/db/package.json index 50cae3d1..6b1f969d 100644 --- a/packages/db/package.json +++ b/packages/db/package.json @@ -10,7 +10,8 @@ "exports": { "./helpers": "./src/helpers.ts", "./schema": "./src/schema.ts", - "./client": "./src/client.ts" + "./client": "./src/client.ts", + "./context": "./src/context.ts" }, "scripts": { "bump:deps": "bun update --interactive", diff --git a/packages/db/src/__tests__/context.test.ts b/packages/db/src/__tests__/context.test.ts new file mode 100644 index 00000000..ed9d3d3f --- /dev/null +++ b/packages/db/src/__tests__/context.test.ts @@ -0,0 +1,86 @@ +import { describe, expect, test } from "bun:test" +import type { Database } from "#client.ts" +import { type DatabaseTransaction, useDatabase, withDatabase, withTransaction } from "#context.ts" + +function createDatabaseDouble(transaction: DatabaseTransaction) { + let transactionCount = 0 + + const client = { + transaction: async (operation: (currentTransaction: DatabaseTransaction) => Promise) => { + transactionCount += 1 + return operation(transaction) + }, + } as unknown as Database + + return { + client, + transactionCount: () => transactionCount, + } +} + +describe("useDatabase", () => { + test("returns the database bound to the current async context", async () => { + const first = {} as Database + const second = {} as Database + + await Promise.all([ + withDatabase(first, async () => { + await Promise.resolve() + expect(useDatabase()).toBe(first) + }), + withDatabase(second, async () => { + await Promise.resolve() + expect(useDatabase()).toBe(second) + }), + ]) + }) + + test("restores the parent database after a nested context finishes", () => { + const parent = {} as Database + const child = {} as Database + + withDatabase(parent, () => { + expect(useDatabase()).toBe(parent) + + withDatabase(child, () => { + expect(useDatabase()).toBe(child) + }) + + expect(useDatabase()).toBe(parent) + }) + }) +}) + +describe("withTransaction", () => { + test("binds a transaction to the current async context", async () => { + const transaction = {} as DatabaseTransaction + const root = createDatabaseDouble(transaction) + + await withDatabase(root.client, async () => { + await withTransaction((currentTransaction) => { + expect(currentTransaction).toBe(transaction) + expect(useDatabase()).toBe(transaction) + return Promise.resolve() + }) + }) + + expect(root.transactionCount()).toBe(1) + }) + + test("reuses the active transaction for nested operations", async () => { + const transaction = {} as DatabaseTransaction + const root = createDatabaseDouble(transaction) + + await withDatabase(root.client, async () => { + await withTransaction(async () => { + await withTransaction((currentTransaction) => { + expect(currentTransaction).toBe(transaction) + expect(useDatabase()).toBe(transaction) + return Promise.resolve() + }) + }) + }) + + expect(root.transactionCount()).toBe(1) + }) +}) diff --git a/packages/db/src/context.ts b/packages/db/src/context.ts new file mode 100644 index 00000000..d1644d7a --- /dev/null +++ b/packages/db/src/context.ts @@ -0,0 +1,35 @@ +import { AsyncLocalStorage } from "node:async_hooks" +import { database, type Database } from "#client.ts" + +export type DatabaseTransaction = Parameters[0]>[0] +export type DatabaseClient = Database | DatabaseTransaction + +type DatabaseContext = + | { client: Database; isTransaction: false } + | { client: DatabaseTransaction; isTransaction: true } + +const storage = new AsyncLocalStorage() + +export function useDatabase(): DatabaseClient { + return storage.getStore()?.client ?? database() +} + +export function withDatabase(client: Database, operation: () => T): T { + return storage.run({ client, isTransaction: false }, operation) +} + +export async function withTransaction( + operation: (transaction: DatabaseTransaction) => Promise +): Promise { + const current = storage.getStore() + + if (current?.isTransaction) { + return operation(current.client) + } + + const client = current?.client ?? database() + + return client.transaction((transaction) => + storage.run({ client: transaction, isTransaction: true }, () => operation(transaction)) + ) +}