Rls fix (#3500)

* Fix previous state for RLS * Add unsquash policy push * Make policy unsquash work for push * Fix linked policies that are not in the table for push * Add release notes
drizzle-team · Nov 6, 2024 · 68edb2b · 68edb2b
1 parent b762f63
commit 68edb2b
Show file tree

Hide file tree

Showing 12 changed files with 326 additions and 90 deletions.
diff --git a/changelogs/drizzle-kit/0.27.2.md b/changelogs/drizzle-kit/0.27.2.md
@@ -0,0 +1,3 @@
+- Fix [[BUG]: Undefined properties when using drizzle-kit push](https://github.com/drizzle-team/drizzle-orm/issues/3391)
+- Fix TypeError: Cannot read properties of undefined (reading 'isRLSEnabled')
+- Fix push bugs, when pushing a schema with linked policy to a table from `drizzle-orm/supabase`
diff --git a/drizzle-kit/package.json b/drizzle-kit/package.json
@@ -1,6 +1,6 @@
 {
 	"name": "drizzle-kit",
-	"version": "0.27.1",
+	"version": "0.27.2",
 	"homepage": "https://orm.drizzle.team",
 	"keywords": [
 		"drizzle",
@@ -138,4 +138,4 @@
 			"default": "./api.mjs"
 		}
 	}
-}
+}
diff --git a/drizzle-kit/src/cli/commands/migrate.ts b/drizzle-kit/src/cli/commands/migrate.ts
@@ -343,18 +343,9 @@ export const prepareAndMigratePg = async (config: GenerateConfig) => {
 };
 
 export const preparePgPush = async (
-	schemaPath: string | string[],
-	snapshot: PgSchema,
-	schemaFilter: string[],
-	casing: CasingType | undefined,
+	cur: PgSchema,
+	prev: PgSchema,
 ) => {
-	const { prev, cur } = await preparePgDbPushSnapshot(
-		snapshot,
-		schemaPath,
-		casing,
-		schemaFilter,
-	);
-
 	const validatedPrev = pgSchema.parse(prev);
 	const validatedCur = pgSchema.parse(cur);
 

diff --git a/drizzle-kit/src/cli/commands/pgIntrospect.ts b/drizzle-kit/src/cli/commands/pgIntrospect.ts
@@ -1,7 +1,7 @@
 import { renderWithTask } from 'hanji';
 import { Minimatch } from 'minimatch';
 import { originUUID } from '../../global';
-import type { PgSchema } from '../../serializer/pgSchema';
+import type { PgSchema, PgSchemaInternal } from '../../serializer/pgSchema';
 import { fromDatabase } from '../../serializer/pgSerializer';
 import type { DB } from '../../utils';
 import { Entities } from '../validations/cli';
@@ -12,6 +12,7 @@ export const pgPushIntrospect = async (
 	filters: string[],
 	schemaFilters: string[],
 	entities: Entities,
+	tsSchema?: PgSchemaInternal,
 ) => {
 	const matchers = filters.map((it) => {
 		return new Minimatch(it);
@@ -45,7 +46,7 @@ export const pgPushIntrospect = async (
 	);
 	const res = await renderWithTask(
 		progress,
-		fromDatabase(db, filter, schemaFilters, entities),
+		fromDatabase(db, filter, schemaFilters, entities, undefined, tsSchema),
 	);
 
 	const schema = { id: originUUID, prevId: '', ...res } as PgSchema;

diff --git a/drizzle-kit/src/cli/commands/pgPushUtils.ts b/drizzle-kit/src/cli/commands/pgPushUtils.ts
@@ -250,7 +250,7 @@ export const pgSuggestions = async (db: DB, statements: JsonStatement[]) => {
 				}
 			}
 		}
-		const stmnt = fromJson([statement], 'postgresql');
+		const stmnt = fromJson([statement], 'postgresql', 'push');
 		if (typeof stmnt !== 'undefined') {
 			statementsToExecute.push(...stmnt);
 		}

diff --git a/drizzle-kit/src/cli/commands/push.ts b/drizzle-kit/src/cli/commands/push.ts
@@ -1,5 +1,7 @@
 import chalk from 'chalk';
+import { randomUUID } from 'crypto';
 import { render } from 'hanji';
+import { serializePg } from 'src/serializer';
 import { fromJson } from '../../sqlgenerator';
 import { Select } from '../selector-ui';
 import { Entities } from '../validations/cli';
@@ -169,11 +171,17 @@ export const pgPush = async (
 	const { pgPushIntrospect } = await import('./pgIntrospect');
 
 	const db = await preparePostgresDB(credentials);
-	const { schema } = await pgPushIntrospect(db, tablesFilter, schemasFilter, entities);
+
+	const serialized = await serializePg(schemaPath, casing, schemasFilter);
+
+	const { schema } = await pgPushIntrospect(db, tablesFilter, schemasFilter, entities, serialized);
 
 	const { preparePgPush } = await import('./migrate');
 
-	const statements = await preparePgPush(schemaPath, schema, schemasFilter, casing);
+	const statements = await preparePgPush(
+		{ id: randomUUID(), prevId: schema.id, ...serialized },
+		schema,
+	);
 
 	try {
 		if (statements.sqlStatements.length === 0) {

diff --git a/drizzle-kit/src/serializer/pgSchema.ts b/drizzle-kit/src/serializer/pgSchema.ts
@@ -240,6 +240,7 @@ export const policy = object({
 	using: string().optional(),
 	withCheck: string().optional(),
 	on: string().optional(),
+	schema: string().optional(),
 }).strict();
 
 export const policySquashed = object({
@@ -657,6 +658,16 @@ export const PgSquasher = {
 	squashPolicyPush: (policy: Policy) => {
 		return `${policy.name}--${policy.as}--${policy.for}--${policy.to?.join(',')}--${policy.on}`;
 	},
+	unsquashPolicyPush: (policy: string): Policy => {
+		const splitted = policy.split('--');
+		return {
+			name: splitted[0],
+			as: splitted[1] as Policy['as'],
+			for: splitted[2] as Policy['for'],
+			to: splitted[3].split(','),
+			on: splitted[4] !== 'undefined' ? splitted[4] : undefined,
+		};
+	},
 	squashPK: (pk: PrimaryKey) => {
 		return `${pk.columns.join(',')};${pk.name}`;
 	},

diff --git a/drizzle-kit/src/serializer/pgSerializer.ts b/drizzle-kit/src/serializer/pgSerializer.ts
@@ -652,6 +652,7 @@ export const generatePgSnapshot = (
 		} else {
 			policiesToReturn[policy.name] = {
 				...mappedPolicy,
+				schema: tableConfig.schema ?? 'public',
 				on: `"${tableConfig.schema ?? 'public'}"."${tableConfig.name}"`,
 			};
 		}
@@ -972,9 +973,11 @@ export const fromDatabase = async (
 		count: number,
 		status: IntrospectStatus,
 	) => void,
+	tsSchema?: PgSchemaInternal,
 ): Promise<PgSchemaInternal> => {
 	const result: Record<string, Table> = {};
 	const views: Record<string, View> = {};
+	const policies: Record<string, Policy> = {};
 	const internals: PgKitInternals = { tables: {} };
 
 	const where = schemaFilters.map((t) => `n.nspname = '${t}'`).join(' or ');
@@ -1134,7 +1137,9 @@ WHERE
 		}
 	}
 
-	const wherePolicies = schemaFilters
+	const schemasForLinkedPoliciesInSchema = Object.values(tsSchema?.policies ?? {}).map((it) => it.schema!);
+
+	const wherePolicies = [...schemaFilters, ...schemasForLinkedPoliciesInSchema]
 		.map((t) => `schemaname = '${t}'`)
 		.join(' or ');
 
@@ -1171,6 +1176,16 @@ WHERE
 				[dbPolicy.name]: { ...rest, to: parsedTo, withCheck: parsedWithCheck, using: parsedUsing } as Policy,
 			};
 		}
+
+		if (tsSchema?.policies[dbPolicy.name]) {
+			policies[dbPolicy.name] = {
+				...rest,
+				to: parsedTo,
+				withCheck: parsedWithCheck,
+				using: parsedUsing,
+				on: tsSchema?.policies[dbPolicy.name].on,
+			} as Policy;
+		}
 	}
 
 	if (progressCallback) {
@@ -1907,7 +1922,7 @@ WHERE
 		schemas: schemasObject,
 		sequences: sequencesToReturn,
 		roles: rolesToReturn,
-		policies: {},
+		policies,
 		views: views,
 		_meta: {
 			schemas: {},

diff --git a/drizzle-kit/src/snapshotsDiffer.ts b/drizzle-kit/src/snapshotsDiffer.ts
@@ -988,8 +988,10 @@ export const applyPgSnapshotsDiff = async (
 		const { renamed, created, deleted } = await policyResolver({
 			tableName: entry.name,
 			schema: entry.schema,
-			deleted: entry.policies.deleted.map(PgSquasher.unsquashPolicy),
-			created: entry.policies.added.map(PgSquasher.unsquashPolicy),
+			deleted: entry.policies.deleted.map(
+				action === 'push' ? PgSquasher.unsquashPolicyPush : PgSquasher.unsquashPolicy,
+			),
+			created: entry.policies.added.map(action === 'push' ? PgSquasher.unsquashPolicyPush : PgSquasher.unsquashPolicy),
 		});
 
 		if (created.length > 0) {
@@ -1043,7 +1045,9 @@ export const applyPgSnapshotsDiff = async (
 					] || [];
 
 					const newName = columnChangeFor(policyKey, rens);
-					const unsquashedPolicy = PgSquasher.unsquashPolicy(policy);
+					const unsquashedPolicy = action === 'push'
+						? PgSquasher.unsquashPolicyPush(policy)
+						: PgSquasher.unsquashPolicy(policy);
 					unsquashedPolicy.name = newName;
 					policy = PgSquasher.squashPolicy(unsquashedPolicy);
 					return newName;
@@ -1068,8 +1072,12 @@ export const applyPgSnapshotsDiff = async (
 	}[];
 
 	const { renamed: indPolicyRenames, created, deleted } = await indPolicyResolver({
-		deleted: indPolicyRes.deleted.map((t) => PgSquasher.unsquashPolicy(t.values)),
-		created: indPolicyRes.added.map((t) => PgSquasher.unsquashPolicy(t.values)),
+		deleted: indPolicyRes.deleted.map((t) =>
+			action === 'push' ? PgSquasher.unsquashPolicyPush(t.values) : PgSquasher.unsquashPolicy(t.values)
+		),
+		created: indPolicyRes.added.map((t) =>
+			action === 'push' ? PgSquasher.unsquashPolicyPush(t.values) : PgSquasher.unsquashPolicy(t.values)
+		),
 	});
 
 	if (created.length > 0) {
@@ -1421,10 +1429,14 @@ export const applyPgSnapshotsDiff = async (
 	typedResult.alteredPolicies.forEach(({ values }) => {
 		// return prepareAlterIndPolicyJson(json1.policies[it.name], json2.policies[it.name]);
 
-		const policy = PgSquasher.unsquashPolicy(values);
+		const policy = action === 'push' ? PgSquasher.unsquashPolicyPush(values) : PgSquasher.unsquashPolicy(values);
 
-		const newPolicy = PgSquasher.unsquashPolicy(json2.policies[policy.name].values);
-		const oldPolicy = PgSquasher.unsquashPolicy(json1.policies[policy.name].values);
+		const newPolicy = action === 'push'
+			? PgSquasher.unsquashPolicyPush(json2.policies[policy.name].values)
+			: PgSquasher.unsquashPolicy(json2.policies[policy.name].values);
+		const oldPolicy = action === 'push'
+			? PgSquasher.unsquashPolicyPush(json2.policies[policy.name].values)
+			: PgSquasher.unsquashPolicy(json1.policies[policy.name].values);
 
 		if (newPolicy.as !== oldPolicy.as) {
 			jsonDropIndPoliciesStatements.push(
@@ -1494,8 +1506,12 @@ export const applyPgSnapshotsDiff = async (
 	alteredTables.forEach((it) => {
 		// handle policies
 		Object.keys(it.alteredPolicies).forEach((policyName: string) => {
-			const newPolicy = PgSquasher.unsquashPolicy(it.alteredPolicies[policyName].__new);
-			const oldPolicy = PgSquasher.unsquashPolicy(it.alteredPolicies[policyName].__old);
+			const newPolicy = action === 'push'
+				? PgSquasher.unsquashPolicyPush(it.alteredPolicies[policyName].__new)
+				: PgSquasher.unsquashPolicy(it.alteredPolicies[policyName].__new);
+			const oldPolicy = action === 'push'
+				? PgSquasher.unsquashPolicyPush(it.alteredPolicies[policyName].__old)
+				: PgSquasher.unsquashPolicy(it.alteredPolicies[policyName].__old);
 
 			if (newPolicy.as !== oldPolicy.as) {
 				jsonDropPoliciesStatements.push(
@@ -1569,7 +1585,8 @@ export const applyPgSnapshotsDiff = async (
 			}
 
 			// handle table.isRLSEnabled
-			if (table.isRLSEnabled !== tableInPreviousState.isRLSEnabled) {
+			const wasRlsEnabled = tableInPreviousState ? tableInPreviousState.isRLSEnabled : false;
+			if (table.isRLSEnabled !== wasRlsEnabled) {
 				if (table.isRLSEnabled) {
 					// was force enabled
 					jsonEnableRLSStatements.push({ type: 'enable_rls', tableName: table.name, schema: table.schema });
@@ -1762,7 +1779,11 @@ export const applyPgSnapshotsDiff = async (
 
 	jsonCreatePoliciesStatements.push(...([] as JsonCreatePolicyStatement[]).concat(
 		...(createdTables.map((it) =>
-			prepareCreatePolicyJsons(it.name, it.schema, Object.values(it.policies).map(PgSquasher.unsquashPolicy))
+			prepareCreatePolicyJsons(
+				it.name,
+				it.schema,
+				Object.values(it.policies).map(action === 'push' ? PgSquasher.unsquashPolicyPush : PgSquasher.unsquashPolicy),
+			)
 		)),
 	));
 	const createViews: JsonCreatePgViewStatement[] = [];
@@ -2037,7 +2058,7 @@ export const applyPgSnapshotsDiff = async (
 		return true;
 	});
 
-	const sqlStatements = fromJson(filteredEnumsJsonStatements, 'postgresql');
+	const sqlStatements = fromJson(filteredEnumsJsonStatements, 'postgresql', action);
 
 	const uniqueSqlStatements: string[] = [];
 	sqlStatements.forEach((ss) => {

diff --git a/drizzle-kit/src/sqlgenerator.ts b/drizzle-kit/src/sqlgenerator.ts
@@ -270,9 +270,13 @@ class PgAlterPolicyConvertor extends Convertor {
 	override can(statement: JsonStatement, dialect: Dialect): boolean {
 		return statement.type === 'alter_policy' && dialect === 'postgresql';
 	}
-	override convert(statement: JsonAlterPolicyStatement): string | string[] {
-		const newPolicy = PgSquasher.unsquashPolicy(statement.newData);
-		const oldPolicy = PgSquasher.unsquashPolicy(statement.oldData);
+	override convert(statement: JsonAlterPolicyStatement, _dialect: any, action?: string): string | string[] {
+		const newPolicy = action === 'push'
+			? PgSquasher.unsquashPolicyPush(statement.newData)
+			: PgSquasher.unsquashPolicy(statement.newData);
+		const oldPolicy = action === 'push'
+			? PgSquasher.unsquashPolicyPush(statement.oldData)
+			: PgSquasher.unsquashPolicy(statement.oldData);
 
 		const tableNameWithSchema = statement.schema
 			? `"${statement.schema}"."${statement.tableName}"`
@@ -1408,7 +1412,7 @@ class PgDropTableConvertor extends Convertor {
 		return statement.type === 'drop_table' && dialect === 'postgresql';
 	}
 
-	convert(statement: JsonDropTableStatement) {
+	convert(statement: JsonDropTableStatement, _d: any, action?: string) {
 		const { tableName, schema, policies } = statement;
 
 		const tableNameWithSchema = schema
@@ -1420,7 +1424,9 @@ class PgDropTableConvertor extends Convertor {
 			return dropPolicyConvertor.convert({
 				type: 'drop_policy',
 				tableName,
-				data: PgSquasher.unsquashPolicy(p),
+				data: action === 'push'
+					? PgSquasher.unsquashPolicyPush(p)
+					: PgSquasher.unsquashPolicy(p),
 				schema,
 			}) as string;
 		}) ?? [];
@@ -3334,17 +3340,6 @@ convertors.push(new MySqlAlterTableCreateCompositePrimaryKeyConvertor());
 convertors.push(new MySqlAlterTableAddPk());
 convertors.push(new MySqlAlterTableAlterCompositePrimaryKeyConvertor());
 
-export function fromJson(
-	statements: JsonStatement[],
-	dialect: Exclude<Dialect, 'sqlite' | 'turso'>,
-): string[];
-export function fromJson(
-	statements: JsonStatement[],
-	dialect: 'sqlite' | 'turso',
-	action?: 'push',
-	json2?: SQLiteSchemaSquashed,
-): string[];
-
 export function fromJson(
 	statements: JsonStatement[],
 	dialect: Dialect,