feat: add transaction-aware database context

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",

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 <T>(operation: (currentTransaction: DatabaseTransaction) => Promise<T>) => {
+      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)
+  })
+})

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<Parameters<Database["transaction"]>[0]>[0]
+export type DatabaseClient = Database | DatabaseTransaction
+
+type DatabaseContext =
+  | { client: Database; isTransaction: false }
+  | { client: DatabaseTransaction; isTransaction: true }
+
+const storage = new AsyncLocalStorage<DatabaseContext>()
+
+export function useDatabase(): DatabaseClient {
+  return storage.getStore()?.client ?? database()
+}
+
+export function withDatabase<T>(client: Database, operation: () => T): T {
+  return storage.run({ client, isTransaction: false }, operation)
+}
+
+export async function withTransaction<T>(
+  operation: (transaction: DatabaseTransaction) => Promise<T>
+): Promise<T> {
+  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))
+  )
+}