Write a database adapter

import { IamError, type RecordDriver, type StorageRow } from 'better-iam/core';

/** The "database": rows by collection, then by id. Replace it with your database client. */
export type Tables = Map<string, Map<string, StorageRow>>;

const conflict = () =>
  new IamError('CONFLICT', 'A record with this identifier or tenant key already exists', 409);

/** Row operations over `tables`. Each method is one statement in a SQL database. */
export function driverFor(tables: Tables): RecordDriver {
  const table = (collection: string) => {
    let rows = tables.get(collection);
    if (!rows) tables.set(collection, (rows = new Map()));
    return rows;
  };
  const keyTaken = (row: StorageRow) =>
    row.unique_key !== null &&
    [...table(row.collection).values()].some(
      (other) =>
        other.id !== row.id &&
        other.tenant_id === row.tenant_id &&
        other.unique_key === row.unique_key,
    );
  return {
    async select(collection, id, tenantId) {
      // SELECT ... FROM iam_records WHERE collection = ? [AND id = ?] [AND tenant_id = ?]
      return [...table(collection).values()].filter(
        (row) =>
          (id === undefined || row.id === id) &&
          (tenantId === undefined || row.tenant_id === tenantId),
      );
    },
    async insert(row) {
      // INSERT INTO iam_records ...; the primary key (collection, id) and the unique constraint
      // (collection, tenant_id, unique_key) turn duplicates into a conflict.
      if (table(row.collection).has(row.id) || keyTaken(row)) throw conflict();
      table(row.collection).set(row.id, { ...row });
    },
    async update(row) {
      // UPDATE iam_records SET data = ?, unique_key = ? WHERE collection = ? AND id = ? AND tenant_id = ?
      const existing = table(row.collection).get(row.id);
      if (!existing || existing.tenant_id !== row.tenant_id) return false;
      if (keyTaken(row)) throw conflict();
      table(row.collection).set(row.id, { ...row });
      return true;
    },
    async delete(collection, id) {
      // DELETE FROM iam_records WHERE collection = ? AND id = ?
      table(collection).delete(id);
    },
    async collections() {
      // SELECT DISTINCT collection FROM iam_records ORDER BY collection
      return [...tables].filter(([, rows]) => rows.size > 0).map(([name]) => name);
    },
  };
}

import { AsyncLocalStorage } from 'node:async_hooks';
import { IamError, RecordStore, type IamStore, type RecordDriver } from 'better-iam/core';
import { driverFor, type Tables } from './driver';

/** One transaction in progress. `active` turns false when it ends, so leaked handles fail. */
interface TransactionContext {
  store: MapStore;
  active: boolean;
  rollbackOnly: boolean;
}

/**
 * Runs one operation at a time, in arrival order: the serialization the contract requires. A call chain that
 * already holds the lock runs straight through, because RecordStore's in-memory `findOrdered` fallback calls
 * `find` from inside `read`.
 */
class Mutex {
  #tail: Promise<void> = Promise.resolve();
  #holder = new AsyncLocalStorage<Mutex>();
  async run<T>(operation: () => Promise<T>): Promise<T> {
    if (this.#holder.getStore() === this) return operation();
    const previous = this.#tail;
    let release!: () => void;
    this.#tail = new Promise<void>((done) => {
      release = done;
    });
    await previous;
    try {
      return await this.#holder.run(this, operation);
    } finally {
      release();
    }
  }
}

class MapStore extends RecordStore {
  #state: { tables: Tables; closed: boolean };
  #mutex: Mutex;
  #scope: AsyncLocalStorage<TransactionContext>;
  #context: TransactionContext | undefined;

  constructor(
    state: { tables: Tables; closed: boolean },
    mutex: Mutex,
    scope: AsyncLocalStorage<TransactionContext>,
    context?: TransactionContext,
  ) {
    super();
    this.#state = state;
    this.#mutex = mutex;
    this.#scope = scope;
    this.#context = context;
  }

  /** The transaction this call belongs to: this handle's own, or the one the async call chain is in. */
  #active(): TransactionContext | undefined {
    if (this.#state.closed) throw new IamError('STORE_CLOSED', 'Database adapter is closed', 500);
    const context = this.#context ?? this.#scope.getStore();
    if (context && !context.active)
      throw new IamError('TRANSACTION_CLOSED', 'Transaction has already completed', 500);
    return context;
  }

