Drizzle ORM Multi-Tenant RLS Wrapper

drizzle_rls.ts

type DrizzleClient = NodePgDatabase<typeof schema>;

export type Transaction = Parameters<
	Parameters<DrizzleClient["transaction"]>[0]
>[0];

type PathSegment =
	| { type: "get"; prop: PropertyKey }
	| { type: "call"; args: unknown[] };

/**
 * Create a proxy that wraps any Drizzle query chain in an RLS transaction.
 * Tracks property accesses and method calls, replaying them on `tx` when awaited.
 *
 * Examples:
 *   await rlsProxy.select().from(users).where(...)
 *   await rlsProxy.query.users.findMany(...)
 *   await rlsProxy.insert(users).values(...).returning()
 */
function createRLSProxy(
	drizzleDb: DrizzleClient,
	tenantId: string,
	path: PathSegment[] = [],
): any {
	const execute = () => {
		return drizzleDb.transaction(async (tx) => {
			await tx.execute(
				sql`SELECT set_config('app.current_tenant_id', ${tenantId}::text, true)`,
			);

			let current: any = tx;
			let callContext: any = tx;

			for (const segment of path) {
				if (segment.type === "get") {
					callContext = current;
					current = current[segment.prop];
				} else {
					current = current.apply(callContext, segment.args);
					callContext = current;
				}
			}

			return current;
		});
	};

	return new Proxy(function () {}, {
		get(_, prop) {
			// Implement Promise interface - mirrors Drizzle's QueryPromise
			if (prop === "then") {
				return (resolve: any, reject: any) => execute().then(resolve, reject);
			}
			if (prop === "catch") {
				return (onRejected: any) => execute().then(undefined, onRejected);
			}
			if (prop === "finally") {
				return (onFinally: any) =>
					execute().then(
						(value) => {
							onFinally?.();
							return value;
						},
						(reason) => {
							onFinally?.();
							throw reason;
						},
					);
			}
			return createRLSProxy(drizzleDb, tenantId, [
				...path,
				{ type: "get", prop },
			]);
		},
		apply(_, __, args) {
			return createRLSProxy(drizzleDb, tenantId, [
				...path,
				{ type: "call", args },
			]);
		},
	});
}

/**
 * Database client wrapper that provides RLS (Row-Level Security) support.
 * When tenantId is set:
 * - .query.* calls are wrapped in transactions with app.current_tenant_id set
 * - .transaction() automatically sets app.current_tenant_id
 * - Direct CRUD (.insert/.update/.delete/.select) requires explicit transaction() for RLS
 */
export class DatabaseClient {
	private readonly db: DrizzleClient;
	readonly $client: pg.Pool;
	readonly tenantId?: string;

	constructor(pool: pg.Pool, tenantId?: string) {
		this.$client = pool;
		this.db = drizzle<typeof schema>({ client: pool, schema, logger: false });
		this.tenantId = tenantId;
	}

	/** Create a new client with tenant context for RLS */
	withTenant(tenantId: string): DatabaseClient {
		return new DatabaseClient(this.$client, tenantId);
	}

	/** Execute operations in a transaction. Automatically sets RLS context if tenantId is set. */
	transaction: DrizzleClient["transaction"] = async (fn, config) => {
		if (!this.tenantId) {
			return this.db.transaction(fn, config);
		}

		return this.db.transaction(async (tx) => {
			await tx.execute(
				sql`SELECT set_config('app.current_tenant_id', ${this.tenantId}::text, true)`,
			);
			return fn(tx);
		}, config);
	};

	/** Relational query API. Wrapped in RLS transaction when tenantId is set. */
	get query(): DrizzleClient["query"] {
		if (!this.tenantId) return this.db.query;
		return createRLSProxy(this.db, this.tenantId, [
			{ type: "get", prop: "query" },
		]) as DrizzleClient["query"];
	}

	/** Select query builder. Wrapped in RLS transaction when tenantId is set. */
	get select(): DrizzleClient["select"] {
		if (!this.tenantId) return this.db.select.bind(this.db);
		return ((...args: Parameters<DrizzleClient["select"]>) =>
			createRLSProxy(this.db, this.tenantId!, [
				{ type: "get", prop: "select" },
				{ type: "call", args },
			])) as DrizzleClient["select"];
	}

	/** Insert query builder. Wrapped in RLS transaction when tenantId is set. */
	get insert(): DrizzleClient["insert"] {
		if (!this.tenantId) return this.db.insert.bind(this.db);
		return ((...args: Parameters<DrizzleClient["insert"]>) =>
			createRLSProxy(this.db, this.tenantId!, [
				{ type: "get", prop: "insert" },
				{ type: "call", args },
			])) as DrizzleClient["insert"];
	}

	/** Update query builder. Wrapped in RLS transaction when tenantId is set. */
	get update(): DrizzleClient["update"] {
		if (!this.tenantId) return this.db.update.bind(this.db);
		return ((...args: Parameters<DrizzleClient["update"]>) =>
			createRLSProxy(this.db, this.tenantId!, [
				{ type: "get", prop: "update" },
				{ type: "call", args },
			])) as DrizzleClient["update"];
	}

	/** Delete query builder. Wrapped in RLS transaction when tenantId is set. */
	get delete(): DrizzleClient["delete"] {
		if (!this.tenantId) return this.db.delete.bind(this.db);
		return ((...args: Parameters<DrizzleClient["delete"]>) =>
			createRLSProxy(this.db, this.tenantId!, [
				{ type: "get", prop: "delete" },
				{ type: "call", args },
			])) as DrizzleClient["delete"];
	}

	/** Select distinct query builder. Wrapped in RLS transaction when tenantId is set. */
	get selectDistinct(): DrizzleClient["selectDistinct"] {
		if (!this.tenantId) return this.db.selectDistinct.bind(this.db);
		return ((...args: Parameters<DrizzleClient["selectDistinct"]>) =>
			createRLSProxy(this.db, this.tenantId!, [
				{ type: "get", prop: "selectDistinct" },
				{ type: "call", args },
			])) as DrizzleClient["selectDistinct"];
	}

	/**
	 * Raw drizzle client for migrations, better-auth adapter, and complex SQL (e.g., subqueries in templates).
	 * WARNING: Queries to RLS-enabled tables will return empty results unless tenant context is set via transaction().
	 */
	get rawDb(): DrizzleClient {
		return this.db;
	}
}

/** Create database pool without migrations - use once at startup */
export async function createDatabasePool(
	env: string | undefined,
	options: { maxConnections?: number; target?: string } = {
		maxConnections: 5,
	},
): Promise<{ pool: pg.Pool; onEnd: () => Promise<void> }> {
	const config = getDatabaseConfig(env, options.target);

	if (!config.isProductionTarget) {
		const pool = new Pool({
			connectionString: config.dbConnectionString,
			database: config.databaseName,
			max: options.maxConnections,
		});

		return { pool, onEnd: () => pool.end() };
	} else {
		const auth = await getCoreServicesAccountAuth(env);
		const connector = new Connector({ auth });

		const clientOpts = await connector.getOptions({
			instanceConnectionName: config.instanceConnectionName,
			ipType: IpAddressTypes.PUBLIC,
			authType: AuthTypes.IAM,
		});

		const pool = new Pool({
			...clientOpts,
			user: config.applicationsServiceUserName,
			database: config.databaseName,
			max: options.maxConnections,
		});

		return {
			pool,
			onEnd: async () => {
				connector.close();
				await pool.end();
			},
		};
	}
}