@willyim/drizzle-audit 0.4.0 → 0.6.0

package/dist/src/postgres/sql.js

@@ -21,46 +21,58 @@ function assertNonEmpty(value, label) {
     }
     return value;
 }
-function buildTriggerFunctionSql(qualifiedAuditTable, qualifiedTriggerFunction, contextLiteral, options) {
-    const workspaceIdColumn = options.workspaceIdColumn?.trim();
-    const hasWorkspace = Boolean(workspaceIdColumn);
-    const workspaceContextLiteral = hasWorkspace
-        ? quoteLiteral("app." + workspaceIdColumn)
-        : "";
-    const workspaceCol = hasWorkspace ? quoteIdent(workspaceIdColumn) : "";
-    const declWorkspace = hasWorkspace ? "\n  audit_workspace TEXT;" : "";
-    const readWorkspace = hasWorkspace
-        ? `\n  audit_workspace := NULLIF(current_setting(${workspaceContextLiteral}, true), '');`
-        : "";
+/**
+ * Normalizes the context columns from the install options, de-duplicating by
+ * column name and applying defaults for `sessionKey` / `index`.
+ */
+function normalizeContextColumns(options) {
+    const resolved = [];
+    const seen = new Set();
+    for (const entry of options.contextColumns ?? []) {
+        const column = entry.column?.trim();
+        if (!column) {
+            throw new Error("contextColumns[].column must not be empty");
+        }
+        if (seen.has(column)) {
+            continue;
+        }
+        seen.add(column);
+        resolved.push({
+            column,
+            sessionKey: entry.sessionKey?.trim() || `app.${column}`,
+            index: entry.index ?? true,
+        });
+    }
+    return resolved;
+}
+function buildTriggerFunctionSql(qualifiedAuditTable, qualifiedTriggerFunction, contextLiteral, contextColumns) {
+    const declContext = contextColumns
+        .map((_, i) => `\n  audit_ctx_${i} TEXT;`)
+        .join("");
+    const readContext = contextColumns
+        .map((col, i) => `\n  audit_ctx_${i} := NULLIF(current_setting(${quoteLiteral(col.sessionKey)}, true), '');`)
+        .join("");
+    const ctxColIdents = contextColumns.map((col) => quoteIdent(col.column));
+    const ctxColPrefix = ctxColIdents.length ? `, ${ctxColIdents.join(", ")}` : "";
+    const ctxValExprs = contextColumns.map((_, i) => `audit_ctx_${i}`);
+    const ctxValPrefix = ctxValExprs.length ? `, ${ctxValExprs.join(", ")}` : "";
     const insertColsBase = "table_name, operation, row_id, user_id";
-    const insertColsInsert = hasWorkspace
-        ? `${insertColsBase}, ${workspaceCol}, new_data`
-        : `${insertColsBase}, new_data`;
-    const insertColsUpdate = hasWorkspace
-        ? `${insertColsBase}, ${workspaceCol}, old_data, new_data`
-        : `${insertColsBase}, old_data, new_data`;
-    const insertColsDelete = hasWorkspace
-        ? `${insertColsBase}, ${workspaceCol}, old_data`
-        : `${insertColsBase}, old_data`;
-    const valuesInsert = hasWorkspace
-        ? `TG_TABLE_NAME, TG_OP, current_row_id, audit_user, audit_workspace, to_jsonb(NEW)`
-        : `TG_TABLE_NAME, TG_OP, current_row_id, audit_user, to_jsonb(NEW)`;
-    const valuesUpdate = hasWorkspace
-        ? `TG_TABLE_NAME, TG_OP, current_row_id, audit_user, audit_workspace, to_jsonb(OLD), to_jsonb(NEW)`
-        : `TG_TABLE_NAME, TG_OP, current_row_id, audit_user, to_jsonb(OLD), to_jsonb(NEW)`;
-    const valuesDelete = hasWorkspace
-        ? `TG_TABLE_NAME, TG_OP, current_row_id, audit_user, audit_workspace, to_jsonb(OLD)`
-        : `TG_TABLE_NAME, TG_OP, current_row_id, audit_user, to_jsonb(OLD)`;
+    const insertColsInsert = `${insertColsBase}${ctxColPrefix}, new_data`;
+    const insertColsUpdate = `${insertColsBase}${ctxColPrefix}, old_data, new_data`;
+    const insertColsDelete = `${insertColsBase}${ctxColPrefix}, old_data`;
+    const valuesInsert = `TG_TABLE_NAME, TG_OP, current_row_id, audit_user${ctxValPrefix}, to_jsonb(NEW)`;
+    const valuesUpdate = `TG_TABLE_NAME, TG_OP, current_row_id, audit_user${ctxValPrefix}, to_jsonb(OLD), to_jsonb(NEW)`;
+    const valuesDelete = `TG_TABLE_NAME, TG_OP, current_row_id, audit_user${ctxValPrefix}, to_jsonb(OLD)`;
     return `
 CREATE OR REPLACE FUNCTION ${qualifiedTriggerFunction}()
 RETURNS TRIGGER AS $$
 DECLARE
-  audit_user TEXT;${declWorkspace}
+  audit_user TEXT;${declContext}
   row_id_column TEXT;
   current_row JSONB;
   current_row_id TEXT;
 BEGIN
-  audit_user := NULLIF(current_setting(${contextLiteral}, true), '');${readWorkspace}
+  audit_user := NULLIF(current_setting(${contextLiteral}, true), '');${readContext}
 
   row_id_column := COALESCE(NULLIF(TG_ARGV[0], ''), ${quoteLiteral(DEFAULT_ROW_ID_COLUMN)});
 
@@ -102,7 +114,7 @@ export function createAuditInstallSql(options = {}) {
     const auditTable = assertNonEmpty(options.auditTable ?? DEFAULT_AUDIT_TABLE, "auditTable");
     const contextKey = assertNonEmpty(options.contextKey ?? DEFAULT_CONTEXT_KEY, "contextKey");
     const triggerFunctionName = assertNonEmpty(options.triggerFunctionName ?? DEFAULT_TRIGGER_FUNCTION, "triggerFunctionName");
-    const workspaceIdColumn = options.workspaceIdColumn?.trim();
+    const contextColumns = normalizeContextColumns(options);
     const qualifiedAuditTable = qualifyName(auditTable, auditSchema);
     const qualifiedTriggerFunction = qualifyName(triggerFunctionName, auditSchema);
     const contextLiteral = quoteLiteral(contextKey);
@@ -112,7 +124,7 @@ export function createAuditInstallSql(options = {}) {
         "operation TEXT NOT NULL CHECK (operation IN ('INSERT', 'UPDATE', 'DELETE'))",
         "row_id TEXT",
         "user_id TEXT",
-        ...(workspaceIdColumn ? [quoteIdent(workspaceIdColumn) + " TEXT"] : []),
+        ...contextColumns.map((col) => quoteIdent(col.column) + " TEXT"),
         "old_data JSONB",
         "new_data JSONB",
         "created_at TIMESTAMPTZ NOT NULL DEFAULT now()",
@@ -121,11 +133,9 @@ export function createAuditInstallSql(options = {}) {
         `CREATE INDEX IF NOT EXISTS ${quoteIdent(`${auditTable}_table_name_idx`)} ON ${qualifiedAuditTable} (table_name);`,
         `CREATE INDEX IF NOT EXISTS ${quoteIdent(`${auditTable}_row_id_idx`)} ON ${qualifiedAuditTable} (row_id);`,
         `CREATE INDEX IF NOT EXISTS ${quoteIdent(`${auditTable}_user_id_idx`)} ON ${qualifiedAuditTable} (user_id);`,
-        ...(workspaceIdColumn
-            ? [
-                `CREATE INDEX IF NOT EXISTS ${quoteIdent(`${auditTable}_${workspaceIdColumn}_idx`)} ON ${qualifiedAuditTable} (${quoteIdent(workspaceIdColumn)});`,
-            ]
-            : []),
+        ...contextColumns
+            .filter((col) => col.index)
+            .map((col) => `CREATE INDEX IF NOT EXISTS ${quoteIdent(`${auditTable}_${col.column}_idx`)} ON ${qualifiedAuditTable} (${quoteIdent(col.column)});`),
         `CREATE INDEX IF NOT EXISTS ${quoteIdent(`${auditTable}_created_at_idx`)} ON ${qualifiedAuditTable} (created_at DESC);`,
     ];
     return [
@@ -135,7 +145,7 @@ CREATE TABLE IF NOT EXISTS ${qualifiedAuditTable} (
   ${tableColumns.join(",\n  ")}
 );`.trim(),
         ...indexStatements,
-        buildTriggerFunctionSql(qualifiedAuditTable, qualifiedTriggerFunction, contextLiteral, { workspaceIdColumn: workspaceIdColumn ?? undefined }),
+        buildTriggerFunctionSql(qualifiedAuditTable, qualifiedTriggerFunction, contextLiteral, contextColumns),
     ].join("\n\n");
 }
 export function createAttachAuditTriggerSql(target, options = {}) {
@@ -166,25 +176,37 @@ export function createAttachAuditTriggersSql(targets, options = {}) {
         .join("\n\n");
 }
 /**
- * Generates SQL to add the workspace column and replace the trigger function on an
- * existing audit_logs table. Use in a new migration when adding workspace_id after
- * the initial install. Options must match your install (auditSchema, auditTable,
- * triggerFunctionName, contextKey, workspaceIdColumn).
+ * Generates SQL to add context columns and regenerate the trigger function on an
+ * existing audit_logs table. Use in a new migration when adding context columns
+ * after the initial install.
+ *
+ * Pass the FULL set of context columns the audit table should have (the trigger is
+ * a single CREATE OR REPLACE, so it must reference every column). Columns are added
+ * with `ADD COLUMN IF NOT EXISTS`, so already-present columns are left untouched.
+ * Options must match your install (auditSchema, auditTable, triggerFunctionName,
+ * contextKey).
  */
-export function createAuditAddWorkspaceColumnSql(options) {
-    const workspaceIdColumn = assertNonEmpty(options.workspaceIdColumn.trim(), "workspaceIdColumn");
+export function createAuditAddContextColumnsSql(options = {}) {
     const auditSchema = assertNonEmpty(options.auditSchema ?? DEFAULT_AUDIT_SCHEMA, "auditSchema");
     const auditTable = assertNonEmpty(options.auditTable ?? DEFAULT_AUDIT_TABLE, "auditTable");
     const contextKey = assertNonEmpty(options.contextKey ?? DEFAULT_CONTEXT_KEY, "contextKey");
     const triggerFunctionName = assertNonEmpty(options.triggerFunctionName ?? DEFAULT_TRIGGER_FUNCTION, "triggerFunctionName");
+    const contextColumns = normalizeContextColumns(options);
+    if (contextColumns.length === 0) {
+        throw new Error("createAuditAddContextColumnsSql requires at least one context column");
+    }
     const qualifiedAuditTable = qualifyName(auditTable, auditSchema);
     const qualifiedTriggerFunction = qualifyName(triggerFunctionName, auditSchema);
     const contextLiteral = quoteLiteral(contextKey);
     return [
-        `ALTER TABLE ${qualifiedAuditTable} ADD COLUMN IF NOT EXISTS ${quoteIdent(workspaceIdColumn)} TEXT;`,
-        `CREATE INDEX IF NOT EXISTS ${quoteIdent(`${auditTable}_${workspaceIdColumn}_idx`)} ON ${qualifiedAuditTable} (${quoteIdent(workspaceIdColumn)});`,
-        buildTriggerFunctionSql(qualifiedAuditTable, qualifiedTriggerFunction, contextLiteral, {
-            workspaceIdColumn,
-        }),
+        ...contextColumns.flatMap((col) => [
+            `ALTER TABLE ${qualifiedAuditTable} ADD COLUMN IF NOT EXISTS ${quoteIdent(col.column)} TEXT;`,
+            ...(col.index
+                ? [
+                    `CREATE INDEX IF NOT EXISTS ${quoteIdent(`${auditTable}_${col.column}_idx`)} ON ${qualifiedAuditTable} (${quoteIdent(col.column)});`,
+                ]
+                : []),
+        ]),
+        buildTriggerFunctionSql(qualifiedAuditTable, qualifiedTriggerFunction, contextLiteral, contextColumns),
     ].join("\n\n");
 }

package/dist/src/postgres/types.d.ts

@@ -5,13 +5,21 @@ export type AuditSqlExecutor = {
 export type AuditTransactionCapable<TTransaction extends AuditSqlExecutor> = {
     transaction: <TResult>(callback: (tx: TTransaction) => Promise<TResult>) => Promise<TResult>;
 };
+export type AuditContextColumn = {
+    /** Column added to the audit table (TEXT, nullable). */
+    column: string;
+    /** Session GUC the trigger reads. Default `app.${column}`. */
+    sessionKey?: string;
+    /** Create an index on the column. Default true. */
+    index?: boolean;
+};
 export type AuditInstallOptions = {
     auditSchema?: string;
     auditTable?: string;
     contextKey?: string;
     triggerFunctionName?: string;
-    /** When set (e.g. "workspace_id"), the audit table and trigger include this column; trigger reads from session variable app.${workspaceIdColumn}. */
-    workspaceIdColumn?: string;
+    /** Extra context columns added to the audit table and populated by the trigger from session GUCs. */
+    contextColumns?: AuditContextColumn[];
 };
 export type AuditTriggerTarget = {
     table: string;

package/dist/src/postgres/types.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"types.d.ts","sourceRoot":"","sources":["../../../src/postgres/types.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,GAAG,EAAE,MAAM,aAAa,CAAA;AAEtC,MAAM,MAAM,gBAAgB,GAAG;IAC7B,OAAO,EAAE,CAAC,KAAK,EAAE,GAAG,KAAK,OAAO,CAAC,OAAO,CAAC,CAAA;CAC1C,CAAA;AAED,MAAM,MAAM,uBAAuB,CAAC,YAAY,SAAS,gBAAgB,IAAI;IAC3E,WAAW,EAAE,CAAC,OAAO,EACnB,QAAQ,EAAE,CAAC,EAAE,EAAE,YAAY,KAAK,OAAO,CAAC,OAAO,CAAC,KAC7C,OAAO,CAAC,OAAO,CAAC,CAAA;CACtB,CAAA;AAED,MAAM,MAAM,mBAAmB,GAAG;IAChC,WAAW,CAAC,EAAE,MAAM,CAAA;IACpB,UAAU,CAAC,EAAE,MAAM,CAAA;IACnB,UAAU,CAAC,EAAE,MAAM,CAAA;IACnB,mBAAmB,CAAC,EAAE,MAAM,CAAA;IAC5B,qJAAqJ;IACrJ,iBAAiB,CAAC,EAAE,MAAM,CAAA;CAC3B,CAAA;AAED,MAAM,MAAM,kBAAkB,GAAG;IAC/B,KAAK,EAAE,MAAM,CAAA;IACb,MAAM,CAAC,EAAE,MAAM,CAAA;IACf,WAAW,CAAC,EAAE,MAAM,CAAA;IACpB,WAAW,CAAC,EAAE,MAAM,CAAA;CACrB,CAAA"}
+{"version":3,"file":"types.d.ts","sourceRoot":"","sources":["../../../src/postgres/types.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,GAAG,EAAE,MAAM,aAAa,CAAA;AAEtC,MAAM,MAAM,gBAAgB,GAAG;IAC7B,OAAO,EAAE,CAAC,KAAK,EAAE,GAAG,KAAK,OAAO,CAAC,OAAO,CAAC,CAAA;CAC1C,CAAA;AAED,MAAM,MAAM,uBAAuB,CAAC,YAAY,SAAS,gBAAgB,IAAI;IAC3E,WAAW,EAAE,CAAC,OAAO,EACnB,QAAQ,EAAE,CAAC,EAAE,EAAE,YAAY,KAAK,OAAO,CAAC,OAAO,CAAC,KAC7C,OAAO,CAAC,OAAO,CAAC,CAAA;CACtB,CAAA;AAED,MAAM,MAAM,kBAAkB,GAAG;IAC/B,wDAAwD;IACxD,MAAM,EAAE,MAAM,CAAA;IACd,8DAA8D;IAC9D,UAAU,CAAC,EAAE,MAAM,CAAA;IACnB,mDAAmD;IACnD,KAAK,CAAC,EAAE,OAAO,CAAA;CAChB,CAAA;AAED,MAAM,MAAM,mBAAmB,GAAG;IAChC,WAAW,CAAC,EAAE,MAAM,CAAA;IACpB,UAAU,CAAC,EAAE,MAAM,CAAA;IACnB,UAAU,CAAC,EAAE,MAAM,CAAA;IACnB,mBAAmB,CAAC,EAAE,MAAM,CAAA;IAC5B,qGAAqG;IACrG,cAAc,CAAC,EAAE,kBAAkB,EAAE,CAAA;CACtC,CAAA;AAED,MAAM,MAAM,kBAAkB,GAAG;IAC/B,KAAK,EAAE,MAAM,CAAA;IACb,MAAM,CAAC,EAAE,MAAM,CAAA;IACf,WAAW,CAAC,EAAE,MAAM,CAAA;IACpB,WAAW,CAAC,EAAE,MAAM,CAAA;CACrB,CAAA"}

package/dist/test/context.test.d.ts

@@ -0,0 +1,2 @@
+export {};
+//# sourceMappingURL=context.test.d.ts.map

package/dist/test/context.test.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"context.test.d.ts","sourceRoot":"","sources":["../../test/context.test.ts"],"names":[],"mappings":""}

package/dist/test/context.test.js

@@ -0,0 +1,98 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import { currentAudit, ensureAuditedTx, hasAuditContext, maybeCurrentAudit, runWithAuditContext, } from "../src/context/index.js";
+function makeFakeDb() {
+    let txCount = 0;
+    const executed = [];
+    const tx = {
+        async execute(query) {
+            executed.push(query);
+            return undefined;
+        },
+    };
+    const db = {
+        async transaction(cb) {
+            txCount++;
+            return cb(tx);
+        },
+    };
+    return {
+        db,
+        tx,
+        executed,
+        get txCount() {
+            return txCount;
+        },
+    };
+}
+test("ensureAuditedTx opens one tx and writes the actor GUC", async () => {
+    const fake = makeFakeDb();
+    let inside;
+    await runWithAuditContext({ actorId: "user_1" }, () => ensureAuditedTx(fake.db, async (tx) => {
+        inside = tx;
+    }));
+    assert.equal(fake.txCount, 1);
+    assert.equal(inside, fake.tx);
+    // setAuditContext ran at least the actor set_config.
+    assert.ok(fake.executed.length >= 1);
+});
+test("nested ensureAuditedTx reuses the open tx (no second transaction)", async () => {
+    const fake = makeFakeDb();
+    const seen = [];
+    await runWithAuditContext({ actorId: "user_1" }, () => ensureAuditedTx(fake.db, async (tx) => {
+        seen.push(tx);
+        await ensureAuditedTx(fake.db, async (tx2) => {
+            seen.push(tx2);
+            await ensureAuditedTx(fake.db, async (tx3) => seen.push(tx3));
+        });
+    }));
+    assert.equal(fake.txCount, 1, "only one transaction is opened");
+    assert.equal(seen.length, 3);
+    assert.ok(seen.every((t) => t === fake.tx));
+});
+test("lazy resolver runs exactly once, on first write", async () => {
+    const fake = makeFakeDb();
+    let resolved = 0;
+    let seenActor;
+    await runWithAuditContext(() => {
+        resolved++;
+        return { actorId: "lazy" };
+    }, async () => {
+        assert.equal(resolved, 0, "not resolved before any write");
+        assert.equal(maybeCurrentAudit(), null);
+        await ensureAuditedTx(fake.db, async () => {
+            seenActor = currentAudit().actorId;
+            // reentrant — must not re-resolve
+            await ensureAuditedTx(fake.db, async () => { });
+        });
+        // resolved value is memoised on the cell
+        assert.equal(maybeCurrentAudit()?.actorId, "lazy");
+    });
+    assert.equal(resolved, 1);
+    assert.equal(seenActor, "lazy");
+});
+test("eager context resolves immediately (currentAudit before any write)", async () => {
+    await runWithAuditContext({ actorId: "eager", context: { "app.workspace_id": "ws_9" } }, async () => {
+        assert.equal(currentAudit().actorId, "eager");
+        assert.equal(currentAudit().context?.["app.workspace_id"], "ws_9");
+    });
+});
+test("guardrail: writes / reads outside a context fail loudly", async () => {
+    const fake = makeFakeDb();
+    assert.equal(hasAuditContext(), false);
+    assert.equal(maybeCurrentAudit(), null);
+    assert.throws(() => currentAudit(), /no ambient audit context/);
+    await assert.rejects(() => ensureAuditedTx(fake.db, async () => { }), /outside an audit context/);
+    assert.equal(fake.txCount, 0);
+});
+test("context is isolated per runWithAuditContext scope", async () => {
+    const a = runWithAuditContext({ actorId: "A" }, async () => {
+        await Promise.resolve();
+        return currentAudit().actorId;
+    });
+    const b = runWithAuditContext({ actorId: "B" }, async () => {
+        await Promise.resolve();
+        return currentAudit().actorId;
+    });
+    assert.deepEqual(await Promise.all([a, b]), ["A", "B"]);
+});

package/dist/test/d1.integration.test.js

@@ -122,20 +122,22 @@ test("d1 column-aware triggers capture full row data", () => {
 });
 test("d1 workspace_id column and context are stored when enabled", () => {
     const sqlite = new Database(":memory:");
-    const auditLogsWithWorkspace = d1AuditLogTable({ workspaceIdColumn: "workspace_id" });
+    const auditLogsWithWorkspace = d1AuditLogTable({
+        contextColumns: [{ column: "workspace_id" }],
+    });
     const db = drizzle({ client: sqlite, schema: { auditLogs: auditLogsWithWorkspace, auditContext, users } });
     try {
-        sqlite.exec(createD1AuditInstallSql({ workspaceIdColumn: "workspace_id" }));
+        sqlite.exec(createD1AuditInstallSql({ contextColumns: [{ column: "workspace_id" }] }));
         sqlite.exec(`
       CREATE TABLE users (
         id TEXT PRIMARY KEY,
         name TEXT NOT NULL
       );
     `);
-        sqlite.exec(createAttachD1AuditTriggersSql([{ table: "users" }], { workspaceIdColumn: "workspace_id" }));
+        sqlite.exec(createAttachD1AuditTriggersSql([{ table: "users" }], { contextColumns: [{ column: "workspace_id" }] }));
         withD1AuditedTransaction(db, "user_1", (tx) => {
             tx.insert(users).values({ id: "u1", name: "Alice" }).run();
-        }, { workspaceId: "ws_1" });
+        }, { context: { workspace_id: "ws_1" } });
         const logs = db.select().from(auditLogsWithWorkspace).all();
         assert.equal(logs.length, 1);
         assert.equal(logs[0]?.user_id, "user_1");
@@ -157,6 +159,71 @@ test("d1 workspace_id column and context are stored when enabled", () => {
         sqlite.close();
     }
 });
+test("d1 generic contextColumns populate and stay NULL without context", () => {
+    const sqlite = new Database(":memory:");
+    const contextColumns = [
+        { column: "workspace_id" },
+        { column: "tenant_id" },
+        { column: "request_id" },
+    ];
+    const auditLogsWithCtx = d1AuditLogTable({ contextColumns });
+    const db = drizzle({
+        client: sqlite,
+        schema: { auditLogs: auditLogsWithCtx, auditContext, users },
+    });
+    try {
+        sqlite.exec(createD1AuditInstallSql({ contextColumns }));
+        sqlite.exec(`
+      CREATE TABLE users (
+        id TEXT PRIMARY KEY,
+        name TEXT NOT NULL
+      );
+    `);
+        sqlite.exec(createAttachD1AuditTriggersSql([{ table: "users" }], { contextColumns }));
+        withD1AuditedTransaction(db, "user_1", (tx) => {
+            tx.insert(users).values({ id: "u1", name: "Alice" }).run();
+            tx.update(users).set({ name: "Alice Updated" }).where(eq(users.id, "u1")).run();
+            tx.delete(users).where(eq(users.id, "u1")).run();
+        }, {
+            context: {
+                workspace_id: "ws_1",
+                tenant_id: "tenant_1",
+                request_id: "req_1",
+            },
+        });
+        const logs = db
+            .select()
+            .from(auditLogsWithCtx)
+            .orderBy(asc(auditLogsWithCtx.id))
+            .all();
+        assert.equal(logs.length, 3);
+        for (const log of logs) {
+            const row = log;
+            assert.equal(row.user_id, "user_1");
+            assert.equal(row.workspace_id, "ws_1");
+            assert.equal(row.tenant_id, "tenant_1");
+            assert.equal(row.request_id, "req_1");
+        }
+        assert.equal(logs[0]?.operation, "INSERT");
+        assert.equal(logs[1]?.operation, "UPDATE");
+        assert.equal(logs[2]?.operation, "DELETE");
+        // No context: all context columns NULL.
+        db.insert(users).values({ id: "u2", name: "No Context" }).run();
+        const all = db
+            .select()
+            .from(auditLogsWithCtx)
+            .orderBy(asc(auditLogsWithCtx.id))
+            .all();
+        const last = all[all.length - 1];
+        assert.equal(last.user_id, null);
+        assert.equal(last.workspace_id, null);
+        assert.equal(last.tenant_id, null);
+        assert.equal(last.request_id, null);
+    }
+    finally {
+        sqlite.close();
+    }
+});
 test("d1 writes without audit context produce rows with user_id = NULL", () => {
     const { db, sqlite } = setupDb();
     try {

package/dist/test/sqlite.integration.test.js

@@ -165,7 +165,9 @@ test("withAudit handles multi-row update", async () => {
 });
 test("withAudit with workspace_id", async () => {
     const sqlite = new Database(":memory:");
-    const auditLogsWithWs = d1AuditLogTable({ workspaceIdColumn: "workspace_id" });
+    const auditLogsWithWs = d1AuditLogTable({
+        contextColumns: [{ column: "workspace_id" }],
+    });
     const db = drizzle({ client: sqlite, schema: { auditLogs: auditLogsWithWs, users } });
     try {
         sqlite.exec(`
@@ -188,7 +190,7 @@ test("withAudit with workspace_id", async () => {
     `);
         const audited = withAudit(db, auditLogsWithWs, {
             userId: "user_1",
-            workspaceId: "ws_1",
+            context: { workspace_id: "ws_1" },
         });
         await audited.insert(users, { id: "u1", name: "Ada" });
         const logs = db.select().from(auditLogsWithWs).all();
@@ -200,6 +202,67 @@ test("withAudit with workspace_id", async () => {
         sqlite.close();
     }
 });
+test("withAudit with generic context columns", async () => {
+    const sqlite = new Database(":memory:");
+    const auditLogsWithCtx = d1AuditLogTable({
+        contextColumns: [
+            { column: "workspace_id" },
+            { column: "tenant_id" },
+            { column: "request_id" },
+        ],
+    });
+    const db = drizzle({ client: sqlite, schema: { auditLogs: auditLogsWithCtx, users } });
+    try {
+        sqlite.exec(`
+      CREATE TABLE audit_logs (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        table_name TEXT NOT NULL,
+        operation TEXT NOT NULL,
+        row_id TEXT,
+        user_id TEXT,
+        workspace_id TEXT,
+        tenant_id TEXT,
+        request_id TEXT,
+        old_data TEXT,
+        new_data TEXT,
+        created_at TEXT NOT NULL DEFAULT (datetime('now'))
+      );
+      CREATE TABLE users (
+        id TEXT PRIMARY KEY,
+        name TEXT NOT NULL,
+        email TEXT
+      );
+    `);
+        const audited = withAudit(db, auditLogsWithCtx, {
+            userId: "user_1",
+            context: {
+                workspace_id: "ws_1",
+                tenant_id: "tenant_1",
+                request_id: "req_1",
+            },
+        });
+        await audited.insert(users, { id: "u1", name: "Ada" });
+        const logs = db.select().from(auditLogsWithCtx).all();
+        assert.equal(logs.length, 1);
+        const row = logs[0];
+        assert.equal(row.user_id, "user_1");
+        assert.equal(row.workspace_id, "ws_1");
+        assert.equal(row.tenant_id, "tenant_1");
+        assert.equal(row.request_id, "req_1");
+        // No context: extra columns stay NULL.
+        const auditedNoCtx = withAudit(db, auditLogsWithCtx, { userId: "user_2" });
+        await auditedNoCtx.insert(users, { id: "u2", name: "Bob" });
+        const all = db.select().from(auditLogsWithCtx).orderBy(asc(auditLogsWithCtx.id)).all();
+        const last = all[all.length - 1];
+        assert.equal(last.user_id, "user_2");
+        assert.equal(last.workspace_id, null);
+        assert.equal(last.tenant_id, null);
+        assert.equal(last.request_id, null);
+    }
+    finally {
+        sqlite.close();
+    }
+});
 test("withAudit.db gives access to raw db for non-audited ops", async () => {
     const { db, sqlite } = setupDb();
     try {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@willyim/drizzle-audit",
-  "version": "0.4.0",
+  "version": "0.6.0",
   "description": "Lightweight audit logging for Drizzle ORM using database triggers (Postgres + D1/SQLite)",
   "type": "module",
   "main": "./dist/src/index.js",
@@ -8,6 +8,7 @@
   "exports": {
     ".": "./dist/src/index.js",
     "./postgres": "./dist/src/postgres/index.js",
+    "./context": "./dist/src/context/index.js",
     "./d1": "./dist/src/d1/index.js",
     "./d1-runtime": "./dist/src/d1-runtime/index.js"
   },