  protected async read<T>(operation: (driver: RecordDriver) => Promise<T>): Promise<T> {
    const context = this.#active();
    if (context) {
      try {
        return await operation(driverFor(this.#state.tables));
      } catch (error) {
        // A failed statement dooms the whole transaction, even if the caller catches the error.
        context.rollbackOnly = true;
        throw error;
      }
    }
    // Outside a transaction, wait for the running one so uncommitted writes are never visible.
    return this.#mutex.run(async () => {
      this.#active();
      return operation(driverFor(this.#state.tables));
    });
  }

  protected async write<T>(operation: (driver: RecordDriver) => Promise<T>): Promise<T> {
    if (!this.#active())
      throw new IamError('TRANSACTION_REQUIRED', 'Database writes require transaction()', 500);
    return this.read(operation);
  }

  async transaction<T>(fn: (tx: IamStore) => Promise<T>): Promise<T> {
    const current = this.#active();
    if (current) {
      // Nested call: join the outer transaction; a failure inside rolls the outer one back.
      try {
        return await fn(current.store);
      } catch (error) {
        current.rollbackOnly = true;
        throw error;
      }
    }
    return this.#mutex.run(async () => {
      this.#active();
      // BEGIN: remember the committed state so a failure can restore it.
      const snapshot: Tables = new Map(
        [...this.#state.tables].map(([name, rows]) => [name, new Map(rows)]),
      );
      const context = { active: true, rollbackOnly: false } as TransactionContext;
      context.store = new MapStore(this.#state, this.#mutex, this.#scope, context);
      try {
        const result = await this.#scope.run(context, () => fn(context.store));
        if (context.rollbackOnly)
          throw new IamError(
            'TRANSACTION_ABORTED',
            'A nested database operation failed; transaction rolled back',
            409,
          );
        return result; // COMMIT
      } catch (error) {
        this.#state.tables = snapshot; // ROLLBACK
        throw error;
      } finally {
        context.active = false;
      }
    });
  }

  async migrate(): Promise<void> {
    // CREATE TABLE IF NOT EXISTS ...: a Map needs no schema. Keep this idempotent.
    this.#active();
  }

  async close(): Promise<void> {
    if (this.#state.closed) return;
    if (this.#active())
      throw new IamError('TRANSACTION_ACTIVE', 'Cannot close an adapter inside a transaction', 500);
    await this.#mutex.run(async () => {
      this.#state.closed = true;
    });
  }
}

/** An IamStore over JavaScript Maps: the contract with no database in the way. */
export function mapAdapter(): IamStore {
  return new MapStore({ tables: new Map(), closed: false }, new Mutex(), new AsyncLocalStorage());
}

import { describe, it } from 'vitest';
import { adapterConformanceCases } from '@better-iam/core/conformance';
import { mapAdapter } from '../map-adapter/store';

describe('map adapter', () => {
  for (const test of adapterConformanceCases())
    it(test.name, async () => {
      const store = mapAdapter();
      await store.migrate();
      try {
        await test.run(store);
      } finally {
        await store.close();
      }
    });
});

import { runAdapterConformance } from '@better-iam/core/conformance';
import { mapAdapter } from '../map-adapter/store';

const { passed, failed } = await runAdapterConformance(() => mapAdapter());
console.log(`${passed.length} passed`);
for (const failure of failed) console.error(`FAIL ${failure.name}:`, failure.error);
process.exitCode = failed.length ? 1 : 0;

import { expect, it } from 'vitest';
import { betterIam } from 'better-iam';
import { instrumentStore, type StoreCall } from 'better-iam/core';
import { mapAdapter } from '../map-adapter/store';

it('runs an instance and keeps lookups selective', async () => {
  const calls: StoreCall[] = [];
  const database = instrumentStore(mapAdapter(), (call) => calls.push(call));
  const iam = betterIam({
    database,
    secret: 'test-secret-with-at-least-32-characters',
    baseURL: 'http://localhost:3000',
  });
  await iam.initialize();
  const root = await iam.bootstrap({
    email: 'root@example.test',
    name: 'Root',
    password: 'a long root test password',
  });
  expect(root.tenant.type).toBe('root');

  // A find without filter keys reads a whole collection.
  const unfiltered = calls.filter((call) => call.method === 'find' && !call.filterKeys?.length);
  expect(unfiltered).toEqual([]);
  await database.close();
});

CREATE TABLE IF NOT EXISTS iam_records (
  collection TEXT NOT NULL,
  id         TEXT NOT NULL,
  tenant_id  TEXT NOT NULL,
  unique_key TEXT,
  data       TEXT NOT NULL,
  PRIMARY KEY (collection, id),
  UNIQUE (collection, tenant_id, unique_key)
)