[Pg] Fixed imports and added enableRLS() method on PgTable

drizzle-team · Nov 10, 2023 · 095d126 · 095d126
1 parent d3ea325
commit 095d126
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 17 deletions.
diff --git a/drizzle-orm/src/pg-core/policy.ts b/drizzle-orm/src/pg-core/policy.ts
@@ -1,7 +1,7 @@
-import { entityKind } from '~/entity';
-import type { AnyPgTable } from '~/pg-core/table';
-import { SQL, type SQLWrapper } from '~/sql/sql';
-import type { AnyPgRole } from './role';
+import { entityKind } from '~/entity.ts';
+import type { AnyPgTable } from '~/pg-core/table.ts';
+import type { SQL } from '~/sql/sql.ts';
+import type { AnyPgRole } from './role.ts';
 
 export type PgPolicyFor = 'select' | 'insert' | 'update' | 'delete' | 'all';
 
@@ -20,11 +20,12 @@ export type PgPolicyConfig = {
 export const PolicyName = Symbol.for('drizzle:PolicyName');
 export const PolicyTable = Symbol.for('drizzle:PolicyTable');
 
+// Do we need to implement SQLWrapper here? I'm not entirely sure if drizzle-kit will need it
 export class PgPolicy<
 	TName extends string,
 	TTable extends AnyPgTable,
 	TConfig extends PgPolicyConfig | undefined,
-> implements SQLWrapper {
+> {
 	static readonly [entityKind]: string = 'PgPolicy';
 
 	[PolicyName]: TName;
@@ -43,10 +44,6 @@ export class PgPolicy<
 		this[PolicyTable] = table;
 		this.config = config;
 	}
-
-	getSQL(): SQL<unknown> {
-		return new SQL([this]);
-	}
 }
 
 export function pgPolicy<TName extends string, TTable extends AnyPgTable, TConfig extends PgPolicyConfig>(

diff --git a/drizzle-orm/src/pg-core/role.ts b/drizzle-orm/src/pg-core/role.ts
@@ -1,5 +1,4 @@
-import { entityKind } from '~/entity';
-import { SQL, type SQLWrapper } from '~/sql/sql.ts';
+import { entityKind } from '~/entity.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
@@ -23,7 +22,8 @@ export type AnyPgRole = PgRole<string, PgRoleConfig>;
 
 export const RoleName = Symbol.for('drizzle:RoleName');
 
-export class PgRole<TName extends string, TConfig extends PgRoleConfig> implements SQLWrapper {
+// Do we need to implement SQLWrapper here? I'm not entirely sure if drizzle-kit will need it
+export class PgRole<TName extends string, TConfig extends PgRoleConfig> {
 	static readonly [entityKind]: string = 'PgRole';
 
 	declare readonly _: {
@@ -39,11 +39,6 @@ export class PgRole<TName extends string, TConfig extends PgRoleConfig> implemen
 		this[RoleName] = name;
 		this.config = config;
 	}
-
-	// Maybe this will be moved to ~/sql/sql.ts like Table, Column and Subquery?
-	getSQL(): SQL<unknown> {
-		return new SQL([this]);
-	}
 }
 
 export function pgRole<TName extends string, TConfig extends PgRoleConfig>(

diff --git a/drizzle-orm/src/pg-core/table.ts b/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>>;