Feat: [Pg] support for row level security

drizzle-orm/src/aws-data-api/pg/session.ts

@@ -183,6 +183,8 @@ export class AwsDataApiSession<
 			await tx.setTransaction(config);
 		}
 		try {
+			await tx.executeRLSConfig(config);
+
 			const result = await transaction(tx);
 			await this.client.send(new CommitTransactionCommand({ ...this.rawQuery, transactionId }));
 			return result;

drizzle-orm/src/neon-serverless/session.ts

@@ -136,6 +136,8 @@ export class NeonSession<
 		const tx = new NeonTransaction(this.dialect, session, this.schema);
 		await tx.execute(sql`begin ${tx.getTransactionConfigSQL(config)}`);
 		try {
+			await tx.executeRLSConfig(config);
+
 			const result = await transaction(tx);
 			await tx.execute(sql`commit`);
 			return result;

drizzle-orm/src/node-postgres/session.ts

@@ -135,6 +135,8 @@ export class NodePgSession<
 		const tx = new NodePgTransaction(this.dialect, session, this.schema);
 		await tx.execute(sql`begin${config ? sql` ${tx.getTransactionConfigSQL(config)}` : undefined}`);
 		try {
+			await tx.executeRLSConfig(config);
+
 			const result = await transaction(tx);
 			await tx.execute(sql`commit`);
 			return result;

drizzle-orm/src/pg-core/policy.ts

@@ -0,0 +1,53 @@
+import { entityKind } from '~/entity.ts';
+import type { AnyPgTable } from '~/pg-core/table.ts';
+import { Policy } from '~/policy.ts';
+import type { SQL } from '~/sql/sql.ts';
+import type { AnyPgRole } from './role.ts';
+
+export type PgPolicyFor = 'select' | 'insert' | 'update' | 'delete' | 'all';
+
+export type PgPolicyTo = 'public' | 'current_role' | 'current_user' | 'session_user';
+
+export type PgPolicyAs = 'permissive' | 'restrictive';
+
+export type PgPolicyConfig = {
+	as?: PgPolicyAs;
+	for?: PgPolicyFor;
+	to?: (AnyPgRole | PgPolicyTo)[];
+	using?: SQL;
+	withCheck?: SQL;
+};
+
+export const PolicyTable = Symbol.for('drizzle:PolicyTable');
+
+export class PgPolicy<
+	TName extends string,
+	TTable extends AnyPgTable,
+	TConfig extends PgPolicyConfig | undefined,
+> extends Policy<TName> {
+	static readonly [entityKind]: string = 'PgPolicy';
+
+	[PolicyTable]: TTable;
+	config: TConfig;
+
+	declare readonly _: {
+		readonly brand: 'PgPolicy';
+		readonly name: TName;
+		readonly table: TTable;
+		readonly config: TConfig;
+	};
+
+	constructor(name: TName, table: TTable, config: TConfig) {
+		super(name);
+		this[PolicyTable] = table;
+		this.config = config;
+	}
+}
+
+export function pgPolicy<TName extends string, TTable extends AnyPgTable, TConfig extends PgPolicyConfig>(
+	name: TName,
+	table: TTable,
+	config?: TConfig,
+): PgPolicy<TName, TTable, TConfig> {
+	return new PgPolicy(name, table, config as TConfig);
+}

drizzle-orm/src/pg-core/role.ts

@@ -0,0 +1,46 @@
+import { entityKind } from '~/entity.ts';
+import { Role } from '~/role.ts';
+
+// Since the create role clause only allow one of these, I guess drizzle-kit will have to generate
+// alter role clauses if the user defines more than one of these
+export type PgRoleConfig = {
+	superuser?: boolean;
+	createDb?: boolean;
+	createRole?: boolean;
+	inherit?: boolean;
+	login?: boolean;
+	replication?: boolean;
+	bypassRLS?: boolean;
+	connectionLimit?: number;
+	password?: string;
+	validUntil?: Date;
+	inRole?: string;
+	role?: string[]; // Should this be a PgRole[]?
+	admin?: string[]; // Should this be a PgRole[]?
+};
+
+export type AnyPgRole = PgRole<string, PgRoleConfig>;
+
+export class PgRole<TName extends string, TConfig extends PgRoleConfig> extends Role<TName> {
+	static readonly [entityKind]: string = 'PgRole';
+
+	declare readonly _: {
+		readonly brand: 'PgRole';
+		readonly name: TName;
+		readonly config: TConfig;
+	};
+
+	config: TConfig;
+
+	constructor(name: TName, config: TConfig) {
+		super(name);
+		this.config = config;
+	}
+}
+
+export function pgRole<TName extends string, TConfig extends PgRoleConfig>(
+	name: TName,
+	config?: TConfig,
+): PgRole<TName, TConfig> {
+	return new PgRole(name, config ?? {} as TConfig);
+}

drizzle-orm/src/pg-core/session.ts

@@ -6,6 +6,7 @@ import { tracer } from '~/tracing.ts';
 import { PgDatabase } from './db.ts';
 import type { PgDialect } from './dialect.ts';
 import type { SelectedFieldsOrdered } from './query-builders/select.types.ts';
+import type { AnyPgRole } from './role.ts';
 
 export interface PreparedQueryConfig {
 	execute: unknown;
@@ -29,6 +30,14 @@ export interface PgTransactionConfig {
 	isolationLevel?: 'read uncommitted' | 'read committed' | 'repeatable read' | 'serializable';
 	accessMode?: 'read only' | 'read write';
 	deferrable?: boolean;
+	rlsConfig?: {
+		set?: {
+			name: string;
+			value: string;
+			isLocal?: boolean;
+		}[];
+		role?: AnyPgRole;
+	};
 }
 
 export abstract class PgSession<
@@ -114,8 +123,25 @@ export abstract class PgTransaction<
 		return sql.raw(chunks.join(' '));
 	}
 
-	setTransaction(config: PgTransactionConfig): Promise<void> {
-		return this.session.execute(sql`set transaction ${this.getTransactionConfigSQL(config)}`);
+	setTransaction(config: PgTransactionConfig): Promise<void> | void {
+		if (config.accessMode || config.deferrable || config.isolationLevel) {
+			return this.session.execute(sql`set transaction ${this.getTransactionConfigSQL(config)}`);
+		}
+	}
+
+	async executeRLSConfig(config: PgTransactionConfig | undefined): Promise<void> {
+		const rlsConfig = config?.rlsConfig;
+		if (rlsConfig) {
+			if (rlsConfig.set) {
+				for (const { name, value, isLocal } of rlsConfig.set) {
+					await this.session.execute(sql`select set_config(${name}, ${value}, ${isLocal === false ? false : true})`);
+				}
+			}
+
+			if (rlsConfig.role) {
+				await this.session.execute(sql`set local role ${rlsConfig.role}`);
+			}
+		}
 	}
 
 	abstract override transaction<T>(

drizzle-orm/src/pg-core/table.ts

@@ -22,6 +22,9 @@ export type TableConfig = TableConfigBase<PgColumn>;
 /** @internal */
 export const InlineForeignKeys = Symbol.for('drizzle:PgInlineForeignKeys');
 
+/** @internal */
+export const RLSEnabled = Symbol.for('drizzle:RLSEnabled');
+
 export class PgTable<T extends TableConfig = TableConfig> extends Table<T> {
 	static readonly [entityKind]: string = 'PgTable';
 
@@ -33,9 +36,16 @@ export class PgTable<T extends TableConfig = TableConfig> extends Table<T> {
 	/**@internal */
 	[InlineForeignKeys]: ForeignKey[] = [];
 
+	[RLSEnabled]: boolean = false;
+
 	/** @internal */
 	override [Table.Symbol.ExtraConfigBuilder]: ((self: Record<string, PgColumn>) => PgTableExtraConfig) | undefined =
 		undefined;
+
+	enableRLS(): this {
+		this[RLSEnabled] = true;
+		return this;
+	}
 }
 
 export type AnyPgTable<TPartial extends Partial<TableConfig> = {}> = PgTable<UpdateTableConfig<TableConfig, TPartial>>;

drizzle-orm/src/policy.ts

@@ -0,0 +1,24 @@
+import { entityKind } from '~/entity.ts';
+import { SQL, type SQLWrapper } from '~/sql/sql.ts';
+
+export const PolicyName = Symbol.for('drizzle:PolicyName');
+
+export class Policy<
+	TName extends string,
+> implements SQLWrapper {
+	static readonly [entityKind]: string = 'Policy';
+
+	[PolicyName]: TName;
+
+	declare readonly _: {
+		readonly name: TName;
+	};
+
+	constructor(readonly name: TName) {
+		this[PolicyName] = name;
+	}
+
+	getSQL(): SQL {
+		return new SQL([this]);
+	}
+}

drizzle-orm/src/postgres-js/session.ts

@@ -139,7 +139,10 @@ export class PostgresJsSession<
 			const tx = new PostgresJsTransaction(this.dialect, session, this.schema);
 			if (config) {
 				await tx.setTransaction(config);
+
+				await tx.executeRLSConfig(config);
 			}
+
 			return transaction(tx);
 		}) as Promise<T>;
 	}

drizzle-orm/src/role.ts

@@ -0,0 +1,24 @@
+import { entityKind } from '~/entity.ts';
+import { SQL, type SQLWrapper } from '~/sql/sql.ts';
+
+export type AnyRole = Role<string>;
+
+export const RoleName = Symbol.for('drizzle:RoleName');
+
+export class Role<TName extends string> implements SQLWrapper {
+	static readonly [entityKind]: string = 'Role';
+
+	declare readonly _: {
+		readonly name: TName;
+	};
+
+	[RoleName]: TName;
+
+	constructor(readonly name: TName) {
+		this[RoleName] = name;
+	}
+
+	getSQL(): SQL {
+		return new SQL([this]);
+	}
+}

drizzle-orm/src/sql/sql.ts

@@ -1,12 +1,14 @@
 import { entityKind, is } from '~/entity.ts';
+import type { SelectedFields } from '~/operations.ts';
+import { Policy } from '~/policy.ts';
 import { Relation } from '~/relations.ts';
+import { Role } from '~/role.ts';
 import { Subquery, SubqueryConfig } from '~/subquery.ts';
 import { tracer } from '~/tracing.ts';
 import { ViewBaseConfig } from '~/view-common.ts';
 import type { AnyColumn } from '../column.ts';
 import { Column } from '../column.ts';
 import { Table } from '../table.ts';
-import type { SelectedFields } from '~/operations.ts';
 
 /**
  * This class is used to indicate a primitive param value that is used in `sql` tag.
@@ -149,6 +151,10 @@ export class SQL<T = unknown> implements SQLWrapper {
 				return { sql: escapeName(chunk.value), params: [] };
 			}
 
+			if (is(chunk, Role) || is(chunk, Policy)) {
+				return { sql: escapeName(chunk.name), params: [] };
+			}
+
 			if (chunk === undefined) {
 				return { sql: '', params: [] };
 			}

drizzle-orm/src/vercel-postgres/session.ts

@@ -144,6 +144,8 @@ export class VercelPgSession<
 		const tx = new VercelPgTransaction(this.dialect, session, this.schema);
 		await tx.execute(sql`begin${config ? sql` ${tx.getTransactionConfigSQL(config)}` : undefined}`);
 		try {
+			await tx.executeRLSConfig(config);
+
 			const result = await transaction(tx);
 			await tx.execute(sql`commit`);
 			return result;