import { AsyncLocalStorage } from "node:async_hooks";

/**
 * Request-scoped Drizzle session for transaction binding.
 *
 * Pool-backed dialects (PG/MySQL) must not stash the active transaction on a
 * shared `dialect.db` field: concurrent awaits would interleave and steal each
 * other's connection. AsyncLocalStorage keeps the binding per async chain.
 */
const drizzleSessionStorage = new AsyncLocalStorage<unknown>();

/** Run `work` with `session` as the active Drizzle handle for this async chain. */
export function runWithDrizzleSession<TSession, T>(session: TSession, work: () => Promise<T>): Promise<T> {
  return drizzleSessionStorage.run(session, work);
}

/** Active transaction session if any, otherwise the dialect's root handle. */
export function resolveDrizzleSession<TSession>(root: TSession): TSession {
  const current = drizzleSessionStorage.getStore();
  return current === undefined ? root : (current as TSession);
}