@willyim/drizzle-audit 0.4.0 → 0.5.0

package/dist/src/d1-runtime/with-audit.js

@@ -21,7 +21,7 @@ function getRowId(row, pk) {
  *
  * @param db - A Drizzle SQLite database instance (D1, better-sqlite3, libsql)
  * @param auditTable - The Drizzle table definition for audit_logs
- * @param context - The audit context (userId, optional workspaceId)
+ * @param context - The audit context (userId, optional context columns)
  *
  * @example
  * import { withAudit } from "drizzle-audit/d1-runtime"
@@ -43,7 +43,7 @@ function getRowId(row, pk) {
  * audited.db.select().from(users).all()
  */
 export function withAudit(db, auditTable, context) {
-    const workspaceIdColumn = (() => {
+    const extraColumns = (() => {
         const cols = getTableColumns(auditTable);
         const known = new Set([
             "id",
@@ -55,10 +55,17 @@ export function withAudit(db, auditTable, context) {
             "new_data",
             "created_at",
         ]);
-        const extra = Object.keys(cols).find((k) => !known.has(k));
-        return extra ?? null;
+        return Object.keys(cols).filter((k) => !known.has(k));
     })();
+    const contextValues = { ...(context.context ?? {}) };
     function buildAuditRow(tableName, operation, rowId, oldData, newData) {
+        const extra = {};
+        for (const column of extraColumns) {
+            const value = contextValues[column];
+            if (value !== undefined && value !== "") {
+                extra[column] = value;
+            }
+        }
         return {
             table_name: tableName,
             operation,
@@ -66,9 +73,7 @@ export function withAudit(db, auditTable, context) {
             user_id: context.userId,
             old_data: oldData ? JSON.stringify(oldData) : null,
             new_data: newData ? JSON.stringify(newData) : null,
-            ...(workspaceIdColumn && context.workspaceId
-                ? { [workspaceIdColumn]: context.workspaceId }
-                : {}),
+            ...extra,
         };
     }
     return {

package/dist/src/index.d.ts

@@ -1,5 +1,6 @@
 export * from "./postgres/index.js";
-export * from "./d1/index.js";
+export { clearD1AuditContext, createAttachD1AuditTriggerSql, createAttachD1AuditTriggersSql, createAttachD1AuditTriggerSqlWithColumns, createAttachD1AuditTriggersSqlWithColumns, createD1AuditInstallSql, d1AuditContextTable, d1AuditLogTable, setD1AuditContext, withD1AuditedTransaction, } from "./d1/index.js";
+export type { D1AuditContextOptions, D1AuditInstallOptions, D1AuditLogTableOptions, D1AuditSqlExecutor, D1AuditTriggerTarget, D1AuditTriggerTargetWithColumns, } from "./d1/index.js";
 export * from "./d1-runtime/index.js";
 export * from "./compute-diff.js";
 //# sourceMappingURL=index.d.ts.map

package/dist/src/index.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../src/index.ts"],"names":[],"mappings":"AAAA,cAAc,qBAAqB,CAAA;AACnC,cAAc,eAAe,CAAA;AAC7B,cAAc,uBAAuB,CAAA;AACrC,cAAc,mBAAmB,CAAA"}
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../src/index.ts"],"names":[],"mappings":"AAAA,cAAc,qBAAqB,CAAA;AACnC,OAAO,EACL,mBAAmB,EACnB,6BAA6B,EAC7B,8BAA8B,EAC9B,wCAAwC,EACxC,yCAAyC,EACzC,uBAAuB,EACvB,mBAAmB,EACnB,eAAe,EACf,iBAAiB,EACjB,wBAAwB,GACzB,MAAM,eAAe,CAAA;AACtB,YAAY,EACV,qBAAqB,EACrB,qBAAqB,EACrB,sBAAsB,EACtB,kBAAkB,EAClB,oBAAoB,EACpB,+BAA+B,GAChC,MAAM,eAAe,CAAA;AACtB,cAAc,uBAAuB,CAAA;AACrC,cAAc,mBAAmB,CAAA"}

package/dist/src/index.js

@@ -1,4 +1,4 @@
 export * from "./postgres/index.js";
-export * from "./d1/index.js";
+export { clearD1AuditContext, createAttachD1AuditTriggerSql, createAttachD1AuditTriggersSql, createAttachD1AuditTriggerSqlWithColumns, createAttachD1AuditTriggersSqlWithColumns, createD1AuditInstallSql, d1AuditContextTable, d1AuditLogTable, setD1AuditContext, withD1AuditedTransaction, } from "./d1/index.js";
 export * from "./d1-runtime/index.js";
 export * from "./compute-diff.js";

package/dist/src/postgres/audit-log-schema.d.ts

@@ -1,6 +1,7 @@
+import type { AuditContextColumn } from "./types.js";
 export type PgAuditLogTableOptions = {
-    /** When set (e.g. "workspace_id"), the table definition includes this optional column to match the install. */
-    workspaceIdColumn?: string;
+    /** Extra context columns to include in the table definition, matching the install. */
+    contextColumns?: AuditContextColumn[];
 };
 export declare function pgAuditLogTable(options?: PgAuditLogTableOptions): import("drizzle-orm/pg-core").PgTableWithColumns<{
     name: "audit_logs";

package/dist/src/postgres/audit-log-schema.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"audit-log-schema.d.ts","sourceRoot":"","sources":["../../../src/postgres/audit-log-schema.ts"],"names":[],"mappings":"AASA,MAAM,MAAM,sBAAsB,GAAG;IACnC,+GAA+G;IAC/G,iBAAiB,CAAC,EAAE,MAAM,CAAA;CAC3B,CAAA;AAED,wBAAgB,eAAe,CAAC,OAAO,CAAC,EAAE,sBAAsB;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GA2B/D"}
+{"version":3,"file":"audit-log-schema.d.ts","sourceRoot":"","sources":["../../../src/postgres/audit-log-schema.ts"],"names":[],"mappings":"AASA,OAAO,KAAK,EAAE,kBAAkB,EAAE,MAAM,YAAY,CAAA;AAEpD,MAAM,MAAM,sBAAsB,GAAG;IACnC,sFAAsF;IACtF,cAAc,CAAC,EAAE,kBAAkB,EAAE,CAAA;CACtC,CAAA;AAiBD,wBAAgB,eAAe,CAAC,OAAO,CAAC,EAAE,sBAAsB;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GAyB/D"}

package/dist/src/postgres/audit-log-schema.js

@@ -1,15 +1,25 @@
 import { bigserial, index, jsonb, pgTable, text, timestamp, } from "drizzle-orm/pg-core";
+function resolveColumns(options) {
+    const columns = [];
+    const seen = new Set();
+    for (const entry of options?.contextColumns ?? []) {
+        const column = entry.column?.trim();
+        if (column && !seen.has(column)) {
+            seen.add(column);
+            columns.push(column);
+        }
+    }
+    return columns;
+}
 export function pgAuditLogTable(options) {
-    const workspaceIdColumn = options?.workspaceIdColumn?.trim();
+    const contextColumns = resolveColumns(options);
     const columns = {
         id: bigserial("id", { mode: "number" }).primaryKey(),
         table_name: text("table_name").notNull(),
         operation: text("operation").notNull(),
         row_id: text("row_id"),
         user_id: text("user_id"),
-        ...(workspaceIdColumn
-            ? { [workspaceIdColumn]: text(workspaceIdColumn) }
-            : {}),
+        ...Object.fromEntries(contextColumns.map((c) => [c, text(c)])),
         old_data: jsonb("old_data"),
         new_data: jsonb("new_data"),
         created_at: timestamp("created_at", { withTimezone: true })

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

@@ -1,6 +1,7 @@
 export { pgAuditLogTable } from "./audit-log-schema.js";
 export type { PgAuditLogTableOptions } from "./audit-log-schema.js";
-export { createAttachAuditTriggerSql, createAttachAuditTriggersSql, createAuditAddWorkspaceColumnSql, createAuditInstallSql, } from "./sql.js";
+export { createAttachAuditTriggerSql, createAttachAuditTriggersSql, createAuditAddContextColumnsSql, createAuditInstallSql, } from "./sql.js";
 export { setAuditContext, withAuditedTransaction } from "./runtime.js";
-export type { AuditInstallOptions, AuditSqlExecutor, AuditTransactionCapable, AuditTriggerTarget, } from "./types.js";
+export type { AuditContextOptions } from "./runtime.js";
+export type { AuditContextColumn, AuditInstallOptions, AuditSqlExecutor, AuditTransactionCapable, AuditTriggerTarget, } from "./types.js";
 //# sourceMappingURL=index.d.ts.map

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

@@ -1 +1 @@
-{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../../src/postgres/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,eAAe,EAAE,MAAM,uBAAuB,CAAA;AACvD,YAAY,EAAE,sBAAsB,EAAE,MAAM,uBAAuB,CAAA;AACnE,OAAO,EACL,2BAA2B,EAC3B,4BAA4B,EAC5B,gCAAgC,EAChC,qBAAqB,GACtB,MAAM,UAAU,CAAA;AACjB,OAAO,EAAE,eAAe,EAAE,sBAAsB,EAAE,MAAM,cAAc,CAAA;AAEtE,YAAY,EACV,mBAAmB,EACnB,gBAAgB,EAChB,uBAAuB,EACvB,kBAAkB,GACnB,MAAM,YAAY,CAAA"}
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../../src/postgres/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,eAAe,EAAE,MAAM,uBAAuB,CAAA;AACvD,YAAY,EAAE,sBAAsB,EAAE,MAAM,uBAAuB,CAAA;AACnE,OAAO,EACL,2BAA2B,EAC3B,4BAA4B,EAC5B,+BAA+B,EAC/B,qBAAqB,GACtB,MAAM,UAAU,CAAA;AACjB,OAAO,EAAE,eAAe,EAAE,sBAAsB,EAAE,MAAM,cAAc,CAAA;AACtE,YAAY,EAAE,mBAAmB,EAAE,MAAM,cAAc,CAAA;AAEvD,YAAY,EACV,kBAAkB,EAClB,mBAAmB,EACnB,gBAAgB,EAChB,uBAAuB,EACvB,kBAAkB,GACnB,MAAM,YAAY,CAAA"}

package/dist/src/postgres/index.js

@@ -1,3 +1,3 @@
 export { pgAuditLogTable } from "./audit-log-schema.js";
-export { createAttachAuditTriggerSql, createAttachAuditTriggersSql, createAuditAddWorkspaceColumnSql, createAuditInstallSql, } from "./sql.js";
+export { createAttachAuditTriggerSql, createAttachAuditTriggersSql, createAuditAddContextColumnsSql, createAuditInstallSql, } from "./sql.js";
 export { setAuditContext, withAuditedTransaction } from "./runtime.js";

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

@@ -1,10 +1,8 @@
 import type { AuditSqlExecutor, AuditTransactionCapable } from "./types.js";
-export declare function setAuditContext(db: AuditSqlExecutor, actorId: string, contextKey?: string, options?: {
-    workspaceId?: string;
-    workspaceContextKey?: string;
-}): Promise<void>;
-export declare function withAuditedTransaction<TTransaction extends AuditSqlExecutor, TResult>(db: AuditTransactionCapable<TTransaction>, actorId: string, callback: (tx: TTransaction) => Promise<TResult> | TResult, contextKey?: string, options?: {
-    workspaceId?: string;
-    workspaceContextKey?: string;
-}): Promise<TResult>;
+export type AuditContextOptions = {
+    /** Map of session GUC name → value to set for the transaction. */
+    context?: Record<string, string>;
+};
+export declare function setAuditContext(db: AuditSqlExecutor, actorId: string, contextKey?: string, options?: AuditContextOptions): Promise<void>;
+export declare function withAuditedTransaction<TTransaction extends AuditSqlExecutor, TResult>(db: AuditTransactionCapable<TTransaction>, actorId: string, callback: (tx: TTransaction) => Promise<TResult> | TResult, contextKey?: string, options?: AuditContextOptions): Promise<TResult>;
 //# sourceMappingURL=runtime.d.ts.map

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

@@ -1 +1 @@
-{"version":3,"file":"runtime.d.ts","sourceRoot":"","sources":["../../../src/postgres/runtime.ts"],"names":[],"mappings":"AAEA,OAAO,KAAK,EACV,gBAAgB,EAChB,uBAAuB,EACxB,MAAM,YAAY,CAAA;AAQnB,wBAAsB,eAAe,CACnC,EAAE,EAAE,gBAAgB,EACpB,OAAO,EAAE,MAAM,EACf,UAAU,SAAgB,EAC1B,OAAO,CAAC,EAAE;IAAE,WAAW,CAAC,EAAE,MAAM,CAAC;IAAC,mBAAmB,CAAC,EAAE,MAAM,CAAA;CAAE,iBAajE;AAED,wBAAsB,sBAAsB,CAC1C,YAAY,SAAS,gBAAgB,EACrC,OAAO,EAEP,EAAE,EAAE,uBAAuB,CAAC,YAAY,CAAC,EACzC,OAAO,EAAE,MAAM,EACf,QAAQ,EAAE,CAAC,EAAE,EAAE,YAAY,KAAK,OAAO,CAAC,OAAO,CAAC,GAAG,OAAO,EAC1D,UAAU,SAAgB,EAC1B,OAAO,CAAC,EAAE;IAAE,WAAW,CAAC,EAAE,MAAM,CAAC;IAAC,mBAAmB,CAAC,EAAE,MAAM,CAAA;CAAE,oBAQjE"}
+{"version":3,"file":"runtime.d.ts","sourceRoot":"","sources":["../../../src/postgres/runtime.ts"],"names":[],"mappings":"AAEA,OAAO,KAAK,EACV,gBAAgB,EAChB,uBAAuB,EACxB,MAAM,YAAY,CAAA;AAwBnB,MAAM,MAAM,mBAAmB,GAAG;IAChC,kEAAkE;IAClE,OAAO,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAA;CACjC,CAAA;AAED,wBAAsB,eAAe,CACnC,EAAE,EAAE,gBAAgB,EACpB,OAAO,EAAE,MAAM,EACf,UAAU,SAAgB,EAC1B,OAAO,CAAC,EAAE,mBAAmB,iBAa9B;AAED,wBAAsB,sBAAsB,CAC1C,YAAY,SAAS,gBAAgB,EACrC,OAAO,EAEP,EAAE,EAAE,uBAAuB,CAAC,YAAY,CAAC,EACzC,OAAO,EAAE,MAAM,EACf,QAAQ,EAAE,CAAC,EAAE,EAAE,YAAY,KAAK,OAAO,CAAC,OAAO,CAAC,GAAG,OAAO,EAC1D,UAAU,SAAgB,EAC1B,OAAO,CAAC,EAAE,mBAAmB,oBAQ9B"}

package/dist/src/postgres/runtime.js

@@ -4,12 +4,24 @@ function assertActorId(actorId) {
         throw new Error("actorId must not be empty");
     }
 }
+/**
+ * Builds the `context` GUC map from a set of options. Empty/undefined values
+ * are dropped so the trigger's NULLIF yields NULL.
+ */
+function resolveContext(options) {
+    const context = {};
+    for (const [key, value] of Object.entries(options?.context ?? {})) {
+        if (value !== undefined && value !== "") {
+            context[key] = value;
+        }
+    }
+    return context;
+}
 export async function setAuditContext(db, actorId, contextKey = "app.user_id", options) {
     assertActorId(actorId);
     await db.execute(sql `select set_config(${contextKey}, ${actorId}, true) as audit_context`);
-    if (options?.workspaceId !== undefined && options.workspaceId !== "") {
-        const wsKey = options.workspaceContextKey ?? "app.workspace_id";
-        await db.execute(sql `select set_config(${wsKey}, ${options.workspaceId}, true) as workspace_context`);
+    for (const [key, value] of Object.entries(resolveContext(options))) {
+        await db.execute(sql `select set_config(${key}, ${value}, true) as audit_context`);
     }
 }
 export async function withAuditedTransaction(db, actorId, callback, contextKey = "app.user_id", options) {

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

@@ -3,12 +3,15 @@ export declare function createAuditInstallSql(options?: AuditInstallOptions): st
 export declare function createAttachAuditTriggerSql(target: AuditTriggerTarget, options?: AuditInstallOptions): string;
 export declare function createAttachAuditTriggersSql(targets: AuditTriggerTarget[], options?: AuditInstallOptions): string;
 /**
- * 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 declare function createAuditAddWorkspaceColumnSql(options: AuditInstallOptions & {
-    workspaceIdColumn: string;
-}): string;
+export declare function createAuditAddContextColumnsSql(options?: AuditInstallOptions): string;
 //# sourceMappingURL=sql.d.ts.map

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

@@ -1 +1 @@
-{"version":3,"file":"sql.d.ts","sourceRoot":"","sources":["../../../src/postgres/sql.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,mBAAmB,EAAE,kBAAkB,EAAE,MAAM,YAAY,CAAA;AAqHzE,wBAAgB,qBAAqB,CAAC,OAAO,GAAE,mBAAwB,UA6DtE;AAED,wBAAgB,2BAA2B,CACzC,MAAM,EAAE,kBAAkB,EAC1B,OAAO,GAAE,mBAAwB,UAiClC;AAED,wBAAgB,4BAA4B,CAC1C,OAAO,EAAE,kBAAkB,EAAE,EAC7B,OAAO,GAAE,mBAAwB,UASlC;AAED;;;;;GAKG;AACH,wBAAgB,gCAAgC,CAC9C,OAAO,EAAE,mBAAmB,GAAG;IAAE,iBAAiB,EAAE,MAAM,CAAA;CAAE,UAkC7D"}
+{"version":3,"file":"sql.d.ts","sourceRoot":"","sources":["../../../src/postgres/sql.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAEV,mBAAmB,EACnB,kBAAkB,EACnB,MAAM,YAAY,CAAA;AAkJnB,wBAAgB,qBAAqB,CAAC,OAAO,GAAE,mBAAwB,UA8DtE;AAED,wBAAgB,2BAA2B,CACzC,MAAM,EAAE,kBAAkB,EAC1B,OAAO,GAAE,mBAAwB,UAiClC;AAED,wBAAgB,4BAA4B,CAC1C,OAAO,EAAE,kBAAkB,EAAE,EAC7B,OAAO,GAAE,mBAAwB,UASlC;AAED;;;;;;;;;;GAUG;AACH,wBAAgB,+BAA+B,CAAC,OAAO,GAAE,mBAAwB,UA2ChF"}

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/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.5.0",
   "description": "Lightweight audit logging for Drizzle ORM using database triggers (Postgres + D1/SQLite)",
   "type": "module",
   "main": "./dist/src/index.js",