@hammadj/better-auth-core 1.5.0-beta.9

package/src/context/request-state.ts

@@ -0,0 +1,91 @@
+import type { AsyncLocalStorage } from "@better-auth/core/async_hooks";
+import { getAsyncLocalStorage } from "@better-auth/core/async_hooks";
+import { __getBetterAuthGlobal } from "./global";
+
+export type RequestStateWeakMap = WeakMap<object, any>;
+
+const ensureAsyncStorage = async () => {
+	const betterAuthGlobal = __getBetterAuthGlobal();
+	if (!betterAuthGlobal.context.requestStateAsyncStorage) {
+		const AsyncLocalStorage = await getAsyncLocalStorage();
+		betterAuthGlobal.context.requestStateAsyncStorage =
+			new AsyncLocalStorage<RequestStateWeakMap>();
+	}
+	return betterAuthGlobal.context
+		.requestStateAsyncStorage as AsyncLocalStorage<RequestStateWeakMap>;
+};
+
+export async function getRequestStateAsyncLocalStorage() {
+	return ensureAsyncStorage();
+}
+
+export async function hasRequestState() {
+	const als = await ensureAsyncStorage();
+	return als.getStore() !== undefined;
+}
+
+export async function getCurrentRequestState(): Promise<RequestStateWeakMap> {
+	const als = await ensureAsyncStorage();
+	const store = als.getStore();
+	if (!store) {
+		throw new Error(
+			"No request state found. Please make sure you are calling this function within a `runWithRequestState` callback.",
+		);
+	}
+	return store;
+}
+
+export async function runWithRequestState<T>(
+	store: RequestStateWeakMap,
+	fn: () => T,
+): Promise<T> {
+	const als = await ensureAsyncStorage();
+	return als.run(store, fn);
+}
+
+export interface RequestState<T> {
+	get(): Promise<T>;
+	set(value: T): Promise<void>;
+
+	// A unique reference used as a key to identify this state within the request's WeakMap. Useful for debugging purposes.
+	readonly ref: Readonly<object>;
+}
+
+/**
+ * Defines a request-scoped state with lazy initialization.
+ *
+ * @param initFn - A function that initializes the state. It is called the first time `get()` is invoked within each request context, and only once per context.
+ * @returns A RequestState object with `get` and `set` methods, and a unique `ref` for debugging.
+ *
+ * @example
+ * const userState = defineRequestState(() => ({ id: '', name: '' }));
+ * // Later, within a request context:
+ * const user = await userState.get();
+ */
+export function defineRequestState<T>(
+	initFn: () => T | Promise<T>,
+): RequestState<T>;
+export function defineRequestState(
+	initFn: () => any | Promise<any>,
+): RequestState<any> {
+	const ref = Object.freeze({});
+	return {
+		get ref(): Readonly<object> {
+			return ref;
+		},
+		async get() {
+			const store = await getCurrentRequestState();
+			if (!store.has(ref)) {
+				const initialValue = await initFn();
+				store.set(ref, initialValue);
+				return initialValue;
+			}
+			return store.get(ref);
+		},
+
+		async set(value) {
+			const store = await getCurrentRequestState();
+			store.set(ref, value);
+		},
+	};
+}

package/src/context/transaction.ts

@@ -0,0 +1,136 @@
+import type { AsyncLocalStorage } from "node:async_hooks";
+import { getAsyncLocalStorage } from "@better-auth/core/async_hooks";
+import type { DBAdapter, DBTransactionAdapter } from "../db/adapter";
+import { __getBetterAuthGlobal } from "./global";
+
+type HookContext = {
+	adapter: DBTransactionAdapter;
+	pendingHooks: Array<() => Promise<void>>;
+};
+
+const ensureAsyncStorage = async () => {
+	const betterAuthGlobal = __getBetterAuthGlobal();
+	if (!betterAuthGlobal.context.adapterAsyncStorage) {
+		const AsyncLocalStorage = await getAsyncLocalStorage();
+		betterAuthGlobal.context.adapterAsyncStorage = new AsyncLocalStorage();
+	}
+	return betterAuthGlobal.context
+		.adapterAsyncStorage as AsyncLocalStorage<HookContext>;
+};
+
+/**
+ * This is for internal use only. Most users should use `getCurrentAdapter` instead.
+ *
+ * It is exposed for advanced use cases where you need direct access to the AsyncLocalStorage instance.
+ */
+export const getCurrentDBAdapterAsyncLocalStorage = async () => {
+	return ensureAsyncStorage();
+};
+
+export const getCurrentAdapter = async (
+	fallback: DBTransactionAdapter,
+): Promise<DBTransactionAdapter> => {
+	return ensureAsyncStorage()
+		.then((als) => {
+			const store = als.getStore();
+			return store?.adapter || fallback;
+		})
+		.catch(() => {
+			return fallback;
+		});
+};
+
+export const runWithAdapter = async <R>(
+	adapter: DBAdapter,
+	fn: () => R,
+): Promise<R> => {
+	let called = false;
+	return ensureAsyncStorage()
+		.then(async (als) => {
+			called = true;
+			const pendingHooks: Array<() => Promise<void>> = [];
+			let result: Awaited<R>;
+			let error: unknown;
+			let hasError = false;
+			try {
+				result = await als.run({ adapter, pendingHooks }, fn);
+			} catch (err) {
+				error = err;
+				hasError = true;
+			}
+			// Execute pending hooks after the function completes (even if it threw)
+			for (const hook of pendingHooks) {
+				await hook();
+			}
+			if (hasError) {
+				throw error;
+			}
+			return result!;
+		})
+		.catch((err) => {
+			if (!called) {
+				return fn();
+			}
+			throw err;
+		});
+};
+
+export const runWithTransaction = async <R>(
+	adapter: DBAdapter,
+	fn: () => R,
+): Promise<R> => {
+	let called = true;
+	return ensureAsyncStorage()
+		.then(async (als) => {
+			called = true;
+			const pendingHooks: Array<() => Promise<void>> = [];
+			let result: Awaited<R>;
+			let error: unknown;
+			let hasError = false;
+			try {
+				result = await adapter.transaction(async (trx) => {
+					return als.run({ adapter: trx, pendingHooks }, fn);
+				});
+			} catch (e) {
+				hasError = true;
+				error = e;
+			}
+			for (const hook of pendingHooks) {
+				await hook();
+			}
+			if (hasError) {
+				throw error;
+			}
+			return result!;
+		})
+		.catch((err) => {
+			if (!called) {
+				return fn();
+			}
+			throw err;
+		});
+};
+
+/**
+ * Queue a hook to be executed after the current transaction commits.
+ * If not in a transaction, the hook will execute immediately.
+ */
+export const queueAfterTransactionHook = async (
+	hook: () => Promise<void>,
+): Promise<void> => {
+	return ensureAsyncStorage()
+		.then((als) => {
+			const store = als.getStore();
+			if (store) {
+				// We're in a transaction context, queue the hook
+				store.pendingHooks.push(hook);
+			} else {
+				// Not in a transaction, execute immediately
+				return hook();
+			}
+		})
+		.catch(() => {
+			// No async storage available, execute immediately
+			return hook();
+		});
+};