From 4045fbc4f51d6104fb57780cf217cffca165a1d6 Mon Sep 17 00:00:00 2001 From: Dan Kochetov Date: Fri, 23 Aug 2024 20:43:30 +0300 Subject: [PATCH 01/24] Add pgPolicy --- drizzle-orm/src/pg-core/index.ts | 1 + drizzle-orm/src/pg-core/policies.ts | 37 ++++++++++++++ drizzle-orm/src/pg-core/table.ts | 2 + drizzle-orm/src/pg-core/utils.ts | 5 ++ integration-tests/tests/pg/pg-common.ts | 64 +++++++++++++++++++++++++ 5 files changed, 109 insertions(+) create mode 100644 drizzle-orm/src/pg-core/policies.ts diff --git a/drizzle-orm/src/pg-core/index.ts b/drizzle-orm/src/pg-core/index.ts index 084633c4a..2e2d5399d 100644 --- a/drizzle-orm/src/pg-core/index.ts +++ b/drizzle-orm/src/pg-core/index.ts @@ -5,6 +5,7 @@ export * from './db.ts'; export * from './dialect.ts'; export * from './foreign-keys.ts'; export * from './indexes.ts'; +export * from './policies.ts'; export * from './primary-keys.ts'; export * from './query-builders/index.ts'; export * from './schema.ts'; diff --git a/drizzle-orm/src/pg-core/policies.ts b/drizzle-orm/src/pg-core/policies.ts new file mode 100644 index 000000000..7f49662b6 --- /dev/null +++ b/drizzle-orm/src/pg-core/policies.ts @@ -0,0 +1,37 @@ +import { entityKind } from '~/entity.ts'; +import type { SQL } from '~/sql/sql.ts'; + +export interface PgPolicyConfig { + as?: 'permissive' | 'restrictive'; + for?: 'all' | 'select' | 'insert' | 'update' | 'delete'; + to?: 'PUBLIC' | 'CURRENT_ROLE' | 'CURRENT_USER' | 'SESSION_USER' | (string & {}); + using?: SQL; + withCheck?: SQL; +} + +export class PgPolicy implements PgPolicyConfig { + static readonly [entityKind]: string = 'PgPolicy'; + + readonly as: PgPolicyConfig['as']; + readonly for: PgPolicyConfig['for']; + readonly to: PgPolicyConfig['to']; + readonly using: PgPolicyConfig['using']; + readonly withCheck: PgPolicyConfig['withCheck']; + + constructor( + readonly name: string, + config?: PgPolicyConfig, + ) { + if (config) { + this.as = config.as; + this.for = config.for; + this.to = config.to; + this.using = config.using; + this.withCheck = config.withCheck; + } + } +} + +export function pgPolicy(name: string, config?: PgPolicyConfig) { + return new PgPolicy(name, config); +} diff --git a/drizzle-orm/src/pg-core/table.ts b/drizzle-orm/src/pg-core/table.ts index 78cd52019..7588ed2bd 100644 --- a/drizzle-orm/src/pg-core/table.ts +++ b/drizzle-orm/src/pg-core/table.ts @@ -5,6 +5,7 @@ import type { CheckBuilder } from './checks.ts'; import type { PgColumn, PgColumnBuilder, PgColumnBuilderBase } from './columns/common.ts'; import type { ForeignKey, ForeignKeyBuilder } from './foreign-keys.ts'; import type { AnyIndexBuilder } from './indexes.ts'; +import type { PgPolicy } from './policies.ts'; import type { PrimaryKeyBuilder } from './primary-keys.ts'; import type { UniqueConstraintBuilder } from './unique-constraint.ts'; @@ -15,6 +16,7 @@ export type PgTableExtraConfig = Record< | ForeignKeyBuilder | PrimaryKeyBuilder | UniqueConstraintBuilder + | PgPolicy >; export type TableConfig = TableConfigBase; diff --git a/drizzle-orm/src/pg-core/utils.ts b/drizzle-orm/src/pg-core/utils.ts index 90378f249..43978ae58 100644 --- a/drizzle-orm/src/pg-core/utils.ts +++ b/drizzle-orm/src/pg-core/utils.ts @@ -7,6 +7,7 @@ import type { AnyPgColumn } from './columns/index.ts'; import { type ForeignKey, ForeignKeyBuilder } from './foreign-keys.ts'; import type { Index } from './indexes.ts'; import { IndexBuilder } from './indexes.ts'; +import { PgPolicy } from './policies.ts'; import { type PrimaryKey, PrimaryKeyBuilder } from './primary-keys.ts'; import { type UniqueConstraint, UniqueConstraintBuilder } from './unique-constraint.ts'; import { PgViewConfig } from './view-common.ts'; @@ -21,6 +22,7 @@ export function getTableConfig(table: TTable) { const uniqueConstraints: UniqueConstraint[] = []; const name = table[Table.Symbol.Name]; const schema = table[Table.Symbol.Schema]; + const policies: PgPolicy[] = []; const extraConfigBuilder = table[PgTable.Symbol.ExtraConfigBuilder]; @@ -37,6 +39,8 @@ export function getTableConfig(table: TTable) { primaryKeys.push(builder.build(table)); } else if (is(builder, ForeignKeyBuilder)) { foreignKeys.push(builder.build(table)); + } else if (is(builder, PgPolicy)) { + policies.push(builder); } } } @@ -50,6 +54,7 @@ export function getTableConfig(table: TTable) { uniqueConstraints, name, schema, + policies, }; } diff --git a/integration-tests/tests/pg/pg-common.ts b/integration-tests/tests/pg/pg-common.ts index c48a533f9..a7b7941af 100644 --- a/integration-tests/tests/pg/pg-common.ts +++ b/integration-tests/tests/pg/pg-common.ts @@ -18,6 +18,7 @@ import { gte, ilike, inArray, + is, lt, max, min, @@ -54,8 +55,11 @@ import { macaddr, macaddr8, numeric, + PgDialect, pgEnum, pgMaterializedView, + PgPolicy, + pgPolicy, pgSchema, pgTable, pgTableCreator, @@ -4660,5 +4664,65 @@ export function tests() { jsonbNumberField: testNumber, }]); }); + + test('policy', () => { + { + const policy = pgPolicy('test policy'); + + expect(is(policy, PgPolicy)).toBe(true); + expect(policy.name).toBe('test policy'); + } + + { + const policy = pgPolicy('test policy', { + as: 'permissive', + for: 'all', + to: 'PUBLIC', + using: sql`1=1`, + withCheck: sql`1=1`, + }); + + expect(is(policy, PgPolicy)).toBe(true); + expect(policy.name).toBe('test policy'); + expect(policy.as).toBe('permissive'); + expect(policy.for).toBe('all'); + expect(policy.to).toBe('PUBLIC'); + const dialect = new PgDialect(); + expect(is(policy.using, SQL)).toBe(true); + expect(dialect.sqlToQuery(policy.using!).sql).toBe('1=1'); + expect(is(policy.withCheck, SQL)).toBe(true); + expect(dialect.sqlToQuery(policy.withCheck!).sql).toBe('1=1'); + } + + { + const policy = pgPolicy('test policy', { + to: 'custom value', + }); + + expect(policy.to).toBe('custom value'); + } + + { + const p1 = pgPolicy('test policy'); + const p2 = pgPolicy('test policy 2', { + as: 'permissive', + for: 'all', + to: 'PUBLIC', + using: sql`1=1`, + withCheck: sql`1=1`, + }); + const table = pgTable('table_with_policy', { + id: serial('id').primaryKey(), + name: text('name').notNull(), + }, () => ({ + p1, + p2, + })); + const config = getTableConfig(table); + expect(config.policies).toHaveLength(2); + expect(config.policies[0]).toBe(p1); + expect(config.policies[1]).toBe(p2); + } + }); }); } From fb1e950c1c584cff82fb0ef40a5b2997cf672080 Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Tue, 27 Aug 2024 14:59:21 +0300 Subject: [PATCH 02/24] Add all policies for generate logic --- drizzle-kit/src/api.ts | 3 + drizzle-kit/src/cli/commands/migrate.ts | 21 +- drizzle-kit/src/jsonDiffer.js | 68 +++++++ drizzle-kit/src/jsonStatements.ts | 102 +++++++++- drizzle-kit/src/serializer/pgSchema.ts | 31 +++ drizzle-kit/src/serializer/pgSerializer.ts | 16 ++ drizzle-kit/src/snapshotsDiffer.ts | 220 ++++++++++++++++++++- drizzle-kit/src/sqlgenerator.ts | 84 ++++++++ drizzle-kit/tests/rls/pg-policy.test.ts | 23 +++ drizzle-kit/tests/schemaDiffer.ts | 73 ++++++- drizzle-kit/vitest.config.ts | 3 +- 11 files changed, 638 insertions(+), 6 deletions(-) create mode 100644 drizzle-kit/tests/rls/pg-policy.test.ts diff --git a/drizzle-kit/src/api.ts b/drizzle-kit/src/api.ts index 00cdb1b61..d0cceda52 100644 --- a/drizzle-kit/src/api.ts +++ b/drizzle-kit/src/api.ts @@ -6,6 +6,7 @@ import { PgDatabase } from 'drizzle-orm/pg-core'; import { columnsResolver, enumsResolver, + policyResolver, schemasResolver, sequencesResolver, tablesResolver, @@ -71,6 +72,7 @@ export const generateMigration = async ( schemasResolver, enumsResolver, sequencesResolver, + policyResolver, tablesResolver, columnsResolver, validatedPrev, @@ -114,6 +116,7 @@ export const pushSchema = async ( schemasResolver, enumsResolver, sequencesResolver, + policyResolver, tablesResolver, columnsResolver, validatedPrev, diff --git a/drizzle-kit/src/cli/commands/migrate.ts b/drizzle-kit/src/cli/commands/migrate.ts index 8ef469fa1..b5f369bc7 100644 --- a/drizzle-kit/src/cli/commands/migrate.ts +++ b/drizzle-kit/src/cli/commands/migrate.ts @@ -14,7 +14,7 @@ import path, { join } from 'path'; import { TypeOf } from 'zod'; import type { CommonSchema } from '../../schemaValidator'; import { MySqlSchema, mysqlSchema, squashMysqlScheme } from '../../serializer/mysqlSchema'; -import { PgSchema, pgSchema, squashPgScheme } from '../../serializer/pgSchema'; +import { PgSchema, pgSchema, Policy, squashPgScheme } from '../../serializer/pgSchema'; import { SQLiteSchema, sqliteSchema, squashSqliteScheme } from '../../serializer/sqliteSchema'; import { applyMysqlSnapshotsDiff, @@ -113,6 +113,23 @@ export const sequencesResolver = async ( } }; +export const policyResolver = async ( + input: ColumnsResolverInput, +): Promise> => { + const result = await promptColumnsConflicts( + input.tableName, + input.created, + input.deleted, + ); + return { + tableName: input.tableName, + schema: input.schema, + created: result.created, + deleted: result.deleted, + renamed: result.renamed, + }; +}; + export const enumsResolver = async ( input: ResolverInput, ): Promise> => { @@ -195,6 +212,7 @@ export const prepareAndMigratePg = async (config: GenerateConfig) => { schemasResolver, enumsResolver, sequencesResolver, + policyResolver, tablesResolver, columnsResolver, validatedPrev, @@ -238,6 +256,7 @@ export const preparePgPush = async ( schemasResolver, enumsResolver, sequencesResolver, + policyResolver, tablesResolver, columnsResolver, validatedPrev, diff --git a/drizzle-kit/src/jsonDiffer.js b/drizzle-kit/src/jsonDiffer.js index 113d7e0a4..7f72b0c26 100644 --- a/drizzle-kit/src/jsonDiffer.js +++ b/drizzle-kit/src/jsonDiffer.js @@ -146,6 +146,49 @@ export function diffColumns(left, right) { return alteredTables; } +export function diffPolicies(left, right) { + left = JSON.parse(JSON.stringify(left)); + right = JSON.parse(JSON.stringify(right)); + const result = diff(left, right) ?? {}; + + const alteredTables = Object.fromEntries( + Object.entries(result) + .filter((it) => { + return !(it[0].includes('__added') || it[0].includes('__deleted')); + }) + .map((tableEntry) => { + // const entry = { name: it, ...result[it] } + const deletedPolicies = Object.entries(tableEntry[1].policies ?? {}) + .filter((it) => { + return it[0].endsWith('__deleted'); + }) + .map((it) => { + return it[1]; + }); + + const addedPolicies = Object.entries(tableEntry[1].policies ?? {}) + .filter((it) => { + return it[0].endsWith('__added'); + }) + .map((it) => { + return it[1]; + }); + + tableEntry[1].policies = { + added: addedPolicies, + deleted: deletedPolicies, + }; + const table = left[tableEntry[0]]; + return [ + tableEntry[0], + { name: table.name, schema: table.schema, ...tableEntry[1] }, + ]; + }), + ); + + return alteredTables; +} + export function applyJsonDiff(json1, json2) { json1 = JSON.parse(JSON.stringify(json1)); json2 = JSON.parse(JSON.stringify(json2)); @@ -286,6 +329,28 @@ const findAlternationsInTable = (table) => { }), ); + const deletedPolicies = Object.fromEntries( + Object.entries(table.policies__deleted || {}) + .concat( + Object.entries(table.policies || {}).filter((it) => it[0].includes('__deleted')), + ) + .map((entry) => [entry[0].replace('__deleted', ''), entry[1]]), + ); + + const addedPolicies = Object.fromEntries( + Object.entries(table.policies__added || {}) + .concat( + Object.entries(table.policies || {}).filter((it) => it[0].includes('__added')), + ) + .map((entry) => [entry[0].replace('__added', ''), entry[1]]), + ); + + const alteredPolicies = Object.fromEntries( + Object.entries(table.policies || {}).filter((it) => { + return !it[0].endsWith('__deleted') && !it[0].endsWith('__added'); + }), + ); + const deletedForeignKeys = Object.fromEntries( Object.entries(table.foreignKeys__deleted || {}) .concat( @@ -364,6 +429,9 @@ const findAlternationsInTable = (table) => { addedUniqueConstraints, deletedUniqueConstraints, alteredUniqueConstraints, + deletedPolicies, + addedPolicies, + alteredPolicies, }; }; diff --git a/drizzle-kit/src/jsonStatements.ts b/drizzle-kit/src/jsonStatements.ts index ad2afea7f..17dc76456 100644 --- a/drizzle-kit/src/jsonStatements.ts +++ b/drizzle-kit/src/jsonStatements.ts @@ -3,7 +3,7 @@ import { table } from 'console'; import { warning } from './cli/views'; import { CommonSquashedSchema, Dialect } from './schemaValidator'; import { MySqlKitInternals, MySqlSchema, MySqlSquasher } from './serializer/mysqlSchema'; -import { Index, PgSchema, PgSquasher } from './serializer/pgSchema'; +import { Index, PgSchema, PgSquasher, Policy } from './serializer/pgSchema'; import { SQLiteKitInternals, SQLiteSquasher } from './serializer/sqliteSchema'; import { AlteredColumn, Column, Sequence, Table } from './snapshotsDiffer'; @@ -153,6 +153,36 @@ export interface JsonSqliteAddColumnStatement { referenceData?: string; } +export interface JsonCreatePolicyStatement { + type: 'create_policy'; + tableName: string; + data: string; + schema: string; +} + +export interface JsonDropPolicyStatement { + type: 'drop_policy'; + tableName: string; + data: string; + schema: string; +} + +export interface JsonRenamePolicyStatement { + type: 'rename_policy'; + tableName: string; + oldName: string; + newName: string; + schema: string; +} + +export interface JsonAlterPolicyStatement { + type: 'alter_policy'; + tableName: string; + oldData: string; + newData: string; + schema: string; +} + export interface JsonCreateIndexStatement { type: 'create_index'; tableName: string; @@ -555,7 +585,11 @@ export type JsonStatement = | JsonDropSequenceStatement | JsonCreateSequenceStatement | JsonMoveSequenceStatement - | JsonRenameSequenceStatement; + | JsonRenameSequenceStatement + | JsonDropPolicyStatement + | JsonCreatePolicyStatement + | JsonAlterPolicyStatement + | JsonRenamePolicyStatement; export const preparePgCreateTableJson = ( table: Table, @@ -1883,6 +1917,70 @@ export const prepareSqliteAlterColumns = ( return [...dropPkStatements, ...setPkStatements, ...statements]; }; +export const prepareRenamePolicyJsons = ( + tableName: string, + schema: string, + renames: { + from: Policy; + to: Policy; + }[], +): JsonRenamePolicyStatement[] => { + return renames.map((it) => { + return { + type: 'rename_policy', + tableName: tableName, + oldName: it.from.name, + newName: it.to.name, + schema, + }; + }); +}; + +export const prepareCreatePolicyJsons = ( + tableName: string, + schema: string, + policies: Record, +): JsonCreatePolicyStatement[] => { + return Object.values(policies).map((policyData) => { + return { + type: 'create_policy', + tableName, + data: policyData, + schema, + }; + }); +}; + +export const prepareDropPolicyJsons = ( + tableName: string, + schema: string, + policies: Record, +): JsonDropPolicyStatement[] => { + return Object.values(policies).map((policyData) => { + return { + type: 'drop_policy', + tableName, + data: policyData, + schema, + }; + }); +}; + +export const prepareAlterPolicyJson = ( + tableName: string, + schema: string, + oldPolicy: string, + newPolicy: string, +): JsonAlterPolicyStatement => { + return { + type: 'alter_policy', + tableName, + oldData: oldPolicy, + newData: newPolicy, + schema, + }; +}; + export const preparePgCreateIndexesJson = ( tableName: string, schema: string, diff --git a/drizzle-kit/src/serializer/pgSchema.ts b/drizzle-kit/src/serializer/pgSchema.ts index 5860a6fef..e03354ebe 100644 --- a/drizzle-kit/src/serializer/pgSchema.ts +++ b/drizzle-kit/src/serializer/pgSchema.ts @@ -220,6 +220,15 @@ const uniqueConstraint = object({ nullsNotDistinct: boolean(), }).strict(); +const policy = object({ + name: string(), + as: enumType(['permissive', 'restrictive']).optional(), + for: enumType(['all', 'select', 'insert', 'update', 'delete']).optional(), + to: string().optional(), + using: string().optional(), + withCheck: string().optional(), +}).strict(); + const tableV4 = object({ name: string(), schema: string(), @@ -266,6 +275,7 @@ const table = object({ foreignKeys: record(string(), fk), compositePrimaryKeys: record(string(), compositePK), uniqueConstraints: record(string(), uniqueConstraint).default({}), + policies: record(string(), policy).default({}), }).strict(); const schemaHash = object({ @@ -385,6 +395,7 @@ const tableSquashed = object({ foreignKeys: record(string(), string()), compositePrimaryKeys: record(string(), string()), uniqueConstraints: record(string(), string()), + policies: record(string(), string()), }).strict(); const tableSquashedV4 = object({ @@ -445,6 +456,7 @@ export type Index = TypeOf; export type ForeignKey = TypeOf; export type PrimaryKey = TypeOf; export type UniqueConstraint = TypeOf; +export type Policy = TypeOf; export type PgKitInternals = TypeOf; export type PgSchemaV1 = TypeOf; @@ -549,6 +561,20 @@ export const PgSquasher = { fk.onUpdate ?? '' };${fk.onDelete ?? ''};${fk.schemaTo || 'public'}`; }, + squashPolicy: (policy: Policy) => { + return `${policy.name}--${policy.as}--${policy.for}--${policy.to}--${policy.using}--${policy.withCheck}`; + }, + unsquashPolicy: (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], + using: splitted[4] !== 'undefined' ? splitted[4] : undefined, + withCheck: splitted[5] !== 'undefined' ? splitted[5] : undefined, + }; + }, squashPK: (pk: PrimaryKey) => { return `${pk.columns.join(',')};${pk.name}`; }, @@ -671,6 +697,10 @@ export const squashPgScheme = ( }, ); + const squashedPolicies = mapValues(it[1].policies, (policy) => { + return PgSquasher.squashPolicy(policy); + }); + return [ it[0], { @@ -681,6 +711,7 @@ export const squashPgScheme = ( foreignKeys: squashedFKs, compositePrimaryKeys: squashedPKs, uniqueConstraints: squashedUniqueConstraints, + policies: squashedPolicies, }, ]; }), diff --git a/drizzle-kit/src/serializer/pgSerializer.ts b/drizzle-kit/src/serializer/pgSerializer.ts index b479e59e2..91076d78f 100644 --- a/drizzle-kit/src/serializer/pgSerializer.ts +++ b/drizzle-kit/src/serializer/pgSerializer.ts @@ -25,6 +25,7 @@ import type { IndexColumnType, PgKitInternals, PgSchemaInternal, + Policy, PrimaryKey, Sequence, Table, @@ -136,6 +137,7 @@ export const generatePgSnapshot = ( schema, primaryKeys, uniqueConstraints, + policies, } = getTableConfig(table); if (schemaFilter && !schemaFilter.includes(schema ?? 'public')) { @@ -147,6 +149,7 @@ export const generatePgSnapshot = ( const foreignKeysObject: Record = {}; const primaryKeysObject: Record = {}; const uniqueConstraintObject: Record = {}; + const policiesObject: Record = {}; columns.forEach((column) => { const notNull: boolean = column.notNull; @@ -492,6 +495,17 @@ export const generatePgSnapshot = ( }; }); + policies.forEach((policy) => { + policiesObject[policy.name] = { + name: policy.name, + as: policy.as ?? 'permissive', + for: policy.for ?? 'all', + to: policy.to ?? 'PUBLIC', + using: is(policy.using, SQL) ? sqlToStr(policy.using) : undefined, + withCheck: is(policy.withCheck, SQL) ? sqlToStr(policy.withCheck) : undefined, + }; + }); + const tableKey = `${schema ?? 'public'}.${tableName}`; result[tableKey] = { @@ -502,6 +516,7 @@ export const generatePgSnapshot = ( foreignKeys: foreignKeysObject, compositePrimaryKeys: primaryKeysObject, uniqueConstraints: uniqueConstraintObject, + policies: policiesObject, }; } @@ -1181,6 +1196,7 @@ export const fromDatabase = async ( foreignKeys: foreignKeysToReturn, compositePrimaryKeys: primaryKeys, uniqueConstraints: uniqueConstrains, + policies: {}, }; } catch (e) { rej(e); diff --git a/drizzle-kit/src/snapshotsDiffer.ts b/drizzle-kit/src/snapshotsDiffer.ts index 9ad2d9e32..a333556fa 100644 --- a/drizzle-kit/src/snapshotsDiffer.ts +++ b/drizzle-kit/src/snapshotsDiffer.ts @@ -22,16 +22,20 @@ import { _prepareSqliteAddColumns, JsonAddColumnStatement, JsonAlterCompositePK, + JsonAlterPolicyStatement, JsonAlterTableSetSchema, JsonAlterUniqueConstraint, JsonCreateCompositePK, + JsonCreatePolicyStatement, JsonCreateReferenceStatement, JsonCreateUniqueConstraint, JsonDeleteCompositePK, JsonDeleteUniqueConstraint, JsonDropColumnStatement, + JsonDropPolicyStatement, JsonReferenceStatement, JsonRenameColumnStatement, + JsonRenamePolicyStatement, JsonSqliteAddColumnStatement, JsonStatement, prepareAddCompositePrimaryKeyMySql, @@ -43,10 +47,12 @@ import { prepareAlterCompositePrimaryKeyMySql, prepareAlterCompositePrimaryKeyPg, prepareAlterCompositePrimaryKeySqlite, + prepareAlterPolicyJson, prepareAlterReferencesJson, prepareAlterSequenceJson, prepareCreateEnumJson, prepareCreateIndexesJson, + prepareCreatePolicyJsons, prepareCreateReferencesJson, prepareCreateSchemasJson, prepareCreateSequenceJson, @@ -57,6 +63,7 @@ import { prepareDeleteUniqueConstraintPg as prepareDeleteUniqueConstraint, prepareDropEnumJson, prepareDropIndexesJson, + prepareDropPolicyJsons, prepareDropReferencesJson, prepareDropSequenceJson, prepareDropTableJson, @@ -68,6 +75,7 @@ import { preparePgCreateTableJson, prepareRenameColumns, prepareRenameEnumJson, + prepareRenamePolicyJsons, prepareRenameSchemasJson, prepareRenameSequenceJson, prepareRenameTableJson, @@ -78,7 +86,14 @@ import { import { Named, NamedWithSchema } from './cli/commands/migrate'; import { mapEntries, mapKeys, mapValues } from './global'; import { MySqlSchema, MySqlSchemaSquashed, MySqlSquasher } from './serializer/mysqlSchema'; -import { PgSchema, PgSchemaSquashed, PgSquasher, sequenceSchema, sequenceSquashed } from './serializer/pgSchema'; +import { + PgSchema, + PgSchemaSquashed, + PgSquasher, + Policy, + sequenceSchema, + sequenceSquashed, +} from './serializer/pgSchema'; import { SQLiteSchema, SQLiteSchemaSquashed, SQLiteSquasher } from './serializer/sqliteSchema'; import { copy, prepareMigrationMeta } from './utils'; @@ -246,6 +261,15 @@ export const alteredTableScheme = object({ __old: string(), }), ), + addedPolicies: record(string(), string()), + deletedPolicies: record(string(), string()), + alteredPolicies: record( + string(), + object({ + __new: string(), + __old: string(), + }), + ), }).strict(); export const diffResultScheme = object({ @@ -381,6 +405,9 @@ export const applyPgSnapshotsDiff = async ( sequencesResolver: ( input: ResolverInput, ) => Promise>, + policyResolver: ( + input: ColumnsResolverInput, + ) => Promise>, tablesResolver: ( input: ResolverInput, ) => Promise>, @@ -717,6 +744,101 @@ export const applyPgSnapshotsDiff = async ( }, ); + //// Policies + + const policyRes = diffColumns(tablesPatchedSnap1.tables, json2.tables); + + const policyRenames = [] as { + table: string; + schema: string; + renames: { from: Policy; to: Policy }[]; + }[]; + + const policyCreates = [] as { + table: string; + schema: string; + columns: Policy[]; + }[]; + + const policyDeletes = [] as { + table: string; + schema: string; + columns: Policy[]; + }[]; + + for (let entry of Object.values(policyRes)) { + const { renamed, created, deleted } = await policyResolver({ + tableName: entry.name, + schema: entry.schema, + deleted: entry.columns.deleted, + created: entry.columns.added, + }); + + if (created.length > 0) { + policyCreates.push({ + table: entry.name, + schema: entry.schema, + columns: created, + }); + } + + if (deleted.length > 0) { + policyDeletes.push({ + table: entry.name, + schema: entry.schema, + columns: deleted, + }); + } + + if (renamed.length > 0) { + policyRenames.push({ + table: entry.name, + schema: entry.schema, + renames: renamed, + }); + } + } + + const policyRenamesDict = columnRenames.reduce( + (acc, it) => { + acc[`${it.schema || 'public'}.${it.table}`] = it.renames; + return acc; + }, + {} as Record< + string, + { + from: Named; + to: Named; + }[] + >, + ); + + const policyPatchedSnap1 = copy(tablesPatchedSnap1); + policyPatchedSnap1.tables = mapEntries( + policyPatchedSnap1.tables, + (tableKey, tableValue) => { + const patchedPolicies = mapKeys( + tableValue.policies, + (policyKey, policy) => { + const rens = policyRenamesDict[ + `${tableValue.schema || 'public'}.${tableValue.name}` + ] || []; + + const newName = columnChangeFor(policyKey, rens); + const unsquashedPolicy = PgSquasher.unsquashPolicy(policy); + unsquashedPolicy.name = newName; + policy = PgSquasher.squashPolicy(unsquashedPolicy); + return newName; + }, + ); + + tableValue.policies = patchedPolicies; + return [tableKey, tableValue]; + }, + ); + + //// + const diffResult = applyJsonDiff(columnsPatchedSnap1, json2); // no diffs @@ -916,7 +1038,98 @@ export const applyPgSnapshotsDiff = async ( }) .flat(); + const jsonCreatePoliciesStatements: JsonCreatePolicyStatement[] = []; + const jsonDropPoliciesStatements: JsonDropPolicyStatement[] = []; + const jsonAlterPoliciesStatements: JsonAlterPolicyStatement[] = []; + const jsonRenamePoliciesStatements: JsonRenamePolicyStatement[] = []; + + for (let it of policyRenames) { + jsonRenamePoliciesStatements.push( + ...prepareRenamePolicyJsons(it.table, it.schema, it.renames), + ); + } + alteredTables.forEach((it) => { + // handle policies + + jsonCreatePoliciesStatements.push( + ...prepareCreatePolicyJsons( + it.name, + it.schema, + it.addedPolicies || {}, + ), + ); + + jsonDropPoliciesStatements.push( + ...prepareDropPolicyJsons( + it.name, + it.schema, + it.deletedPolicies || {}, + ), + ); + + Object.keys(it.alteredPolicies).forEach((policyName: string) => { + const newPolicy = PgSquasher.unsquashPolicy(it.alteredPolicies[policyName].__new); + const oldPolicy = PgSquasher.unsquashPolicy(it.alteredPolicies[policyName].__old); + + if (newPolicy.as !== oldPolicy.as) { + jsonDropPoliciesStatements.push( + ...prepareDropPolicyJsons( + it.name, + it.schema, + { [oldPolicy.name]: it.alteredPolicies[policyName].__old }, + ), + ); + + jsonCreatePoliciesStatements.push( + ...prepareCreatePolicyJsons( + it.name, + it.schema, + { [newPolicy.name]: it.alteredPolicies[policyName].__new }, + ), + ); + } + + if (newPolicy.for !== oldPolicy.for) { + jsonDropPoliciesStatements.push( + ...prepareDropPolicyJsons( + it.name, + it.schema, + { [oldPolicy.name]: it.alteredPolicies[policyName].__old }, + ), + ); + + jsonCreatePoliciesStatements.push( + ...prepareCreatePolicyJsons( + it.name, + it.schema, + { [newPolicy.name]: it.alteredPolicies[policyName].__new }, + ), + ); + } + + // alter + jsonAlterPoliciesStatements.push( + prepareAlterPolicyJson( + it.name, + it.schema, + it.alteredPolicies[policyName].__old, + it.alteredPolicies[policyName].__new, + ), + ); + }); + + // TODO + // Add to sql generators + // Test generate logic manually + // Add tests + // Add introspect + // add introspect-pg.ts - think about introspecting from supabase and neon + // add push logic - think about using with neon and supabase + // add push and introspect tests + // beta release + + // handle indexes const droppedIndexes = Object.keys(it.alteredIndexes).reduce( (current, item: string) => { current[item] = it.alteredIndexes[item].__old; @@ -1133,6 +1346,11 @@ export const applyPgSnapshotsDiff = async ( jsonStatements.push(...jsonAlteredUniqueConstraints); + jsonStatements.push(...jsonCreatePoliciesStatements); + jsonStatements.push(...jsonRenamePoliciesStatements); + jsonStatements.push(...jsonAlterPoliciesStatements); + jsonStatements.push(...jsonDropPoliciesStatements); + jsonStatements.push(...dropEnums); jsonStatements.push(...dropSequences); jsonStatements.push(...dropSchemas); diff --git a/drizzle-kit/src/sqlgenerator.ts b/drizzle-kit/src/sqlgenerator.ts index ec1a2d69e..5d6e1e002 100644 --- a/drizzle-kit/src/sqlgenerator.ts +++ b/drizzle-kit/src/sqlgenerator.ts @@ -20,6 +20,7 @@ import { JsonAlterColumnSetPrimaryKeyStatement, JsonAlterColumnTypeStatement, JsonAlterCompositePK, + JsonAlterPolicyStatement, JsonAlterReferenceStatement, JsonAlterSequenceStatement, JsonAlterTableRemoveFromSchema, @@ -29,6 +30,7 @@ import { JsonCreateCompositePK, JsonCreateEnumStatement, JsonCreateIndexStatement, + JsonCreatePolicyStatement, JsonCreateReferenceStatement, JsonCreateSchema, JsonCreateSequenceStatement, @@ -39,11 +41,13 @@ import { JsonDeleteUniqueConstraint, JsonDropColumnStatement, JsonDropIndexStatement, + JsonDropPolicyStatement, JsonDropSequenceStatement, JsonDropTableStatement, JsonMoveSequenceStatement, JsonPgCreateIndexStatement, JsonRenameColumnStatement, + JsonRenamePolicyStatement, JsonRenameSchema, JsonRenameSequenceStatement, JsonRenameTableStatement, @@ -131,6 +135,81 @@ abstract class Convertor { abstract convert(statement: JsonStatement): string | string[]; } +class PgCreatePolicyConvertor extends Convertor { + override can(statement: JsonStatement, dialect: Dialect): boolean { + return statement.type === 'create_policy' && dialect === 'postgresql'; + } + override convert(statement: JsonCreatePolicyStatement): string | string[] { + const policy = PgSquasher.unsquashPolicy(statement.data); + + const tableNameWithSchema = statement.schema + ? `"${statement.schema}"."${statement.tableName}"` + : `"${statement.tableName}"`; + + const usingPart = policy.using ? ` USING (${policy.using})` : ''; + + const withCheckPart = policy.withCheck ? ` WITH CHECK (${policy.withCheck})` : ''; + + return `CREATE POLICY ${policy.name} ON ${tableNameWithSchema} AS ${policy.as?.toUpperCase()} FOR ${policy.for?.toUpperCase()} TO ${policy.to?.toUpperCase()}${usingPart}${withCheckPart}`; + } +} + +class PgDropPolicyConvertor extends Convertor { + override can(statement: JsonStatement, dialect: Dialect): boolean { + return statement.type === 'drop_policy' && dialect === 'postgresql'; + } + override convert(statement: JsonDropPolicyStatement): string | string[] { + const policy = PgSquasher.unsquashPolicy(statement.data); + + const tableNameWithSchema = statement.schema + ? `"${statement.schema}"."${statement.tableName}"` + : `"${statement.tableName}"`; + + return `DROP POLICY ${policy.name} ON ${tableNameWithSchema} CASCADE`; + } +} + +class PgRenamePolicyConvertor extends Convertor { + override can(statement: JsonStatement, dialect: Dialect): boolean { + return statement.type === 'rename_policy' && dialect === 'postgresql'; + } + override convert(statement: JsonRenamePolicyStatement): string | string[] { + const tableNameWithSchema = statement.schema + ? `"${statement.schema}"."${statement.tableName}"` + : `"${statement.tableName}"`; + + return `ALTER POLICY ${statement.oldName} ON ${tableNameWithSchema} RENAME TO ${statement.newName}`; + } +} + +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); + + const tableNameWithSchema = statement.schema + ? `"${statement.schema}"."${statement.tableName}"` + : `"${statement.tableName}"`; + + const usingPart = newPolicy.using + ? ` USING (${newPolicy.using})` + : oldPolicy.using + ? ` USING (${oldPolicy.using})` + : ''; + + const withCheckPart = newPolicy.withCheck + ? ` WITH CHECK (${newPolicy.withCheck})` + : oldPolicy.withCheck + ? ` WITH CHECK (${oldPolicy.withCheck})` + : ''; + + return `ALTER POLICY ${oldPolicy.name} ON ${tableNameWithSchema} TO ${newPolicy.to}${usingPart}${withCheckPart}`; + } +} + class PgCreateTableConvertor extends Convertor { can(statement: JsonStatement, dialect: Dialect): boolean { return statement.type === 'create_table' && dialect === 'postgresql'; @@ -2568,6 +2647,11 @@ convertors.push(new PgAlterTableAlterColumnDropNotNullConvertor()); convertors.push(new PgAlterTableAlterColumnSetDefaultConvertor()); convertors.push(new PgAlterTableAlterColumnDropDefaultConvertor()); +convertors.push(new PgAlterPolicyConvertor()); +convertors.push(new PgCreatePolicyConvertor()); +convertors.push(new PgDropPolicyConvertor()); +convertors.push(new PgRenamePolicyConvertor()); + /// generated convertors.push(new PgAlterTableAlterColumnSetExpressionConvertor()); convertors.push(new PgAlterTableAlterColumnDropGeneratedConvertor()); diff --git a/drizzle-kit/tests/rls/pg-policy.test.ts b/drizzle-kit/tests/rls/pg-policy.test.ts new file mode 100644 index 000000000..41289f4e4 --- /dev/null +++ b/drizzle-kit/tests/rls/pg-policy.test.ts @@ -0,0 +1,23 @@ +import { integer, pgPolicy, pgTable } from 'drizzle-orm/pg-core'; +import { diffTestSchemas } from 'tests/schemaDiffer'; +import { expect, test } from 'vitest'; + +test('add policy #1', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + console.log(sqlStatements); +}); diff --git a/drizzle-kit/tests/schemaDiffer.ts b/drizzle-kit/tests/schemaDiffer.ts index 4a14d920b..b9d7b5fde 100644 --- a/drizzle-kit/tests/schemaDiffer.ts +++ b/drizzle-kit/tests/schemaDiffer.ts @@ -10,6 +10,7 @@ import { columnsResolver, enumsResolver, Named, + policyResolver, schemasResolver, sequencesResolver, tablesResolver, @@ -23,7 +24,7 @@ import { mysqlSchema, squashMysqlScheme } from 'src/serializer/mysqlSchema'; import { generateMySqlSnapshot } from 'src/serializer/mysqlSerializer'; import { fromDatabase as fromMySqlDatabase } from 'src/serializer/mysqlSerializer'; import { prepareFromPgImports } from 'src/serializer/pgImports'; -import { pgSchema, squashPgScheme } from 'src/serializer/pgSchema'; +import { pgSchema, Policy, squashPgScheme } from 'src/serializer/pgSchema'; import { fromDatabase, generatePgSnapshot } from 'src/serializer/pgSerializer'; import { prepareFromSqliteImports } from 'src/serializer/sqliteImports'; import { sqliteSchema, squashSqliteScheme } from 'src/serializer/sqliteSchema'; @@ -338,6 +339,70 @@ async ( } }; +export const testPolicyResolver = (renames: Set) => +async ( + input: ColumnsResolverInput, +): Promise> => { + try { + if ( + input.created.length === 0 + || input.deleted.length === 0 + || renames.size === 0 + ) { + return { + tableName: input.tableName, + schema: input.schema, + created: input.created, + renamed: [], + deleted: input.deleted, + }; + } + + let createdPolicies = [...input.created]; + let deletedPolicies = [...input.deleted]; + + const renamed: { from: Policy; to: Policy }[] = []; + + const schema = input.schema || 'public'; + + for (let rename of renames) { + const [from, to] = rename.split('->'); + + const idxFrom = deletedPolicies.findIndex((it) => { + return `${schema}.${input.tableName}.${it.name}` === from; + }); + + if (idxFrom >= 0) { + const idxTo = createdPolicies.findIndex((it) => { + return `${schema}.${input.tableName}.${it.name}` === to; + }); + + renamed.push({ + from: deletedPolicies[idxFrom], + to: createdPolicies[idxTo], + }); + + delete createdPolicies[idxTo]; + delete deletedPolicies[idxFrom]; + + createdPolicies = createdPolicies.filter(Boolean); + deletedPolicies = deletedPolicies.filter(Boolean); + } + } + + return { + tableName: input.tableName, + schema: input.schema, + created: createdPolicies, + deleted: deletedPolicies, + renamed, + }; + } catch (e) { + console.error(e); + throw e; + } +}; + export const testColumnsResolver = (renames: Set) => async ( input: ColumnsResolverInput, @@ -476,6 +541,7 @@ export const diffTestSchemasPush = async ( testSchemasResolver(renames), testEnumsResolver(renames), testSequencesResolver(renames), + testPolicyResolver(renames), testTablesResolver(renames), testColumnsResolver(renames), validatedPrev, @@ -490,6 +556,7 @@ export const diffTestSchemasPush = async ( schemasResolver, enumsResolver, sequencesResolver, + policyResolver, tablesResolver, columnsResolver, validatedPrev, @@ -548,6 +615,7 @@ export const applyPgDiffs = async (sn: PostgresSchema) => { testSchemasResolver(new Set()), testEnumsResolver(new Set()), testSequencesResolver(new Set()), + testPolicyResolver(new Set()), testTablesResolver(new Set()), testColumnsResolver(new Set()), validatedPrev, @@ -625,6 +693,7 @@ export const diffTestSchemas = async ( testSchemasResolver(renames), testEnumsResolver(renames), testSequencesResolver(renames), + testPolicyResolver(renames), testTablesResolver(renames), testColumnsResolver(renames), validatedPrev, @@ -638,6 +707,7 @@ export const diffTestSchemas = async ( schemasResolver, enumsResolver, sequencesResolver, + policyResolver, tablesResolver, columnsResolver, validatedPrev, @@ -1127,6 +1197,7 @@ export const introspectPgToFile = async ( testSchemasResolver(new Set()), testEnumsResolver(new Set()), testSequencesResolver(new Set()), + testPolicyResolver(new Set()), testTablesResolver(new Set()), testColumnsResolver(new Set()), validatedCurAfterImport, diff --git a/drizzle-kit/vitest.config.ts b/drizzle-kit/vitest.config.ts index 602e96ede..23067c3b0 100644 --- a/drizzle-kit/vitest.config.ts +++ b/drizzle-kit/vitest.config.ts @@ -4,7 +4,8 @@ import { defineConfig } from 'vitest/config'; export default defineConfig({ test: { include: [ - 'tests/**/*.test.ts', + // 'tests/**/*.test.ts', + 'tests/rls/pg-policy.test.ts', ], typecheck: { From 474710ed3a6d090f6fe50c2454451f496bb0f90c Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Fri, 30 Aug 2024 13:46:32 +0300 Subject: [PATCH 03/24] Add all policies api and pgRole into generate and pull --- drizzle-kit/src/api.ts | 1 + drizzle-kit/src/cli/commands/introspect.ts | 17 +- drizzle-kit/src/cli/commands/pgIntrospect.ts | 4 +- drizzle-kit/src/cli/commands/push.ts | 4 +- drizzle-kit/src/cli/commands/utils.ts | 6 +- drizzle-kit/src/cli/schema.ts | 2 + drizzle-kit/src/cli/validations/cli.ts | 11 +- drizzle-kit/src/index.ts | 3 + drizzle-kit/src/introspect-pg.ts | 40 ++ drizzle-kit/src/jsonStatements.ts | 38 +- drizzle-kit/src/serializer/index.ts | 4 +- drizzle-kit/src/serializer/pgImports.ts | 13 +- drizzle-kit/src/serializer/pgSchema.ts | 15 +- drizzle-kit/src/serializer/pgSerializer.ts | 172 ++++- drizzle-kit/src/snapshotsDiffer.ts | 75 ++- drizzle-kit/src/sqlgenerator.ts | 82 ++- drizzle-kit/tests/mysql.test.ts | 2 + drizzle-kit/tests/pg-identity.test.ts | 3 + drizzle-kit/tests/pg-tables.test.ts | 13 + drizzle-kit/tests/push/pg.test.ts | 4 + drizzle-kit/tests/rls/pg-policy.test.ts | 658 ++++++++++++++++++- drizzle-kit/tests/rls/pg-role.test.ts | 0 drizzle-kit/tests/schemaDiffer.ts | 26 +- drizzle-kit/tests/sqlite-tables.test.ts | 1 + drizzle-orm/src/pg-core/index.ts | 1 + drizzle-orm/src/pg-core/policies.ts | 12 +- drizzle-orm/src/pg-core/roles.ts | 41 ++ 27 files changed, 1183 insertions(+), 65 deletions(-) create mode 100644 drizzle-kit/tests/rls/pg-role.test.ts create mode 100644 drizzle-orm/src/pg-core/roles.ts diff --git a/drizzle-kit/src/api.ts b/drizzle-kit/src/api.ts index d0cceda52..8922fef0e 100644 --- a/drizzle-kit/src/api.ts +++ b/drizzle-kit/src/api.ts @@ -44,6 +44,7 @@ export const generateDrizzleJson = ( prepared.enums, prepared.schemas, prepared.sequences, + prepared.roles, schemaFilters, ); diff --git a/drizzle-kit/src/cli/commands/introspect.ts b/drizzle-kit/src/cli/commands/introspect.ts index 61ba0b44a..89ced7e67 100644 --- a/drizzle-kit/src/cli/commands/introspect.ts +++ b/drizzle-kit/src/cli/commands/introspect.ts @@ -16,6 +16,7 @@ import { drySQLite, type SQLiteSchema, squashSqliteScheme } from '../../serializ import { fromDatabase as fromSqliteDatabase } from '../../serializer/sqliteSerializer'; import { applyMysqlSnapshotsDiff, applyPgSnapshotsDiff, applySqliteSnapshotsDiff } from '../../snapshotsDiffer'; import { prepareOutFolder } from '../../utils'; +import { Entities } from '../validations/cli'; import type { Casing, Prefix } from '../validations/common'; import type { MysqlCredentials } from '../validations/mysql'; import type { PostgresCredentials } from '../validations/postgres'; @@ -24,6 +25,7 @@ import { IntrospectProgress } from '../views'; import { columnsResolver, enumsResolver, + policyResolver, schemasResolver, sequencesResolver, tablesResolver, @@ -38,6 +40,7 @@ export const introspectPostgres = async ( tablesFilter: string[], schemasFilter: string[], prefix: Prefix, + entities: Entities, ) => { const { preparePostgresDB } = await import('../connections'); const db = await preparePostgresDB(credentials); @@ -70,11 +73,18 @@ export const introspectPostgres = async ( }; const progress = new IntrospectProgress(true); + const res = await renderWithTask( progress, - fromPostgresDatabase(db, filter, schemasFilter, (stage, count, status) => { - progress.update(stage, count, status); - }), + fromPostgresDatabase( + db, + filter, + schemasFilter, + entities, + (stage, count, status) => { + progress.update(stage, count, status); + }, + ), ); const schema = { id: originUUID, prevId: '', ...res } as PgSchema; @@ -97,6 +107,7 @@ export const introspectPostgres = async ( schemasResolver, enumsResolver, sequencesResolver, + policyResolver, tablesResolver, columnsResolver, dryPg, diff --git a/drizzle-kit/src/cli/commands/pgIntrospect.ts b/drizzle-kit/src/cli/commands/pgIntrospect.ts index dbd3ba238..020ed031b 100644 --- a/drizzle-kit/src/cli/commands/pgIntrospect.ts +++ b/drizzle-kit/src/cli/commands/pgIntrospect.ts @@ -4,12 +4,14 @@ import { originUUID } from '../../global'; import type { PgSchema } from '../../serializer/pgSchema'; import { fromDatabase } from '../../serializer/pgSerializer'; import type { DB } from '../../utils'; +import { Entities } from '../validations/cli'; import { ProgressView } from '../views'; export const pgPushIntrospect = async ( db: DB, filters: string[], schemaFilters: string[], + entities: Entities = { roles: true }, ) => { const matchers = filters.map((it) => { return new Minimatch(it); @@ -43,7 +45,7 @@ export const pgPushIntrospect = async ( ); const res = await renderWithTask( progress, - fromDatabase(db, filter, schemaFilters), + fromDatabase(db, filter, schemaFilters, entities), ); const schema = { id: originUUID, prevId: '', ...res } as PgSchema; diff --git a/drizzle-kit/src/cli/commands/push.ts b/drizzle-kit/src/cli/commands/push.ts index e48a5da9e..84d3bce7f 100644 --- a/drizzle-kit/src/cli/commands/push.ts +++ b/drizzle-kit/src/cli/commands/push.ts @@ -2,6 +2,7 @@ import chalk from 'chalk'; import { render } from 'hanji'; import { fromJson } from '../../sqlgenerator'; import { Select } from '../selector-ui'; +import { Entities } from '../validations/cli'; import type { MysqlCredentials } from '../validations/mysql'; import { withStyle } from '../validations/outputs'; import type { PostgresCredentials } from '../validations/postgres'; @@ -157,13 +158,14 @@ export const pgPush = async ( credentials: PostgresCredentials, tablesFilter: string[], schemasFilter: string[], + entities: Entities, force: boolean, ) => { const { preparePostgresDB } = await import('../connections'); const { pgPushIntrospect } = await import('./pgIntrospect'); const db = await preparePostgresDB(credentials); - const { schema } = await pgPushIntrospect(db, tablesFilter, schemasFilter); + const { schema } = await pgPushIntrospect(db, tablesFilter, schemasFilter, entities); const { preparePgPush } = await import('./migrate'); diff --git a/drizzle-kit/src/cli/commands/utils.ts b/drizzle-kit/src/cli/commands/utils.ts index fbfeede70..8ea33883f 100644 --- a/drizzle-kit/src/cli/commands/utils.ts +++ b/drizzle-kit/src/cli/commands/utils.ts @@ -6,7 +6,7 @@ import { object, string } from 'zod'; import { assertUnreachable } from '../../global'; import { type Dialect, dialect } from '../../schemaValidator'; import { prepareFilenames } from '../../serializer'; -import { pullParams, pushParams } from '../validations/cli'; +import { Entities, pullParams, pushParams } from '../validations/cli'; import { Casing, CliConfig, @@ -354,6 +354,7 @@ export const preparePullConfig = async ( tablesFilter: string[]; schemasFilter: string[]; prefix: Prefix; + entities: Entities; } > => { const raw = flattenPull( @@ -413,6 +414,7 @@ export const preparePullConfig = async ( tablesFilter, schemasFilter, prefix: config.migrations?.prefix || 'index', + entities: config.entities, }; } @@ -431,6 +433,7 @@ export const preparePullConfig = async ( tablesFilter, schemasFilter, prefix: config.migrations?.prefix || 'index', + entities: config.entities, }; } @@ -449,6 +452,7 @@ export const preparePullConfig = async ( tablesFilter, schemasFilter, prefix: config.migrations?.prefix || 'index', + entities: config.entities, }; } diff --git a/drizzle-kit/src/cli/schema.ts b/drizzle-kit/src/cli/schema.ts index 4da8af0ac..44eba5c8d 100644 --- a/drizzle-kit/src/cli/schema.ts +++ b/drizzle-kit/src/cli/schema.ts @@ -416,6 +416,7 @@ export const pull = command({ tablesFilter, schemasFilter, prefix, + entities, } = config; mkdirSync(out, { recursive: true }); @@ -462,6 +463,7 @@ export const pull = command({ tablesFilter, schemasFilter, prefix, + entities, ); } else if (dialect === 'mysql') { const { introspectMysql } = await import('./commands/introspect'); diff --git a/drizzle-kit/src/cli/validations/cli.ts b/drizzle-kit/src/cli/validations/cli.ts index c4bbbe530..47105c09c 100644 --- a/drizzle-kit/src/cli/validations/cli.ts +++ b/drizzle-kit/src/cli/validations/cli.ts @@ -1,4 +1,4 @@ -import { boolean, intersection, literal, object, string, TypeOf, union } from 'zod'; +import { array, boolean, intersection, literal, object, string, TypeOf, union } from 'zod'; import { dialect } from '../../schemaValidator'; import { casing, prefix } from './common'; @@ -43,8 +43,17 @@ export const pullParams = object({ migrations: object({ prefix: prefix.optional().default('index'), }).optional(), + entities: object({ + roles: boolean().or(object({ + provider: string().optional(), + include: string().array().optional(), + exclude: string().array().optional(), + })).optional().default(false), + }).optional(), }).passthrough(); +export type Entities = TypeOf['entities']; + export type PullParams = TypeOf; export const configCheck = object({ diff --git a/drizzle-kit/src/index.ts b/drizzle-kit/src/index.ts index 3d29b5c85..9723d655c 100644 --- a/drizzle-kit/src/index.ts +++ b/drizzle-kit/src/index.ts @@ -125,6 +125,9 @@ export type Config = introspect?: { casing: 'camel' | 'preserve'; }; + entities?: { + roles?: boolean | { provider?: string; exclude?: string[]; include?: string[] }[]; + }; } & ( | { diff --git a/drizzle-kit/src/introspect-pg.ts b/drizzle-kit/src/introspect-pg.ts index b7a52b735..d9c909bb5 100644 --- a/drizzle-kit/src/introspect-pg.ts +++ b/drizzle-kit/src/introspect-pg.ts @@ -19,6 +19,7 @@ import { Index, PgKitInternals, PgSchemaInternal, + Policy, PrimaryKey, UniqueConstraint, } from './serializer/pgSchema'; @@ -464,6 +465,19 @@ export const schemaToTypeScript = ( }) .join(''); + const rolesStatements = Object.entries(schema.roles) + .map((it) => { + const fields = it[1]; + return `export const ${withCasing(it[0], casing)} = pgRole("${fields.name}"${ + !fields.createDb && !fields.createRole && !fields.inherit + ? '' + : `, { ${fields.createDb ? `createDb: true, ` : ''}${fields.createRole ? `createRole: true, ` : ''}${ + fields.inherit ? `inherit: true, ` : '' + } }` + });\n`; + }) + .join(''); + const tableStatements = Object.values(schema.tables).map((table) => { const tableSchema = schemas[table.schema]; const paramName = paramNameFor(table.name, tableSchema); @@ -510,6 +524,10 @@ export const schemaToTypeScript = ( Object.values(table.uniqueConstraints), casing, ); + statement += createTablePolicies( + Object.values(table.policies), + casing, + ); statement += '\t}\n'; statement += '}'; } @@ -528,6 +546,7 @@ export const schemaToTypeScript = ( import { sql } from "drizzle-orm"\n\n`; let decalrations = schemaStatements; + decalrations += rolesStatements; decalrations += enumStatements; decalrations += sequencesStatements; decalrations += '\n'; @@ -1248,6 +1267,27 @@ const createTablePKs = (pks: PrimaryKey[], casing: Casing): string => { return statement; }; +const createTablePolicies = ( + policies: Policy[], + casing: Casing, +): string => { + let statement = ''; + + policies.forEach((it) => { + const idxKey = withCasing(it.name, casing); + + statement += `\t\t${idxKey}: `; + statement += 'pgPolicy('; + statement += `"${it.name}", { `; + statement += `as: "${it.as}", for: "${it.for}", to: "${it.to?.join(',')}"${ + it.using ? `, using: sql\`${it.using}\`` : '' + }${it.withCheck ? `, withCheck: sql\`${it.withCheck}\`` : ''}`; + statement += ` },\n`; + }); + + return statement; +}; + const createTableUniques = ( unqs: UniqueConstraint[], casing: Casing, diff --git a/drizzle-kit/src/jsonStatements.ts b/drizzle-kit/src/jsonStatements.ts index 17dc76456..d98cde0e6 100644 --- a/drizzle-kit/src/jsonStatements.ts +++ b/drizzle-kit/src/jsonStatements.ts @@ -32,6 +32,7 @@ export interface JsonCreateTableStatement { compositePKs: string[]; compositePkName?: string; uniqueConstraints?: string[]; + policies?: string[]; internals?: MySqlKitInternals; } @@ -39,6 +40,7 @@ export interface JsonDropTableStatement { type: 'drop_table'; tableName: string; schema: string; + policies?: string[]; } export interface JsonRenameTableStatement { @@ -156,14 +158,14 @@ export interface JsonSqliteAddColumnStatement { export interface JsonCreatePolicyStatement { type: 'create_policy'; tableName: string; - data: string; + data: Policy; schema: string; } export interface JsonDropPolicyStatement { type: 'drop_policy'; tableName: string; - data: string; + data: Policy; schema: string; } @@ -175,6 +177,18 @@ export interface JsonRenamePolicyStatement { schema: string; } +export interface JsonEnableRLSStatement { + type: 'enable_rls'; + tableName: string; + schema: string; +} + +export interface JsonDisableRLSStatement { + type: 'disable_rls'; + tableName: string; + schema: string; +} + export interface JsonAlterPolicyStatement { type: 'alter_policy'; tableName: string; @@ -589,14 +603,16 @@ export type JsonStatement = | JsonDropPolicyStatement | JsonCreatePolicyStatement | JsonAlterPolicyStatement - | JsonRenamePolicyStatement; + | JsonRenamePolicyStatement + | JsonEnableRLSStatement + | JsonDisableRLSStatement; export const preparePgCreateTableJson = ( table: Table, // TODO: remove? json2: PgSchema, ): JsonCreateTableStatement => { - const { name, schema, columns, compositePrimaryKeys, uniqueConstraints } = table; + const { name, schema, columns, compositePrimaryKeys, uniqueConstraints, policies } = table; const tableKey = `${schema || 'public'}.${name}`; // TODO: @AndriiSherman. We need this, will add test cases @@ -614,6 +630,7 @@ export const preparePgCreateTableJson = ( compositePKs: Object.values(compositePrimaryKeys), compositePkName: compositePkName, uniqueConstraints: Object.values(uniqueConstraints), + policies: Object.values(policies), }; }; @@ -678,6 +695,7 @@ export const prepareDropTableJson = (table: Table): JsonDropTableStatement => { type: 'drop_table', tableName: table.name, schema: table.schema, + policies: table.policies ? Object.values(table.policies) : [], }; }; @@ -1939,13 +1957,13 @@ export const prepareRenamePolicyJsons = ( export const prepareCreatePolicyJsons = ( tableName: string, schema: string, - policies: Record, + policies: Policy[], ): JsonCreatePolicyStatement[] => { - return Object.values(policies).map((policyData) => { + return policies.map((it) => { return { type: 'create_policy', tableName, - data: policyData, + data: it, schema, }; }); @@ -1954,13 +1972,13 @@ export const prepareCreatePolicyJsons = ( export const prepareDropPolicyJsons = ( tableName: string, schema: string, - policies: Record, + policies: Policy[], ): JsonDropPolicyStatement[] => { - return Object.values(policies).map((policyData) => { + return policies.map((it) => { return { type: 'drop_policy', tableName, - data: policyData, + data: it, schema, }; }); diff --git a/drizzle-kit/src/serializer/index.ts b/drizzle-kit/src/serializer/index.ts index 214ca38c7..beacd3d5c 100644 --- a/drizzle-kit/src/serializer/index.ts +++ b/drizzle-kit/src/serializer/index.ts @@ -60,11 +60,11 @@ export const serializePg = async ( const { prepareFromPgImports } = await import('./pgImports'); const { generatePgSnapshot } = await import('./pgSerializer'); - const { tables, enums, schemas, sequences } = await prepareFromPgImports( + const { tables, enums, schemas, sequences, roles } = await prepareFromPgImports( filenames, ); - return generatePgSnapshot(tables, enums, schemas, sequences, schemaFilter); + return generatePgSnapshot(tables, enums, schemas, sequences, roles, schemaFilter); }; export const serializeSQLite = async ( diff --git a/drizzle-kit/src/serializer/pgImports.ts b/drizzle-kit/src/serializer/pgImports.ts index ffedd084c..c954504b2 100644 --- a/drizzle-kit/src/serializer/pgImports.ts +++ b/drizzle-kit/src/serializer/pgImports.ts @@ -1,5 +1,5 @@ import { is } from 'drizzle-orm'; -import { AnyPgTable, isPgEnum, isPgSequence, PgEnum, PgSchema, PgSequence, PgTable } from 'drizzle-orm/pg-core'; +import { AnyPgTable, isPgEnum, isPgSequence, PgEnum, PgRole, PgSchema, PgSequence, PgTable } from 'drizzle-orm/pg-core'; import { safeRegister } from '../cli/commands/utils'; export const prepareFromExports = (exports: Record) => { @@ -7,6 +7,7 @@ export const prepareFromExports = (exports: Record) => { const enums: PgEnum[] = []; const schemas: PgSchema[] = []; const sequences: PgSequence[] = []; + const roles: PgRole[] = []; const i0values = Object.values(exports); i0values.forEach((t) => { @@ -25,9 +26,13 @@ export const prepareFromExports = (exports: Record) => { if (isPgSequence(t)) { sequences.push(t); } + + if (is(t, PgRole)) { + roles.push(t); + } }); - return { tables, enums, schemas, sequences }; + return { tables, enums, schemas, sequences, roles }; }; export const prepareFromPgImports = async (imports: string[]) => { @@ -35,6 +40,7 @@ export const prepareFromPgImports = async (imports: string[]) => { let enums: PgEnum[] = []; let schemas: PgSchema[] = []; let sequences: PgSequence[] = []; + let roles: PgRole[] = []; const { unregister } = await safeRegister(); for (let i = 0; i < imports.length; i++) { @@ -47,8 +53,9 @@ export const prepareFromPgImports = async (imports: string[]) => { enums.push(...prepared.enums); schemas.push(...prepared.schemas); sequences.push(...prepared.sequences); + roles.push(...prepared.roles); } unregister(); - return { tables: Array.from(new Set(tables)), enums, schemas, sequences }; + return { tables: Array.from(new Set(tables)), enums, schemas, sequences, roles }; }; diff --git a/drizzle-kit/src/serializer/pgSchema.ts b/drizzle-kit/src/serializer/pgSchema.ts index e03354ebe..b579ebdab 100644 --- a/drizzle-kit/src/serializer/pgSchema.ts +++ b/drizzle-kit/src/serializer/pgSchema.ts @@ -148,6 +148,13 @@ export const sequenceSchema = object({ schema: string(), }).strict(); +export const roleSchema = object({ + name: string(), + createDb: boolean().optional(), + createRole: boolean().optional(), + inherit: boolean().optional(), +}).strict(); + export const sequenceSquashed = object({ name: string(), schema: string(), @@ -224,7 +231,7 @@ const policy = object({ name: string(), as: enumType(['permissive', 'restrictive']).optional(), for: enumType(['all', 'select', 'insert', 'update', 'delete']).optional(), - to: string().optional(), + to: string().array().optional(), using: string().optional(), withCheck: string().optional(), }).strict(); @@ -379,6 +386,7 @@ export const pgSchemaInternal = object({ enums: record(string(), enumSchema), schemas: record(string(), string()), sequences: record(string(), sequenceSchema).default({}), + roles: record(string(), roleSchema).default({}), _meta: object({ schemas: record(string(), string()), tables: record(string(), string()), @@ -440,6 +448,7 @@ export const pgSchema = pgSchemaInternal.merge(schemaHash); export type Enum = TypeOf; export type Sequence = TypeOf; +export type Role = TypeOf; export type Column = TypeOf; export type TableV3 = TypeOf; export type TableV4 = TypeOf; @@ -562,7 +571,7 @@ export const PgSquasher = { };${fk.onDelete ?? ''};${fk.schemaTo || 'public'}`; }, squashPolicy: (policy: Policy) => { - return `${policy.name}--${policy.as}--${policy.for}--${policy.to}--${policy.using}--${policy.withCheck}`; + return `${policy.name}--${policy.as}--${policy.for}--${policy.to?.join(',')}--${policy.using}--${policy.withCheck}`; }, unsquashPolicy: (policy: string): Policy => { const splitted = policy.split('--'); @@ -570,7 +579,7 @@ export const PgSquasher = { name: splitted[0], as: splitted[1] as Policy['as'], for: splitted[2] as Policy['for'], - to: splitted[3], + to: splitted[3].split(','), using: splitted[4] !== 'undefined' ? splitted[4] : undefined, withCheck: splitted[5] !== 'undefined' ? splitted[5] : undefined, }; diff --git a/drizzle-kit/src/serializer/pgSerializer.ts b/drizzle-kit/src/serializer/pgSerializer.ts index 91076d78f..f9487684d 100644 --- a/drizzle-kit/src/serializer/pgSerializer.ts +++ b/drizzle-kit/src/serializer/pgSerializer.ts @@ -9,6 +9,8 @@ import { PgEnum, PgEnumColumn, PgInteger, + PgPolicy, + PgRole, PgSchema, PgSequence, uniqueKeyName, @@ -27,6 +29,7 @@ import type { PgSchemaInternal, Policy, PrimaryKey, + Role, Sequence, Table, UniqueConstraint, @@ -58,7 +61,7 @@ function maxRangeForIdentityBasedOn(columnType: string) { : '32767'; } -function minRangeForIdentityBasedOn(columnType: string) { +export function minRangeForIdentityBasedOn(columnType: string) { return columnType === 'integer' ? '-2147483648' : columnType === 'bitint' @@ -76,7 +79,7 @@ function stringFromDatabaseIdentityProperty(field: any): string | undefined { : String(field); } -function buildArrayString(array: any[], sqlType: string): string { +export function buildArrayString(array: any[], sqlType: string): string { sqlType = sqlType.split('[')[0]; const values = array .map((value) => { @@ -118,10 +121,12 @@ export const generatePgSnapshot = ( enums: PgEnum[], schemas: PgSchema[], sequences: PgSequence[], + roles: PgRole[], schemaFilter?: string[], ): PgSchemaInternal => { const result: Record = {}; const sequencesToReturn: Record = {}; + const rolesToReturn: Record = {}; // This object stores unique names for indexes and will be used to detect if you have the same names for indexes // within the same PostgreSQL schema @@ -496,11 +501,31 @@ export const generatePgSnapshot = ( }); policies.forEach((policy) => { + const mappedTo = []; + + if (!policy.to) { + mappedTo.push('PUBLIC'); + } else { + if (policy.to && typeof policy.to === 'string') { + mappedTo.push(policy.to); + } else if (policy.to && is(policy.to, PgRole)) { + mappedTo.push(policy.to.name); + } else if (policy.to && Array.isArray(policy.to)) { + policy.to.forEach((it) => { + if (typeof it === 'string') { + mappedTo.push(it); + } else if (is(it, PgRole)) { + mappedTo.push(it.name); + } + }); + } + } + policiesObject[policy.name] = { name: policy.name, as: policy.as ?? 'permissive', for: policy.for ?? 'all', - to: policy.to ?? 'PUBLIC', + to: mappedTo, using: is(policy.using, SQL) ? sqlToStr(policy.using) : undefined, withCheck: is(policy.withCheck, SQL) ? sqlToStr(policy.withCheck) : undefined, }; @@ -550,6 +575,17 @@ export const generatePgSnapshot = ( } } + for (const role of roles) { + if (!(role as any)._existing) { + rolesToReturn[role.name] = { + name: role.name, + createDb: (role as any).createDb ?? false, + createRole: (role as any).createRole ?? false, + inherit: (role as any).inherit ?? false, + }; + } + } + const enumsToReturn: Record = enums.reduce<{ [key: string]: Enum; }>((map, obj) => { @@ -584,6 +620,7 @@ export const generatePgSnapshot = ( enums: enumsToReturn, schemas: schemasObject, sequences: sequencesToReturn, + roles: rolesToReturn, _meta: { schemas: {}, tables: {}, @@ -605,10 +642,59 @@ const trimChar = (str: string, char: string) => { : str.toString(); }; +function prepareRoles(entities?: { + roles: boolean | { + provider?: string | undefined; + include?: string[] | undefined; + exclude?: string[] | undefined; + }; +}) { + let useRoles: boolean = false; + const includeRoles: string[] = []; + const excludeRoles: string[] = []; + + if (entities && entities.roles) { + if (typeof entities.roles === 'object') { + if (entities.roles.provider) { + if (entities.roles.provider === 'supabase') { + excludeRoles.push(...[ + 'anon', + 'authenticator', + 'authenticated', + 'service_role', + 'supabase_auth_admin', + 'supabase_storage_admin', + 'dashboard_user', + 'supabase_admin', + ]); + } else if (entities.roles.provider === 'neon') { + excludeRoles.push(...['authenticated', 'anonymous']); + } + } + if (entities.roles.include) { + includeRoles.push(...entities.roles.include); + } + if (entities.roles.exclude) { + excludeRoles.push(...entities.roles.exclude); + } + } else { + useRoles = entities.roles; + } + } + return { useRoles, includeRoles, excludeRoles }; +} + export const fromDatabase = async ( db: DB, tablesFilter: (table: string) => boolean = () => true, schemaFilters: string[], + entities?: { + roles: boolean | { + provider?: string | undefined; + include?: string[] | undefined; + exclude?: string[] | undefined; + }; + }, progressCallback?: ( stage: IntrospectStage, count: number, @@ -719,6 +805,83 @@ export const fromDatabase = async ( progressCallback('enums', Object.keys(enumsToReturn).length, 'done'); } + const allRoles = await db.query< + { rolname: string; rolinherit: boolean; rolcreatedb: boolean; rolcreaterole: boolean } + >( + `SELECT rolname, rolinherit, rolcreatedb, rolcreaterole FROM pg_roles;`, + ); + + const rolesToReturn: Record = {}; + + const preparedRoles = prepareRoles(entities); + + if ( + preparedRoles.useRoles || !(preparedRoles.includeRoles.length === 0 && preparedRoles.excludeRoles.length === 0) + ) { + for (const dbRole of allRoles) { + if ( + preparedRoles.useRoles + ) { + rolesToReturn[dbRole.rolname] = { + createDb: dbRole.rolcreatedb, + createRole: dbRole.rolcreatedb, + inherit: dbRole.rolinherit, + name: dbRole.rolname, + }; + } else { + if (preparedRoles.includeRoles.length === 0 && preparedRoles.excludeRoles.length === 0) continue; + if ( + preparedRoles.includeRoles.includes(dbRole.rolname) && preparedRoles.excludeRoles.includes(dbRole.rolname) + ) continue; + if (preparedRoles.excludeRoles.includes(dbRole.rolname)) continue; + if (!preparedRoles.includeRoles.includes(dbRole.rolname)) continue; + + rolesToReturn[dbRole.rolname] = { + createDb: dbRole.rolcreatedb, + createRole: dbRole.rolcreatedb, + inherit: dbRole.rolinherit, + name: dbRole.rolname, + }; + } + } + } + + const wherePolicies = schemaFilters + .map((t) => `schemaname = '${t}'`) + .join(' or '); + + const policiesByTable: Record> = {}; + + const allPolicies = await db.query< + { + schemaname: string; + tablename: string; + name: string; + as: string; + to: string; + for: string; + using: string; + withCheck: string; + } + >(`SELECT schemaname, tablename, policyname as name, permissive as "as", roles as to, cmd as for, qual as using, "withCheck" FROM pg_policies${wherePolicies};`); + + for (const dbPolicy of allPolicies) { + const { tablename, schemaname, to, ...rest } = dbPolicy; + const tableForPolicy = policiesByTable[`${schemaname}.${tablename}`]; + + const parsedTo = to === '{}' + ? [] + : to.substring(1, to.length - 1).split(/\s*,\s*/g); + + if (tableForPolicy) { + tableForPolicy[dbPolicy.name] = { ...rest, to: parsedTo } as Policy; + } else { + policiesByTable[`${schemaname}.${tablename}`] = { + [dbPolicy.name]: { ...rest, to: parsedTo } as Policy, + }; + } + } + const sequencesInColumns: string[] = []; const all = allTables.map((row) => { @@ -1196,7 +1359,7 @@ export const fromDatabase = async ( foreignKeys: foreignKeysToReturn, compositePrimaryKeys: primaryKeys, uniqueConstraints: uniqueConstrains, - policies: {}, + policies: policiesByTable[`${tableSchema}.${tableName}`], }; } catch (e) { rej(e); @@ -1228,6 +1391,7 @@ export const fromDatabase = async ( enums: enumsToReturn, schemas: schemasObject, sequences: sequencesToReturn, + roles: rolesToReturn, _meta: { schemas: {}, tables: {}, diff --git a/drizzle-kit/src/snapshotsDiffer.ts b/drizzle-kit/src/snapshotsDiffer.ts index a333556fa..a4f5a3f03 100644 --- a/drizzle-kit/src/snapshotsDiffer.ts +++ b/drizzle-kit/src/snapshotsDiffer.ts @@ -13,7 +13,7 @@ import { union, ZodTypeAny, } from 'zod'; -import { applyJsonDiff, diffColumns, diffSchemasOrTables } from './jsonDiffer'; +import { applyJsonDiff, diffColumns, diffPolicies, diffSchemasOrTables } from './jsonDiffer'; import { fromJson } from './sqlgenerator'; import { @@ -31,8 +31,10 @@ import { JsonCreateUniqueConstraint, JsonDeleteCompositePK, JsonDeleteUniqueConstraint, + JsonDisableRLSStatement, JsonDropColumnStatement, JsonDropPolicyStatement, + JsonEnableRLSStatement, JsonReferenceStatement, JsonRenameColumnStatement, JsonRenamePolicyStatement, @@ -219,6 +221,7 @@ const tableScheme = object({ foreignKeys: record(string(), string()), compositePrimaryKeys: record(string(), string()).default({}), uniqueConstraints: record(string(), string()).default({}), + policies: record(string(), string()).default({}), }).strict(); export const alteredTableScheme = object({ @@ -746,7 +749,7 @@ export const applyPgSnapshotsDiff = async ( //// Policies - const policyRes = diffColumns(tablesPatchedSnap1.tables, json2.tables); + const policyRes = diffPolicies(tablesPatchedSnap1.tables, json2.tables); const policyRenames = [] as { table: string; @@ -770,8 +773,8 @@ export const applyPgSnapshotsDiff = async ( const { renamed, created, deleted } = await policyResolver({ tableName: entry.name, schema: entry.schema, - deleted: entry.columns.deleted, - created: entry.columns.added, + deleted: entry.policies.deleted.map(PgSquasher.unsquashPolicy), + created: entry.policies.added.map(PgSquasher.unsquashPolicy), }); if (created.length > 0) { @@ -1043,31 +1046,37 @@ export const applyPgSnapshotsDiff = async ( const jsonAlterPoliciesStatements: JsonAlterPolicyStatement[] = []; const jsonRenamePoliciesStatements: JsonRenamePolicyStatement[] = []; + const jsonEnableRLSStatements: JsonEnableRLSStatement[] = []; + const jsonDisableRLSStatements: JsonDisableRLSStatement[] = []; + for (let it of policyRenames) { jsonRenamePoliciesStatements.push( ...prepareRenamePolicyJsons(it.table, it.schema, it.renames), ); } - alteredTables.forEach((it) => { - // handle policies - + for (const it of policyCreates) { jsonCreatePoliciesStatements.push( ...prepareCreatePolicyJsons( - it.name, + it.table, it.schema, - it.addedPolicies || {}, + it.columns, ), ); + } + for (const it of policyDeletes) { jsonDropPoliciesStatements.push( ...prepareDropPolicyJsons( - it.name, + it.table, it.schema, - it.deletedPolicies || {}, + it.columns, ), ); + } + 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); @@ -1077,7 +1086,7 @@ export const applyPgSnapshotsDiff = async ( ...prepareDropPolicyJsons( it.name, it.schema, - { [oldPolicy.name]: it.alteredPolicies[policyName].__old }, + [oldPolicy], ), ); @@ -1085,9 +1094,10 @@ export const applyPgSnapshotsDiff = async ( ...prepareCreatePolicyJsons( it.name, it.schema, - { [newPolicy.name]: it.alteredPolicies[policyName].__new }, + [newPolicy], ), ); + return; } if (newPolicy.for !== oldPolicy.for) { @@ -1095,7 +1105,7 @@ export const applyPgSnapshotsDiff = async ( ...prepareDropPolicyJsons( it.name, it.schema, - { [oldPolicy.name]: it.alteredPolicies[policyName].__old }, + [oldPolicy], ), ); @@ -1103,9 +1113,10 @@ export const applyPgSnapshotsDiff = async ( ...prepareCreatePolicyJsons( it.name, it.schema, - { [newPolicy.name]: it.alteredPolicies[policyName].__new }, + [newPolicy], ), ); + return; } // alter @@ -1119,6 +1130,33 @@ export const applyPgSnapshotsDiff = async ( ); }); + // if there were no tables in json1 and is a table in json2 and it has policies -> enable + // if there was a table in json1 and is no table in json2 and it had policies -> disable + + // Handle enabling and disabling RLS + for (const table of Object.values(json2.tables)) { + const policiesInCurrentState = Object.keys(table.policies); + const tableInPreviousState = + columnsPatchedSnap1.tables[`${table.schema === '' ? 'public' : table.schema}.${table.name}`]; + const policiesInPreviousState = tableInPreviousState ? Object.keys(tableInPreviousState.policies) : []; + + if (policiesInPreviousState.length === 0 && policiesInCurrentState.length > 0) { + jsonEnableRLSStatements.push({ type: 'enable_rls', tableName: table.name, schema: table.schema }); + } + + if (policiesInPreviousState.length > 0 && policiesInCurrentState.length === 0) { + jsonDisableRLSStatements.push({ type: 'disable_rls', tableName: table.name, schema: table.schema }); + } + } + + for (const table of Object.values(columnsPatchedSnap1.tables)) { + const tableInCurrentState = json2.tables[`${table.schema === '' ? 'public' : table.schema}.${table.name}`]; + + if (tableInCurrentState === undefined) { + jsonDisableRLSStatements.push({ type: 'disable_rls', tableName: table.name, schema: table.schema }); + } + } + // TODO // Add to sql generators // Test generate logic manually @@ -1315,6 +1353,9 @@ export const applyPgSnapshotsDiff = async ( jsonStatements.push(...createTables); + jsonStatements.push(...jsonEnableRLSStatements); + jsonStatements.push(...jsonDisableRLSStatements); + jsonStatements.push(...jsonDropTables); jsonStatements.push(...jsonSetTableSchemas); jsonStatements.push(...jsonRenameTables); @@ -1346,10 +1387,10 @@ export const applyPgSnapshotsDiff = async ( jsonStatements.push(...jsonAlteredUniqueConstraints); - jsonStatements.push(...jsonCreatePoliciesStatements); jsonStatements.push(...jsonRenamePoliciesStatements); - jsonStatements.push(...jsonAlterPoliciesStatements); jsonStatements.push(...jsonDropPoliciesStatements); + jsonStatements.push(...jsonCreatePoliciesStatements); + jsonStatements.push(...jsonAlterPoliciesStatements); jsonStatements.push(...dropEnums); jsonStatements.push(...dropSequences); diff --git a/drizzle-kit/src/sqlgenerator.ts b/drizzle-kit/src/sqlgenerator.ts index 5d6e1e002..dad888c20 100644 --- a/drizzle-kit/src/sqlgenerator.ts +++ b/drizzle-kit/src/sqlgenerator.ts @@ -39,11 +39,13 @@ import { JsonDeleteCompositePK, JsonDeleteReferenceStatement, JsonDeleteUniqueConstraint, + JsonDisableRLSStatement, JsonDropColumnStatement, JsonDropIndexStatement, JsonDropPolicyStatement, JsonDropSequenceStatement, JsonDropTableStatement, + JsonEnableRLSStatement, JsonMoveSequenceStatement, JsonPgCreateIndexStatement, JsonRenameColumnStatement, @@ -140,7 +142,7 @@ class PgCreatePolicyConvertor extends Convertor { return statement.type === 'create_policy' && dialect === 'postgresql'; } override convert(statement: JsonCreatePolicyStatement): string | string[] { - const policy = PgSquasher.unsquashPolicy(statement.data); + const policy = statement.data; const tableNameWithSchema = statement.schema ? `"${statement.schema}"."${statement.tableName}"` @@ -150,7 +152,9 @@ class PgCreatePolicyConvertor extends Convertor { const withCheckPart = policy.withCheck ? ` WITH CHECK (${policy.withCheck})` : ''; - return `CREATE POLICY ${policy.name} ON ${tableNameWithSchema} AS ${policy.as?.toUpperCase()} FOR ${policy.for?.toUpperCase()} TO ${policy.to?.toUpperCase()}${usingPart}${withCheckPart}`; + return `CREATE POLICY "${policy.name}" ON ${tableNameWithSchema} AS ${policy.as?.toUpperCase()} FOR ${policy.for?.toUpperCase()} TO ${ + policy.to?.join(', ') + }${usingPart}${withCheckPart};`; } } @@ -159,13 +163,13 @@ class PgDropPolicyConvertor extends Convertor { return statement.type === 'drop_policy' && dialect === 'postgresql'; } override convert(statement: JsonDropPolicyStatement): string | string[] { - const policy = PgSquasher.unsquashPolicy(statement.data); + const policy = statement.data; const tableNameWithSchema = statement.schema ? `"${statement.schema}"."${statement.tableName}"` : `"${statement.tableName}"`; - return `DROP POLICY ${policy.name} ON ${tableNameWithSchema} CASCADE`; + return `DROP POLICY "${policy.name}" ON ${tableNameWithSchema} CASCADE;`; } } @@ -178,7 +182,7 @@ class PgRenamePolicyConvertor extends Convertor { ? `"${statement.schema}"."${statement.tableName}"` : `"${statement.tableName}"`; - return `ALTER POLICY ${statement.oldName} ON ${tableNameWithSchema} RENAME TO ${statement.newName}`; + return `ALTER POLICY "${statement.oldName}" ON ${tableNameWithSchema} RENAME TO "${statement.newName}";`; } } @@ -206,7 +210,33 @@ class PgAlterPolicyConvertor extends Convertor { ? ` WITH CHECK (${oldPolicy.withCheck})` : ''; - return `ALTER POLICY ${oldPolicy.name} ON ${tableNameWithSchema} TO ${newPolicy.to}${usingPart}${withCheckPart}`; + return `ALTER POLICY "${oldPolicy.name}" ON ${tableNameWithSchema} TO ${newPolicy.to}${usingPart}${withCheckPart};`; + } +} + +class PgEnableRlsConvertor extends Convertor { + override can(statement: JsonStatement, dialect: Dialect): boolean { + return statement.type === 'enable_rls' && dialect === 'postgresql'; + } + override convert(statement: JsonEnableRLSStatement): string { + const tableNameWithSchema = statement.schema + ? `"${statement.schema}"."${statement.tableName}"` + : `"${statement.tableName}"`; + + return `ALTER TABLE ${tableNameWithSchema} ENABLE ROW LEVEL SECURITY;`; + } +} + +class PgDisableRlsConvertor extends Convertor { + override can(statement: JsonStatement, dialect: Dialect): boolean { + return statement.type === 'disable_rls' && dialect === 'postgresql'; + } + override convert(statement: JsonDisableRLSStatement): string { + const tableNameWithSchema = statement.schema + ? `"${statement.schema}"."${statement.tableName}"` + : `"${statement.tableName}"`; + + return `ALTER TABLE ${tableNameWithSchema} DISABLE ROW LEVEL SECURITY;`; } } @@ -215,8 +245,8 @@ class PgCreateTableConvertor extends Convertor { return statement.type === 'create_table' && dialect === 'postgresql'; } - convert(st: JsonCreateTableStatement) { - const { tableName, schema, columns, compositePKs, uniqueConstraints } = st; + convert(st: JsonCreateTableStatement): string[] { + const { tableName, schema, columns, compositePKs, uniqueConstraints, policies } = st; let statement = ''; const name = schema ? `"${schema}"."${tableName}"` : `"${tableName}"`; @@ -304,7 +334,22 @@ class PgCreateTableConvertor extends Convertor { statement += `\n);`; statement += `\n`; - return statement; + const createPolicyConvertor = new PgCreatePolicyConvertor(); + const createPolicies = policies?.map((p) => { + return createPolicyConvertor.convert({ + type: 'create_policy', + tableName, + data: PgSquasher.unsquashPolicy(p), + schema, + }) as string; + }) ?? []; + const enableRls = new PgEnableRlsConvertor().convert({ + type: 'enable_rls', + tableName, + schema, + }); + + return [statement, ...(createPolicies.length > 0 ? [enableRls] : []), ...createPolicies]; } } @@ -835,13 +880,26 @@ class PgDropTableConvertor extends Convertor { } convert(statement: JsonDropTableStatement) { - const { tableName, schema } = statement; + const { tableName, schema, policies } = statement; const tableNameWithSchema = schema ? `"${schema}"."${tableName}"` : `"${tableName}"`; - return `DROP TABLE ${tableNameWithSchema};`; + const dropPolicyConvertor = new PgDropPolicyConvertor(); + const droppedPolicies = policies?.map((p) => { + return dropPolicyConvertor.convert({ + type: 'drop_policy', + tableName, + data: PgSquasher.unsquashPolicy(p), + schema, + }) as string; + }) ?? []; + + return [ + ...droppedPolicies, + `DROP TABLE ${tableNameWithSchema};`, + ]; } } @@ -2651,6 +2709,8 @@ convertors.push(new PgAlterPolicyConvertor()); convertors.push(new PgCreatePolicyConvertor()); convertors.push(new PgDropPolicyConvertor()); convertors.push(new PgRenamePolicyConvertor()); +convertors.push(new PgEnableRlsConvertor()); +convertors.push(new PgDisableRlsConvertor()); /// generated convertors.push(new PgAlterTableAlterColumnSetExpressionConvertor()); diff --git a/drizzle-kit/tests/mysql.test.ts b/drizzle-kit/tests/mysql.test.ts index e7b0b32a5..6d5406018 100644 --- a/drizzle-kit/tests/mysql.test.ts +++ b/drizzle-kit/tests/mysql.test.ts @@ -183,6 +183,7 @@ test('add table #6', async () => { }); expect(statements[1]).toStrictEqual({ type: 'drop_table', + policies: [], tableName: 'users1', schema: undefined, }); @@ -275,6 +276,7 @@ test('change table schema #1', async () => { expect(statements.length).toBe(1); expect(statements[0]).toStrictEqual({ type: 'drop_table', + policies: [], tableName: 'users', schema: undefined, }); diff --git a/drizzle-kit/tests/pg-identity.test.ts b/drizzle-kit/tests/pg-identity.test.ts index 906d812d4..400096671 100644 --- a/drizzle-kit/tests/pg-identity.test.ts +++ b/drizzle-kit/tests/pg-identity.test.ts @@ -45,6 +45,7 @@ test('create table: identity always/by default - no params', async () => { compositePKs: [], compositePkName: '', schema: '', + policies: [], tableName: 'users', type: 'create_table', uniqueConstraints: [], @@ -82,6 +83,7 @@ test('create table: identity always/by default - few params', async () => { ], compositePKs: [], compositePkName: '', + policies: [], schema: '', tableName: 'users', type: 'create_table', @@ -124,6 +126,7 @@ test('create table: identity always/by default - all params', async () => { ], compositePKs: [], compositePkName: '', + policies: [], schema: '', tableName: 'users', type: 'create_table', diff --git a/drizzle-kit/tests/pg-tables.test.ts b/drizzle-kit/tests/pg-tables.test.ts index 4171af333..6c862d211 100644 --- a/drizzle-kit/tests/pg-tables.test.ts +++ b/drizzle-kit/tests/pg-tables.test.ts @@ -31,6 +31,7 @@ test('add table #1', async () => { schema: '', columns: [], compositePKs: [], + policies: [], uniqueConstraints: [], compositePkName: '', }); @@ -59,6 +60,7 @@ test('add table #2', async () => { }, ], compositePKs: [], + policies: [], uniqueConstraints: [], compositePkName: '', }); @@ -98,6 +100,7 @@ test('add table #3', async () => { }, ], compositePKs: ['id;users_pk'], + policies: [], uniqueConstraints: [], compositePkName: 'users_pk', }); @@ -118,12 +121,14 @@ test('add table #4', async () => { schema: '', columns: [], compositePKs: [], + policies: [], uniqueConstraints: [], compositePkName: '', }); expect(statements[1]).toStrictEqual({ type: 'create_table', tableName: 'posts', + policies: [], schema: '', columns: [], compositePKs: [], @@ -152,6 +157,7 @@ test('add table #5', async () => { schema: 'folder', columns: [], compositePKs: [], + policies: [], uniqueConstraints: [], compositePkName: '', }); @@ -176,10 +182,12 @@ test('add table #6', async () => { columns: [], compositePKs: [], uniqueConstraints: [], + policies: [], compositePkName: '', }); expect(statements[1]).toStrictEqual({ type: 'drop_table', + policies: [], tableName: 'users1', schema: '', }); @@ -206,6 +214,7 @@ test('add table #7', async () => { schema: '', columns: [], compositePKs: [], + policies: [], uniqueConstraints: [], compositePkName: '', }); @@ -262,6 +271,7 @@ test('multiproject schema add table #1', async () => { }, ], compositePKs: [], + policies: [], compositePkName: '', uniqueConstraints: [], }); @@ -284,6 +294,7 @@ test('multiproject schema drop table #1', async () => { schema: '', tableName: 'prefix_users', type: 'drop_table', + policies: [], }); }); @@ -353,6 +364,7 @@ test('add schema + table #1', async () => { type: 'create_table', tableName: 'users', schema: 'folder', + policies: [], columns: [], compositePKs: [], uniqueConstraints: [], @@ -610,6 +622,7 @@ test('drop table + rename schema #1', async () => { type: 'drop_table', tableName: 'users', schema: 'folder2', + policies: [], }); }); diff --git a/drizzle-kit/tests/push/pg.test.ts b/drizzle-kit/tests/push/pg.test.ts index cb1a97122..d6d1212e6 100644 --- a/drizzle-kit/tests/push/pg.test.ts +++ b/drizzle-kit/tests/push/pg.test.ts @@ -618,6 +618,7 @@ const pgSuite: DialectSuite = { compositePkName: '', schema: '', tableName: 'users', + policies: [], type: 'create_table', uniqueConstraints: [], }, @@ -1386,6 +1387,7 @@ test('create table: identity always/by default - no params', async () => { compositePkName: '', schema: '', tableName: 'users', + policies: [], type: 'create_table', uniqueConstraints: [], }, @@ -1451,6 +1453,7 @@ test('create table: identity always/by default - few params', async () => { ], compositePKs: [], compositePkName: '', + policies: [], schema: '', tableName: 'users', type: 'create_table', @@ -1527,6 +1530,7 @@ test('create table: identity always/by default - all params', async () => { schema: '', tableName: 'users', type: 'create_table', + policies: [], uniqueConstraints: [], }, ]); diff --git a/drizzle-kit/tests/rls/pg-policy.test.ts b/drizzle-kit/tests/rls/pg-policy.test.ts index 41289f4e4..7403fcce0 100644 --- a/drizzle-kit/tests/rls/pg-policy.test.ts +++ b/drizzle-kit/tests/rls/pg-policy.test.ts @@ -1,8 +1,9 @@ -import { integer, pgPolicy, pgTable } from 'drizzle-orm/pg-core'; +import { sql } from 'drizzle-orm'; +import { integer, pgPolicy, pgRole, pgTable } from 'drizzle-orm/pg-core'; import { diffTestSchemas } from 'tests/schemaDiffer'; import { expect, test } from 'vitest'; -test('add policy #1', async (t) => { +test('add policy + enable rls', async (t) => { const schema1 = { users: pgTable('users', { id: integer('id').primaryKey(), @@ -19,5 +20,656 @@ test('add policy #1', async (t) => { const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); - console.log(sqlStatements); + expect(sqlStatements).toStrictEqual([ + 'ALTER TABLE "users" ENABLE ROW LEVEL SECURITY;', + 'CREATE POLICY "test" ON "users" AS PERMISSIVE FOR ALL TO PUBLIC;', + ]); + expect(statements).toStrictEqual([ + { + schema: '', + tableName: 'users', + type: 'enable_rls', + }, + { + data: { + as: 'permissive', + for: 'all', + name: 'test', + to: ['PUBLIC'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'create_policy', + }, + ]); +}); + +test('drop policy + disable rls', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'ALTER TABLE "users" DISABLE ROW LEVEL SECURITY;', + 'DROP POLICY "test" ON "users" CASCADE;', + ]); + expect(statements).toStrictEqual([ + { + schema: '', + tableName: 'users', + type: 'disable_rls', + }, + { + data: { + as: 'permissive', + for: 'all', + name: 'test', + to: ['PUBLIC'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'drop_policy', + }, + ]); +}); + +test('add policy without enable rls', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + newrls: pgPolicy('newRls'), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'CREATE POLICY "newRls" ON "users" AS PERMISSIVE FOR ALL TO PUBLIC;', + ]); + expect(statements).toStrictEqual([ + { + data: { + as: 'permissive', + for: 'all', + name: 'newRls', + to: ['PUBLIC'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'create_policy', + }, + ]); +}); + +test('drop policy without disable rls', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + oldRls: pgPolicy('oldRls'), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'DROP POLICY "oldRls" ON "users" CASCADE;', + ]); + expect(statements).toStrictEqual([ + { + data: { + as: 'permissive', + for: 'all', + name: 'oldRls', + to: ['PUBLIC'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'drop_policy', + }, + ]); +}); + +test('alter policy without recreation: changing roles', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive', to: 'CURRENT_ROLE' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'ALTER POLICY "test" ON "users" TO CURRENT_ROLE;', + ]); + expect(statements).toStrictEqual([ + { + newData: 'test--permissive--all--CURRENT_ROLE--undefined--undefined', + oldData: 'test--permissive--all--PUBLIC--undefined--undefined', + schema: '', + tableName: 'users', + type: 'alter_policy', + }, + ]); +}); + +test('alter policy without recreation: changing using', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive', using: sql`true` }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'ALTER POLICY "test" ON "users" TO PUBLIC USING (true);', + ]); + expect(statements).toStrictEqual([ + { + newData: 'test--permissive--all--PUBLIC--true--undefined', + oldData: 'test--permissive--all--PUBLIC--undefined--undefined', + schema: '', + tableName: 'users', + type: 'alter_policy', + }, + ]); +}); + +test('alter policy without recreation: changing with check', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive', withCheck: sql`true` }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'ALTER POLICY "test" ON "users" TO PUBLIC WITH CHECK (true);', + ]); + expect(statements).toStrictEqual([ + { + newData: 'test--permissive--all--PUBLIC--undefined--true', + oldData: 'test--permissive--all--PUBLIC--undefined--undefined', + schema: '', + tableName: 'users', + type: 'alter_policy', + }, + ]); +}); + +/// + +test('alter policy with recreation: changing as', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'restrictive' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'DROP POLICY "test" ON "users" CASCADE;', + 'CREATE POLICY "test" ON "users" AS RESTRICTIVE FOR ALL TO PUBLIC;', + ]); + expect(statements).toStrictEqual([ + { + data: { + as: 'permissive', + for: 'all', + name: 'test', + to: ['PUBLIC'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'drop_policy', + }, + { + data: { + as: 'restrictive', + for: 'all', + name: 'test', + to: ['PUBLIC'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'create_policy', + }, + ]); +}); + +test('alter policy with recreation: changing for', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive', for: 'delete' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'DROP POLICY "test" ON "users" CASCADE;', + 'CREATE POLICY "test" ON "users" AS PERMISSIVE FOR DELETE TO PUBLIC;', + ]); + expect(statements).toStrictEqual([ + { + data: { + as: 'permissive', + for: 'all', + name: 'test', + to: ['PUBLIC'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'drop_policy', + }, + { + data: { + as: 'permissive', + for: 'delete', + name: 'test', + to: ['PUBLIC'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'create_policy', + }, + ]); +}); + +test('alter policy with recreation: changing both "as" and "for"', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'restrictive', for: 'insert' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'DROP POLICY "test" ON "users" CASCADE;', + 'CREATE POLICY "test" ON "users" AS RESTRICTIVE FOR INSERT TO PUBLIC;', + ]); + expect(statements).toStrictEqual([ + { + data: { + as: 'permissive', + for: 'all', + name: 'test', + to: ['PUBLIC'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'drop_policy', + }, + { + data: { + as: 'restrictive', + for: 'insert', + name: 'test', + to: ['PUBLIC'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'create_policy', + }, + ]); +}); + +test('alter policy with recreation: changing all fields', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive', for: 'select', using: sql`true` }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'restrictive', to: 'CURRENT_ROLE', withCheck: sql`true` }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'DROP POLICY "test" ON "users" CASCADE;', + 'CREATE POLICY "test" ON "users" AS RESTRICTIVE FOR ALL TO CURRENT_ROLE WITH CHECK (true);', + ]); + expect(statements).toStrictEqual([ + { + data: { + as: 'permissive', + for: 'select', + name: 'test', + to: ['PUBLIC'], + using: 'true', + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'drop_policy', + }, + { + data: { + as: 'restrictive', + for: 'all', + name: 'test', + to: ['CURRENT_ROLE'], + using: undefined, + withCheck: 'true', + }, + schema: '', + tableName: 'users', + type: 'create_policy', + }, + ]); +}); + +test('rename policy', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('newName', { as: 'permissive' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, [ + 'public.users.test->public.users.newName', + ]); + + expect(sqlStatements).toStrictEqual([ + 'ALTER POLICY "test" ON "users" RENAME TO "newName";', + ]); + expect(statements).toStrictEqual([ + { + newName: 'newName', + oldName: 'test', + schema: '', + tableName: 'users', + type: 'rename_policy', + }, + ]); +}); + +test('rename policy in renamed table', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users2', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('newName', { as: 'permissive' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, [ + 'public.users->public.users2', + 'public.users2.test->public.users2.newName', + ]); + + expect(sqlStatements).toStrictEqual([ + 'ALTER TABLE "users" RENAME TO "users2";', + 'ALTER POLICY "test" ON "users2" RENAME TO "newName";', + ]); + expect(statements).toStrictEqual([ + { + fromSchema: '', + tableNameFrom: 'users', + tableNameTo: 'users2', + toSchema: '', + type: 'rename_table', + }, + { + newName: 'newName', + oldName: 'test', + schema: '', + tableName: 'users2', + type: 'rename_policy', + }, + ]); +}); + +test('create table with a policy', async (t) => { + const schema1 = {}; + + const schema2 = { + users: pgTable('users2', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'CREATE TABLE IF NOT EXISTS "users2" (\n\t"id" integer PRIMARY KEY NOT NULL\n);\n', + 'ALTER TABLE "users2" ENABLE ROW LEVEL SECURITY;', + 'CREATE POLICY "test" ON "users2" AS PERMISSIVE FOR ALL TO PUBLIC;', + ]); + expect(statements).toStrictEqual([ + { + columns: [ + { + name: 'id', + notNull: true, + primaryKey: true, + type: 'integer', + }, + ], + compositePKs: [], + compositePkName: '', + policies: [ + 'test--permissive--all--PUBLIC--undefined--undefined', + ], + schema: '', + tableName: 'users2', + type: 'create_table', + uniqueConstraints: [], + }, + ]); +}); + +test('drop table with a policy', async (t) => { + const schema1 = { + users: pgTable('users2', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = {}; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'DROP POLICY "test" ON "users2" CASCADE;', + 'DROP TABLE "users2";', + ]); + expect(statements).toStrictEqual([ + { + policies: [ + 'test--permissive--all--PUBLIC--undefined--undefined', + ], + schema: '', + tableName: 'users2', + type: 'drop_table', + }, + ]); +}); + +test('add policy with multiple "to" roles', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { to: ['CURRENT_ROLE', pgRole('manager').existing()] }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'ALTER TABLE "users" ENABLE ROW LEVEL SECURITY;', + 'CREATE POLICY "test" ON "users" AS PERMISSIVE FOR ALL TO CURRENT_ROLE, manager;', + ]); + expect(statements).toStrictEqual([ + { + schema: '', + tableName: 'users', + type: 'enable_rls', + }, + { + data: { + as: 'permissive', + for: 'all', + name: 'test', + to: ['CURRENT_ROLE', 'manager'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'create_policy', + }, + ]); }); diff --git a/drizzle-kit/tests/rls/pg-role.test.ts b/drizzle-kit/tests/rls/pg-role.test.ts new file mode 100644 index 000000000..e69de29bb diff --git a/drizzle-kit/tests/schemaDiffer.ts b/drizzle-kit/tests/schemaDiffer.ts index b9d7b5fde..7b912c6cb 100644 --- a/drizzle-kit/tests/schemaDiffer.ts +++ b/drizzle-kit/tests/schemaDiffer.ts @@ -2,7 +2,7 @@ import { PGlite } from '@electric-sql/pglite'; import { Database } from 'better-sqlite3'; import { is } from 'drizzle-orm'; import { MySqlSchema, MySqlTable } from 'drizzle-orm/mysql-core'; -import { isPgEnum, isPgSequence, PgEnum, PgSchema, PgSequence, PgTable } from 'drizzle-orm/pg-core'; +import { isPgEnum, isPgSequence, PgEnum, PgRole, PgSchema, PgSequence, PgTable } from 'drizzle-orm/pg-core'; import { SQLiteTable } from 'drizzle-orm/sqlite-core'; import * as fs from 'fs'; import { Connection } from 'mysql2/promise'; @@ -16,6 +16,7 @@ import { tablesResolver, } from 'src/cli/commands/migrate'; import { logSuggestionsAndReturn } from 'src/cli/commands/sqlitePushUtils'; +import { Entities } from 'src/cli/validations/cli'; import { schemaToTypeScript as schemaToTypeScriptMySQL } from 'src/introspect-mysql'; import { schemaToTypeScript } from 'src/introspect-pg'; import { schemaToTypeScript as schemaToTypeScriptSQLite } from 'src/introspect-sqlite'; @@ -24,7 +25,7 @@ import { mysqlSchema, squashMysqlScheme } from 'src/serializer/mysqlSchema'; import { generateMySqlSnapshot } from 'src/serializer/mysqlSerializer'; import { fromDatabase as fromMySqlDatabase } from 'src/serializer/mysqlSerializer'; import { prepareFromPgImports } from 'src/serializer/pgImports'; -import { pgSchema, Policy, squashPgScheme } from 'src/serializer/pgSchema'; +import { pgSchema, PgSquasher, Policy, squashPgScheme } from 'src/serializer/pgSchema'; import { fromDatabase, generatePgSnapshot } from 'src/serializer/pgSerializer'; import { prepareFromSqliteImports } from 'src/serializer/sqliteImports'; import { sqliteSchema, squashSqliteScheme } from 'src/serializer/sqliteSchema'; @@ -47,7 +48,7 @@ import { export type PostgresSchema = Record< string, - PgTable | PgEnum | PgSchema | PgSequence + PgTable | PgEnum | PgSchema | PgSequence | PgRole >; export type MysqlSchema = Record | MySqlSchema>; export type SqliteSchema = Record>; @@ -500,11 +501,14 @@ export const diffTestSchemasPush = async ( const leftSequences = Object.values(right).filter((it) => isPgSequence(it)) as PgSequence[]; + const leftRoles = Object.values(right).filter((it) => is(it, PgRole)) as PgRole[]; + const serialized2 = generatePgSnapshot( leftTables, leftEnums, leftSchemas, leftSequences, + leftRoles, ); const { version: v1, dialect: d1, ...rest1 } = introspectedSchema; @@ -592,7 +596,9 @@ export const applyPgDiffs = async (sn: PostgresSchema) => { const sequences = Object.values(sn).filter((it) => isPgSequence(it)) as PgSequence[]; - const serialized1 = generatePgSnapshot(tables, enums, schemas, sequences); + const roles = Object.values(sn).filter((it) => is(it, PgRole)) as PgRole[]; + + const serialized1 = generatePgSnapshot(tables, enums, schemas, sequences, roles); const { version: v1, dialect: d1, ...rest1 } = serialized1; @@ -646,17 +652,23 @@ export const diffTestSchemas = async ( const rightSequences = Object.values(right).filter((it) => isPgSequence(it)) as PgSequence[]; + const leftRoles = Object.values(left).filter((it) => is(it, PgRole)) as PgRole[]; + + const rightRoles = Object.values(right).filter((it) => is(it, PgRole)) as PgRole[]; + const serialized1 = generatePgSnapshot( leftTables, leftEnums, leftSchemas, leftSequences, + leftRoles, ); const serialized2 = generatePgSnapshot( rightTables, rightEnums, rightSchemas, rightSequences, + rightRoles, ); const { version: v1, dialect: d1, ...rest1 } = serialized1; @@ -1113,6 +1125,7 @@ export const introspectPgToFile = async ( initSchema: PostgresSchema, testName: string, schemas: string[] = ['public'], + entities: Entities, ) => { // put in db const { sqlStatements } = await applyPgDiffs(initSchema); @@ -1130,6 +1143,7 @@ export const introspectPgToFile = async ( }, undefined, schemas, + entities, ); const file = schemaToTypeScript(introspectedSchema, 'camel'); @@ -1145,6 +1159,7 @@ export const introspectPgToFile = async ( response.enums, response.schemas, response.sequences, + response.roles, ); const { version: v2, dialect: d2, ...rest2 } = afterFileImports; @@ -1168,11 +1183,14 @@ export const introspectPgToFile = async ( const leftSequences = Object.values(initSchema).filter((it) => isPgSequence(it)) as PgSequence[]; + const leftRoles = Object.values(initSchema).filter((it) => is(it, PgRole)) as PgRole[]; + const initSnapshot = generatePgSnapshot( leftTables, leftEnums, leftSchemas, leftSequences, + leftRoles, ); const { version: initV, dialect: initD, ...initRest } = initSnapshot; diff --git a/drizzle-kit/tests/sqlite-tables.test.ts b/drizzle-kit/tests/sqlite-tables.test.ts index d7781f150..9afbf1a40 100644 --- a/drizzle-kit/tests/sqlite-tables.test.ts +++ b/drizzle-kit/tests/sqlite-tables.test.ts @@ -143,6 +143,7 @@ test('add table #6', async () => { type: 'drop_table', tableName: 'users1', schema: undefined, + policies: [], }); }); diff --git a/drizzle-orm/src/pg-core/index.ts b/drizzle-orm/src/pg-core/index.ts index 2e2d5399d..ebc436bc1 100644 --- a/drizzle-orm/src/pg-core/index.ts +++ b/drizzle-orm/src/pg-core/index.ts @@ -8,6 +8,7 @@ export * from './indexes.ts'; export * from './policies.ts'; export * from './primary-keys.ts'; export * from './query-builders/index.ts'; +export * from './roles.ts'; export * from './schema.ts'; export * from './sequence.ts'; export * from './session.ts'; diff --git a/drizzle-orm/src/pg-core/policies.ts b/drizzle-orm/src/pg-core/policies.ts index 7f49662b6..11e898eda 100644 --- a/drizzle-orm/src/pg-core/policies.ts +++ b/drizzle-orm/src/pg-core/policies.ts @@ -1,10 +1,20 @@ import { entityKind } from '~/entity.ts'; import type { SQL } from '~/sql/sql.ts'; +import type { PgRole } from './roles'; + +export type PgPolicyToOption = + | 'PUBLIC' + | 'CURRENT_ROLE' + | 'CURRENT_USER' + | 'SESSION_USER' + | (string & {}) + | PgPolicyToOption[] + | PgRole; export interface PgPolicyConfig { as?: 'permissive' | 'restrictive'; for?: 'all' | 'select' | 'insert' | 'update' | 'delete'; - to?: 'PUBLIC' | 'CURRENT_ROLE' | 'CURRENT_USER' | 'SESSION_USER' | (string & {}); + to?: PgPolicyToOption; using?: SQL; withCheck?: SQL; } diff --git a/drizzle-orm/src/pg-core/roles.ts b/drizzle-orm/src/pg-core/roles.ts new file mode 100644 index 000000000..a2c77c303 --- /dev/null +++ b/drizzle-orm/src/pg-core/roles.ts @@ -0,0 +1,41 @@ +import { entityKind } from '~/entity.ts'; + +export interface PgRoleConfig { + createDb?: boolean; + createRole?: boolean; + inherit?: boolean; +} + +export class PgRole implements PgRoleConfig { + static readonly [entityKind]: string = 'PgRole'; + + /** @internal */ + _existing?: boolean; + + /** @internal */ + readonly createDb: PgRoleConfig['createDb']; + /** @internal */ + readonly createRole: PgRoleConfig['createRole']; + /** @internal */ + readonly inherit: PgRoleConfig['inherit']; + + constructor( + readonly name: string, + config?: PgRoleConfig, + ) { + if (config) { + this.createDb = config.createDb; + this.createRole = config.createRole; + this.inherit = config.inherit; + } + } + + existing(): this { + this._existing = true; + return this; + } +} + +export function pgRole(name: string, config?: PgRoleConfig) { + return new PgRole(name, config); +} From 5ec6b993163d54b5531c5469f773ffbc294eb4cb Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Sun, 1 Sep 2024 17:09:12 +0300 Subject: [PATCH 04/24] Add roles generate and introspect logic Add generate tests for roles --- drizzle-kit/src/cli/commands/migrate.ts | 94 ++++++++- drizzle-kit/src/cli/views.ts | 73 ++++++- drizzle-kit/src/jsonDiffer.js | 9 + drizzle-kit/src/jsonStatements.ts | 92 +++++++- drizzle-kit/src/serializer/pgSchema.ts | 2 + drizzle-kit/src/serializer/pgSerializer.ts | 2 +- drizzle-kit/src/snapshotsDiffer.ts | 107 ++++++++++ drizzle-kit/src/sqlgenerator.ts | 55 +++++ drizzle-kit/tests/rls/pg-policy.test.ts | 5 +- drizzle-kit/tests/rls/pg-role.test.ts | 234 +++++++++++++++++++++ drizzle-kit/tests/schemaDiffer.ts | 70 +++++- drizzle-kit/vitest.config.ts | 1 + 12 files changed, 737 insertions(+), 7 deletions(-) diff --git a/drizzle-kit/src/cli/commands/migrate.ts b/drizzle-kit/src/cli/commands/migrate.ts index b5f369bc7..9f679bc76 100644 --- a/drizzle-kit/src/cli/commands/migrate.ts +++ b/drizzle-kit/src/cli/commands/migrate.ts @@ -14,7 +14,7 @@ import path, { join } from 'path'; import { TypeOf } from 'zod'; import type { CommonSchema } from '../../schemaValidator'; import { MySqlSchema, mysqlSchema, squashMysqlScheme } from '../../serializer/mysqlSchema'; -import { PgSchema, pgSchema, Policy, squashPgScheme } from '../../serializer/pgSchema'; +import { PgSchema, pgSchema, Policy, Role, squashPgScheme } from '../../serializer/pgSchema'; import { SQLiteSchema, sqliteSchema, squashSqliteScheme } from '../../serializer/sqliteSchema'; import { applyMysqlSnapshotsDiff, @@ -27,6 +27,8 @@ import { ResolverInput, ResolverOutput, ResolverOutputWithMoved, + RolesResolverInput, + RolesResolverOutput, Sequence, Table, } from '../../snapshotsDiffer'; @@ -40,6 +42,7 @@ import { ResolveColumnSelect, ResolveSchemasSelect, ResolveSelect, + ResolveSelectNamed, schema, } from '../views'; import { GenerateConfig } from './utils'; @@ -113,6 +116,21 @@ export const sequencesResolver = async ( } }; +export const roleResolver = async ( + input: RolesResolverInput, +): Promise> => { + const result = await promptNamedConflict( + input.created, + input.deleted, + 'role', + ); + return { + created: result.created, + deleted: result.deleted, + renamed: result.renamed, + }; +}; + export const policyResolver = async ( input: ColumnsResolverInput, ): Promise> => { @@ -213,6 +231,7 @@ export const prepareAndMigratePg = async (config: GenerateConfig) => { enumsResolver, sequencesResolver, policyResolver, + roleResolver, tablesResolver, columnsResolver, validatedPrev, @@ -257,6 +276,7 @@ export const preparePgPush = async ( enumsResolver, sequencesResolver, policyResolver, + roleResolver, tablesResolver, columnsResolver, validatedPrev, @@ -576,6 +596,78 @@ export const promptColumnsConflicts = async ( return result; }; +export const promptNamedConflict = async ( + newItems: T[], + missingItems: T[], + entity: 'role', +): Promise<{ + created: T[]; + renamed: { from: T; to: T }[]; + deleted: T[]; +}> => { + if (missingItems.length === 0 || newItems.length === 0) { + return { + created: newItems, + renamed: [], + deleted: missingItems, + }; + } + + const result: { + created: T[]; + renamed: { from: T; to: T }[]; + deleted: T[]; + } = { created: [], renamed: [], deleted: [] }; + let index = 0; + let leftMissing = [...missingItems]; + do { + const created = newItems[index]; + const renames: RenamePropmtItem[] = leftMissing.map((it) => { + return { from: it, to: created }; + }); + + const promptData: (RenamePropmtItem | T)[] = [created, ...renames]; + + const { status, data } = await render( + new ResolveSelectNamed(created, promptData, entity), + ); + if (status === 'aborted') { + console.error('ERROR'); + process.exit(1); + } + + if (isRenamePromptItem(data)) { + console.log( + `${chalk.yellow('~')} ${data.from.name} › ${data.to.name} ${ + chalk.gray( + `${entity} will be renamed/moved`, + ) + }`, + ); + + if (data.from.name !== data.to.name) { + result.renamed.push(data); + } + + delete leftMissing[leftMissing.indexOf(data.from)]; + leftMissing = leftMissing.filter(Boolean); + } else { + console.log( + `${chalk.green('+')} ${data.name} ${ + chalk.gray( + `${entity} will be created`, + ) + }`, + ); + result.created.push(created); + } + index += 1; + } while (index < newItems.length); + console.log(chalk.gray(`--- all ${entity} conflicts resolved ---\n`)); + result.deleted.push(...leftMissing); + return result; +}; + export const promptNamedWithSchemasConflict = async ( newItems: T[], missingItems: T[], diff --git a/drizzle-kit/src/cli/views.ts b/drizzle-kit/src/cli/views.ts index 56e0331df..5cae2ac3f 100644 --- a/drizzle-kit/src/cli/views.ts +++ b/drizzle-kit/src/cli/views.ts @@ -158,6 +158,77 @@ export const tableKey = (it: NamedWithSchema) => { : `${it.schema}.${it.name}`; }; +export class ResolveSelectNamed extends Prompt< + RenamePropmtItem | T +> { + private readonly state: SelectState | T>; + + constructor( + private readonly base: T, + data: (RenamePropmtItem | T)[], + private readonly entityType: 'role', + ) { + super(); + this.on('attach', (terminal) => terminal.toggleCursor('hide')); + this.state = new SelectState(data); + this.state.bind(this); + this.base = base; + } + + render(status: 'idle' | 'submitted' | 'aborted'): string { + if (status === 'submitted' || status === 'aborted') { + return ''; + } + const key = this.base.name; + + let text = `\nIs ${chalk.bold.blue(key)} ${this.entityType} created or renamed from another ${this.entityType}?\n`; + + const isSelectedRenamed = isRenamePromptItem( + this.state.items[this.state.selectedIdx], + ); + + const selectedPrefix = isSelectedRenamed + ? chalk.yellow('❯ ') + : chalk.green('❯ '); + + const labelLength: number = this.state.items + .filter((it) => isRenamePromptItem(it)) + .map((_) => { + const it = _ as RenamePropmtItem; + const keyFrom = it.from.name; + return key.length + 3 + keyFrom.length; + }) + .reduce((a, b) => { + if (a > b) { + return a; + } + return b; + }, 0); + + const entityType = this.entityType; + this.state.items.forEach((it, idx) => { + const isSelected = idx === this.state.selectedIdx; + const isRenamed = isRenamePromptItem(it); + + const title = isRenamed + ? `${it.from.name} › ${it.to.name}`.padEnd(labelLength, ' ') + : it.name.padEnd(labelLength, ' '); + + const label = isRenamed + ? `${chalk.yellow('~')} ${title} ${chalk.gray(`rename ${entityType}`)}` + : `${chalk.green('+')} ${title} ${chalk.gray(`create ${entityType}`)}`; + + text += isSelected ? `${selectedPrefix}${label}` : ` ${label}`; + text += idx != this.state.items.length - 1 ? '\n' : ''; + }); + return text; + } + + result(): RenamePropmtItem | T { + return this.state.items[this.state.selectedIdx]!; + } +} + export class ResolveSelect extends Prompt< RenamePropmtItem | T > { @@ -166,7 +237,7 @@ export class ResolveSelect extends Prompt< constructor( private readonly base: T, data: (RenamePropmtItem | T)[], - private readonly entityType: 'table' | 'enum' | 'sequence', + private readonly entityType: 'table' | 'enum' | 'sequence' | 'role', ) { super(); this.on('attach', (terminal) => terminal.toggleCursor('hide')); diff --git a/drizzle-kit/src/jsonDiffer.js b/drizzle-kit/src/jsonDiffer.js index 7f72b0c26..15c636c09 100644 --- a/drizzle-kit/src/jsonDiffer.js +++ b/drizzle-kit/src/jsonDiffer.js @@ -201,6 +201,7 @@ export function applyJsonDiff(json1, json2) { difference.tables = difference.tables || {}; difference.enums = difference.enums || {}; difference.sequences = difference.sequences || {}; + difference.roles = difference.roles || {}; // remove added/deleted schemas const schemaKeys = Object.keys(difference.schemas); @@ -282,6 +283,13 @@ export function applyJsonDiff(json1, json2) { return json2.sequences[it[0]]; }); + const rolesEntries = Object.entries(difference.roles); + const alteredRoles = rolesEntries + .filter((it) => !(it[0].includes('__added') || it[0].includes('__deleted'))) + .map((it) => { + return json2.roles[it[0]]; + }); + const alteredTablesWithColumns = Object.values(difference.tables).map( (table) => { return findAlternationsInTable(table); @@ -292,6 +300,7 @@ export function applyJsonDiff(json1, json2) { alteredTablesWithColumns, alteredEnums, alteredSequences, + alteredRoles, }; } diff --git a/drizzle-kit/src/jsonStatements.ts b/drizzle-kit/src/jsonStatements.ts index d98cde0e6..ef5f55c27 100644 --- a/drizzle-kit/src/jsonStatements.ts +++ b/drizzle-kit/src/jsonStatements.ts @@ -3,7 +3,7 @@ import { table } from 'console'; import { warning } from './cli/views'; import { CommonSquashedSchema, Dialect } from './schemaValidator'; import { MySqlKitInternals, MySqlSchema, MySqlSquasher } from './serializer/mysqlSchema'; -import { Index, PgSchema, PgSquasher, Policy } from './serializer/pgSchema'; +import { Index, PgSchema, PgSquasher, Policy, Role } from './serializer/pgSchema'; import { SQLiteKitInternals, SQLiteSquasher } from './serializer/sqliteSchema'; import { AlteredColumn, Column, Sequence, Table } from './snapshotsDiffer'; @@ -86,6 +86,40 @@ export interface JsonAddValueToEnumStatement { before: string; } +////// + +export interface JsonCreateRoleStatement { + type: 'create_role'; + name: string; + values: { + inherit?: boolean; + createDb?: boolean; + createRole?: boolean; + }; +} + +export interface JsonDropRoleStatement { + type: 'drop_role'; + name: string; +} +export interface JsonRenameRoleStatement { + type: 'rename_role'; + nameFrom: string; + nameTo: string; +} + +export interface JsonAlterRoleStatement { + type: 'alter_role'; + name: string; + values: { + inherit?: boolean; + createDb?: boolean; + createRole?: boolean; + }; +} + +////// + export interface JsonCreateSequenceStatement { type: 'create_sequence'; name: string; @@ -605,7 +639,11 @@ export type JsonStatement = | JsonAlterPolicyStatement | JsonRenamePolicyStatement | JsonEnableRLSStatement - | JsonDisableRLSStatement; + | JsonDisableRLSStatement + | JsonRenameRoleStatement + | JsonCreateRoleStatement + | JsonDropRoleStatement + | JsonAlterRoleStatement; export const preparePgCreateTableJson = ( table: Table, @@ -846,6 +884,56 @@ export const prepareRenameSequenceJson = ( //////////// +export const prepareCreateRoleJson = ( + role: Role, +): JsonCreateRoleStatement => { + return { + type: 'create_role', + name: role.name, + values: { + createDb: role.createDb, + createRole: role.createRole, + inherit: role.inherit, + }, + }; +}; + +export const prepareAlterRoleJson = ( + role: Role, +): JsonAlterRoleStatement => { + return { + type: 'alter_role', + name: role.name, + values: { + createDb: role.createDb, + createRole: role.createRole, + inherit: role.inherit, + }, + }; +}; + +export const prepareDropRoleJson = ( + name: string, +): JsonDropRoleStatement => { + return { + type: 'drop_role', + name: name, + }; +}; + +export const prepareRenameRoleJson = ( + nameFrom: string, + nameTo: string, +): JsonRenameRoleStatement => { + return { + type: 'rename_role', + nameFrom, + nameTo, + }; +}; + +////////// + export const prepareCreateSchemasJson = ( values: string[], ): JsonCreateSchema[] => { diff --git a/drizzle-kit/src/serializer/pgSchema.ts b/drizzle-kit/src/serializer/pgSchema.ts index b579ebdab..494dc2a78 100644 --- a/drizzle-kit/src/serializer/pgSchema.ts +++ b/drizzle-kit/src/serializer/pgSchema.ts @@ -437,6 +437,7 @@ export const pgSchemaSquashed = object({ enums: record(string(), enumSchema), schemas: record(string(), string()), sequences: record(string(), sequenceSquashed), + roles: record(string(), roleSchema).default({}), }).strict(); export const pgSchemaV3 = pgSchemaInternalV3.merge(schemaHash); @@ -746,6 +747,7 @@ export const squashPgScheme = ( enums: json.enums, schemas: json.schemas, sequences: mappedSequences, + roles: json.roles, }; }; diff --git a/drizzle-kit/src/serializer/pgSerializer.ts b/drizzle-kit/src/serializer/pgSerializer.ts index f9487684d..8490ca6aa 100644 --- a/drizzle-kit/src/serializer/pgSerializer.ts +++ b/drizzle-kit/src/serializer/pgSerializer.ts @@ -581,7 +581,7 @@ export const generatePgSnapshot = ( name: role.name, createDb: (role as any).createDb ?? false, createRole: (role as any).createRole ?? false, - inherit: (role as any).inherit ?? false, + inherit: (role as any).inherit ?? true, }; } } diff --git a/drizzle-kit/src/snapshotsDiffer.ts b/drizzle-kit/src/snapshotsDiffer.ts index a4f5a3f03..cd7ab95dd 100644 --- a/drizzle-kit/src/snapshotsDiffer.ts +++ b/drizzle-kit/src/snapshotsDiffer.ts @@ -38,6 +38,7 @@ import { JsonReferenceStatement, JsonRenameColumnStatement, JsonRenamePolicyStatement, + JsonRenameRoleStatement, JsonSqliteAddColumnStatement, JsonStatement, prepareAddCompositePrimaryKeyMySql, @@ -51,11 +52,13 @@ import { prepareAlterCompositePrimaryKeySqlite, prepareAlterPolicyJson, prepareAlterReferencesJson, + prepareAlterRoleJson, prepareAlterSequenceJson, prepareCreateEnumJson, prepareCreateIndexesJson, prepareCreatePolicyJsons, prepareCreateReferencesJson, + prepareCreateRoleJson, prepareCreateSchemasJson, prepareCreateSequenceJson, prepareDeleteCompositePrimaryKeyMySql, @@ -67,6 +70,7 @@ import { prepareDropIndexesJson, prepareDropPolicyJsons, prepareDropReferencesJson, + prepareDropRoleJson, prepareDropSequenceJson, prepareDropTableJson, prepareMoveEnumJson, @@ -78,6 +82,7 @@ import { prepareRenameColumns, prepareRenameEnumJson, prepareRenamePolicyJsons, + prepareRenameRoleJson, prepareRenameSchemasJson, prepareRenameSequenceJson, prepareRenameTableJson, @@ -93,6 +98,8 @@ import { PgSchemaSquashed, PgSquasher, Policy, + Role, + roleSchema, sequenceSchema, sequenceSquashed, } from './serializer/pgSchema'; @@ -279,6 +286,7 @@ export const diffResultScheme = object({ alteredTablesWithColumns: alteredTableScheme.array(), alteredEnums: changedEnumSchema.array(), alteredSequences: sequenceSquashed.array(), + alteredRoles: roleSchema.array(), }).strict(); export const diffResultSchemeMysql = object({ @@ -326,6 +334,17 @@ export interface ColumnsResolverInput { deleted: T[]; } +export interface RolesResolverInput { + created: T[]; + deleted: T[]; +} + +export interface RolesResolverOutput { + created: T[]; + renamed: { from: T; to: T }[]; + deleted: T[]; +} + export interface ColumnsResolverOutput { tableName: string; schema: string; @@ -396,6 +415,12 @@ const columnChangeFor = ( return column; }; +// resolve roles same as enums +// create new json statements +// sql generators + +// tests everything! + export const applyPgSnapshotsDiff = async ( json1: PgSchemaSquashed, json2: PgSchemaSquashed, @@ -411,6 +436,9 @@ export const applyPgSnapshotsDiff = async ( policyResolver: ( input: ColumnsResolverInput, ) => Promise>, + roleResolver: ( + input: RolesResolverInput, + ) => Promise>, tablesResolver: ( input: ResolverInput
, ) => Promise>, @@ -635,6 +663,60 @@ export const applyPgSnapshotsDiff = async ( }, ); + const rolesDiff = diffSchemasOrTables( + schemasPatchedSnap1.roles, + json2.roles, + ); + + const { + created: createdRoles, + deleted: deletedRoles, + renamed: renamedRoles, + } = await roleResolver({ + created: rolesDiff.added, + deleted: rolesDiff.deleted, + }); + + schemasPatchedSnap1.roles = mapEntries( + schemasPatchedSnap1.roles, + (_, it) => { + const { name } = nameChangeFor(it, renamedRoles); + it.name = name; + return [name, it]; + }, + ); + + const rolesChangeMap = renamedRoles.reduce( + (acc, it) => { + acc[it.from.name] = { + nameFrom: it.from.name, + nameTo: it.to.name, + }; + return acc; + }, + {} as Record< + string, + { + nameFrom: string; + nameTo: string; + } + >, + ); + + schemasPatchedSnap1.roles = mapEntries( + schemasPatchedSnap1.roles, + (roleKey, roleValue) => { + const key = roleKey; + const change = rolesChangeMap[key]; + + if (change) { + roleValue.name = change.nameTo; + } + + return [roleKey, roleValue]; + }, + ); + const tablesDiff = diffSchemasOrTables( schemasPatchedSnap1.tables as Record, json2.tables, @@ -1323,6 +1405,26 @@ export const applyPgSnapshotsDiff = async ( //////////// + const createRoles = createdRoles.map((it) => { + return prepareCreateRoleJson(it); + }) ?? []; + + const dropRoles = deletedRoles.map((it) => { + return prepareDropRoleJson(it.name); + }); + + const renameRoles = renamedRoles.map((it) => { + return prepareRenameRoleJson(it.from.name, it.to.name); + }); + + const jsonAlterRoles = typedResult.alteredRoles + .map((it) => { + return prepareAlterRoleJson(it); + }) + .flat() ?? []; + + //////////// + const createSchemas = prepareCreateSchemasJson( createdSchemas.map((it) => it.name), ); @@ -1392,6 +1494,11 @@ export const applyPgSnapshotsDiff = async ( jsonStatements.push(...jsonCreatePoliciesStatements); jsonStatements.push(...jsonAlterPoliciesStatements); + jsonStatements.push(...renameRoles); + jsonStatements.push(...dropRoles); + jsonStatements.push(...createRoles); + jsonStatements.push(...jsonAlterRoles); + jsonStatements.push(...dropEnums); jsonStatements.push(...dropSequences); jsonStatements.push(...dropSchemas); diff --git a/drizzle-kit/src/sqlgenerator.ts b/drizzle-kit/src/sqlgenerator.ts index dad888c20..a8d962c4d 100644 --- a/drizzle-kit/src/sqlgenerator.ts +++ b/drizzle-kit/src/sqlgenerator.ts @@ -22,6 +22,7 @@ import { JsonAlterCompositePK, JsonAlterPolicyStatement, JsonAlterReferenceStatement, + JsonAlterRoleStatement, JsonAlterSequenceStatement, JsonAlterTableRemoveFromSchema, JsonAlterTableSetNewSchema, @@ -32,6 +33,7 @@ import { JsonCreateIndexStatement, JsonCreatePolicyStatement, JsonCreateReferenceStatement, + JsonCreateRoleStatement, JsonCreateSchema, JsonCreateSequenceStatement, JsonCreateTableStatement, @@ -43,6 +45,7 @@ import { JsonDropColumnStatement, JsonDropIndexStatement, JsonDropPolicyStatement, + JsonDropRoleStatement, JsonDropSequenceStatement, JsonDropTableStatement, JsonEnableRLSStatement, @@ -50,6 +53,7 @@ import { JsonPgCreateIndexStatement, JsonRenameColumnStatement, JsonRenamePolicyStatement, + JsonRenameRoleStatement, JsonRenameSchema, JsonRenameSequenceStatement, JsonRenameTableStatement, @@ -137,6 +141,52 @@ abstract class Convertor { abstract convert(statement: JsonStatement): string | string[]; } +class PgCreateRoleConvertor extends Convertor { + override can(statement: JsonStatement, dialect: Dialect): boolean { + return statement.type === 'create_role' && dialect === 'postgresql'; + } + override convert(statement: JsonCreateRoleStatement): string | string[] { + return `CREATE ROLE "${statement.name}"${ + statement.values.createDb || statement.values.createRole || !statement.values.inherit + ? ` WITH${statement.values.createDb ? ' CREATEDB' : ''}${statement.values.createRole ? ' CREATEROLE' : ''}${ + statement.values.inherit ? '' : ' NOINHERIT' + }` + : '' + };`; + } +} + +class PgDropRoleConvertor extends Convertor { + override can(statement: JsonStatement, dialect: Dialect): boolean { + return statement.type === 'drop_role' && dialect === 'postgresql'; + } + override convert(statement: JsonDropRoleStatement): string | string[] { + return `DROP ROLE "${statement.name}";`; + } +} + +class PgRenameRoleConvertor extends Convertor { + override can(statement: JsonStatement, dialect: Dialect): boolean { + return statement.type === 'rename_role' && dialect === 'postgresql'; + } + override convert(statement: JsonRenameRoleStatement): string | string[] { + return `ALTER ROLE "${statement.nameFrom}" RENAME TO "${statement.nameTo}";`; + } +} + +class PgAlterRoleConvertor extends Convertor { + override can(statement: JsonStatement, dialect: Dialect): boolean { + return statement.type === 'alter_role' && dialect === 'postgresql'; + } + override convert(statement: JsonAlterRoleStatement): string | string[] { + return `ALTER ROLE "${statement.name}"${` WITH${statement.values.createDb ? ' CREATEDB' : ' NOCREATEDB'}${ + statement.values.createRole ? ' CREATEROLE' : ' NOCREATEROLE' + }${statement.values.inherit ? ' INHERIT' : ' NOINHERIT'}`};`; + } +} + +///// + class PgCreatePolicyConvertor extends Convertor { override can(statement: JsonStatement, dialect: Dialect): boolean { return statement.type === 'create_policy' && dialect === 'postgresql'; @@ -2712,6 +2762,11 @@ convertors.push(new PgRenamePolicyConvertor()); convertors.push(new PgEnableRlsConvertor()); convertors.push(new PgDisableRlsConvertor()); +convertors.push(new PgDropRoleConvertor()); +convertors.push(new PgAlterRoleConvertor()); +convertors.push(new PgCreateRoleConvertor()); +convertors.push(new PgRenameRoleConvertor()); + /// generated convertors.push(new PgAlterTableAlterColumnSetExpressionConvertor()); convertors.push(new PgAlterTableAlterColumnDropGeneratedConvertor()); diff --git a/drizzle-kit/tests/rls/pg-policy.test.ts b/drizzle-kit/tests/rls/pg-policy.test.ts index 7403fcce0..3f12b3e64 100644 --- a/drizzle-kit/tests/rls/pg-policy.test.ts +++ b/drizzle-kit/tests/rls/pg-policy.test.ts @@ -638,11 +638,14 @@ test('add policy with multiple "to" roles', async (t) => { }), }; + const role = pgRole('manager').existing(); + const schema2 = { + role, users: pgTable('users', { id: integer('id').primaryKey(), }, () => ({ - rls: pgPolicy('test', { to: ['CURRENT_ROLE', pgRole('manager').existing()] }), + rls: pgPolicy('test', { to: ['CURRENT_ROLE', role] }), })), }; diff --git a/drizzle-kit/tests/rls/pg-role.test.ts b/drizzle-kit/tests/rls/pg-role.test.ts index e69de29bb..a6b762955 100644 --- a/drizzle-kit/tests/rls/pg-role.test.ts +++ b/drizzle-kit/tests/rls/pg-role.test.ts @@ -0,0 +1,234 @@ +import { pgRole } from 'drizzle-orm/pg-core'; +import { diffTestSchemas } from 'tests/schemaDiffer'; +import { expect, test } from 'vitest'; + +test('create role', async (t) => { + const schema1 = {}; + + const schema2 = { + manager: pgRole('manager'), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual(['CREATE ROLE "manager";']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'create_role', + values: { + createDb: false, + createRole: false, + inherit: true, + }, + }, + ]); +}); + +test('create role with properties', async (t) => { + const schema1 = {}; + + const schema2 = { + manager: pgRole('manager', { createDb: true, inherit: false, createRole: true }), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual(['CREATE ROLE "manager" WITH CREATEDB CREATEROLE NOINHERIT;']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'create_role', + values: { + createDb: true, + createRole: true, + inherit: false, + }, + }, + ]); +}); + +test('create role with some properties', async (t) => { + const schema1 = {}; + + const schema2 = { + manager: pgRole('manager', { createDb: true, inherit: false }), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual(['CREATE ROLE "manager" WITH CREATEDB NOINHERIT;']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'create_role', + values: { + createDb: true, + createRole: false, + inherit: false, + }, + }, + ]); +}); + +test('drop role', async (t) => { + const schema1 = { manager: pgRole('manager') }; + + const schema2 = {}; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual(['DROP ROLE "manager";']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'drop_role', + }, + ]); +}); + +test('create and drop role', async (t) => { + const schema1 = { + manager: pgRole('manager'), + }; + + const schema2 = { + admin: pgRole('admin'), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual(['DROP ROLE "manager";', 'CREATE ROLE "admin";']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'drop_role', + }, + { + name: 'admin', + type: 'create_role', + values: { + createDb: false, + createRole: false, + inherit: true, + }, + }, + ]); +}); + +test('rename role', async (t) => { + const schema1 = { + manager: pgRole('manager'), + }; + + const schema2 = { + admin: pgRole('admin'), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, ['manager->admin']); + + expect(sqlStatements).toStrictEqual(['ALTER ROLE "manager" RENAME TO "admin";']); + expect(statements).toStrictEqual([ + { nameFrom: 'manager', nameTo: 'admin', type: 'rename_role' }, + ]); +}); + +test('alter all role field', async (t) => { + const schema1 = { + manager: pgRole('manager'), + }; + + const schema2 = { + manager: pgRole('manager', { createDb: true, createRole: true, inherit: false }), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual(['ALTER ROLE "manager" WITH CREATEDB CREATEROLE NOINHERIT;']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'alter_role', + values: { + createDb: true, + createRole: true, + inherit: false, + }, + }, + ]); +}); + +test('alter createdb in role', async (t) => { + const schema1 = { + manager: pgRole('manager'), + }; + + const schema2 = { + manager: pgRole('manager', { createDb: true }), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual(['ALTER ROLE "manager" WITH CREATEDB NOCREATEROLE INHERIT;']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'alter_role', + values: { + createDb: true, + createRole: false, + inherit: true, + }, + }, + ]); +}); + +test('alter createrole in role', async (t) => { + const schema1 = { + manager: pgRole('manager'), + }; + + const schema2 = { + manager: pgRole('manager', { createRole: true }), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual(['ALTER ROLE "manager" WITH NOCREATEDB CREATEROLE INHERIT;']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'alter_role', + values: { + createDb: false, + createRole: true, + inherit: true, + }, + }, + ]); +}); + +test('alter inherit in role', async (t) => { + const schema1 = { + manager: pgRole('manager'), + }; + + const schema2 = { + manager: pgRole('manager', { inherit: false }), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual(['ALTER ROLE "manager" WITH NOCREATEDB NOCREATEROLE NOINHERIT;']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'alter_role', + values: { + createDb: false, + createRole: false, + inherit: false, + }, + }, + ]); +}); diff --git a/drizzle-kit/tests/schemaDiffer.ts b/drizzle-kit/tests/schemaDiffer.ts index 7b912c6cb..092dc73a8 100644 --- a/drizzle-kit/tests/schemaDiffer.ts +++ b/drizzle-kit/tests/schemaDiffer.ts @@ -11,6 +11,7 @@ import { enumsResolver, Named, policyResolver, + roleResolver, schemasResolver, sequencesResolver, tablesResolver, @@ -25,7 +26,7 @@ import { mysqlSchema, squashMysqlScheme } from 'src/serializer/mysqlSchema'; import { generateMySqlSnapshot } from 'src/serializer/mysqlSerializer'; import { fromDatabase as fromMySqlDatabase } from 'src/serializer/mysqlSerializer'; import { prepareFromPgImports } from 'src/serializer/pgImports'; -import { pgSchema, PgSquasher, Policy, squashPgScheme } from 'src/serializer/pgSchema'; +import { pgSchema, PgSquasher, Policy, Role, squashPgScheme } from 'src/serializer/pgSchema'; import { fromDatabase, generatePgSnapshot } from 'src/serializer/pgSerializer'; import { prepareFromSqliteImports } from 'src/serializer/sqliteImports'; import { sqliteSchema, squashSqliteScheme } from 'src/serializer/sqliteSchema'; @@ -42,6 +43,8 @@ import { ResolverInput, ResolverOutput, ResolverOutputWithMoved, + RolesResolverInput, + RolesResolverOutput, Sequence, Table, } from 'src/snapshotsDiffer'; @@ -404,6 +407,64 @@ async ( } }; +export const testRolesResolver = (renames: Set) => +async ( + input: RolesResolverInput, +): Promise> => { + try { + if ( + input.created.length === 0 + || input.deleted.length === 0 + || renames.size === 0 + ) { + return { + created: input.created, + renamed: [], + deleted: input.deleted, + }; + } + + let createdPolicies = [...input.created]; + let deletedPolicies = [...input.deleted]; + + const renamed: { from: Policy; to: Policy }[] = []; + + for (let rename of renames) { + const [from, to] = rename.split('->'); + + const idxFrom = deletedPolicies.findIndex((it) => { + return `${it.name}` === from; + }); + + if (idxFrom >= 0) { + const idxTo = createdPolicies.findIndex((it) => { + return `${it.name}` === to; + }); + + renamed.push({ + from: deletedPolicies[idxFrom], + to: createdPolicies[idxTo], + }); + + delete createdPolicies[idxTo]; + delete deletedPolicies[idxFrom]; + + createdPolicies = createdPolicies.filter(Boolean); + deletedPolicies = deletedPolicies.filter(Boolean); + } + } + + return { + created: createdPolicies, + deleted: deletedPolicies, + renamed, + }; + } catch (e) { + console.error(e); + throw e; + } +}; + export const testColumnsResolver = (renames: Set) => async ( input: ColumnsResolverInput, @@ -546,6 +607,7 @@ export const diffTestSchemasPush = async ( testEnumsResolver(renames), testSequencesResolver(renames), testPolicyResolver(renames), + testRolesResolver(renames), testTablesResolver(renames), testColumnsResolver(renames), validatedPrev, @@ -561,6 +623,7 @@ export const diffTestSchemasPush = async ( enumsResolver, sequencesResolver, policyResolver, + roleResolver, tablesResolver, columnsResolver, validatedPrev, @@ -581,6 +644,7 @@ export const applyPgDiffs = async (sn: PostgresSchema) => { enums: {}, schemas: {}, sequences: {}, + roles: {}, _meta: { schemas: {}, tables: {}, @@ -622,6 +686,7 @@ export const applyPgDiffs = async (sn: PostgresSchema) => { testEnumsResolver(new Set()), testSequencesResolver(new Set()), testPolicyResolver(new Set()), + testRolesResolver(new Set()), testTablesResolver(new Set()), testColumnsResolver(new Set()), validatedPrev, @@ -706,6 +771,7 @@ export const diffTestSchemas = async ( testEnumsResolver(renames), testSequencesResolver(renames), testPolicyResolver(renames), + testRolesResolver(renames), testTablesResolver(renames), testColumnsResolver(renames), validatedPrev, @@ -720,6 +786,7 @@ export const diffTestSchemas = async ( enumsResolver, sequencesResolver, policyResolver, + roleResolver, tablesResolver, columnsResolver, validatedPrev, @@ -1216,6 +1283,7 @@ export const introspectPgToFile = async ( testEnumsResolver(new Set()), testSequencesResolver(new Set()), testPolicyResolver(new Set()), + testRolesResolver(new Set()), testTablesResolver(new Set()), testColumnsResolver(new Set()), validatedCurAfterImport, diff --git a/drizzle-kit/vitest.config.ts b/drizzle-kit/vitest.config.ts index 23067c3b0..9855269a0 100644 --- a/drizzle-kit/vitest.config.ts +++ b/drizzle-kit/vitest.config.ts @@ -5,6 +5,7 @@ export default defineConfig({ test: { include: [ // 'tests/**/*.test.ts', + 'tests/rls/pg-role.test.ts', 'tests/rls/pg-policy.test.ts', ], From d934de8d34eb88caad276deb711df0170e1e94e9 Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Mon, 2 Sep 2024 18:13:36 +0300 Subject: [PATCH 05/24] Finish all logic for roles and policies in Postgres --- drizzle-kit/package.json | 2 +- drizzle-kit/src/api.ts | 3 + drizzle-kit/src/introspect-pg.ts | 52 +- drizzle-kit/src/serializer/pgSchema.ts | 11 +- drizzle-kit/src/serializer/pgSerializer.ts | 31 +- drizzle-kit/src/snapshotsDiffer.ts | 10 +- drizzle-kit/src/sqlgenerator.ts | 8 +- drizzle-kit/tests/introspect/pg.test.ts | 240 ++++ drizzle-kit/tests/push/pg.test.ts | 1266 ++++++++++++++++++++ drizzle-kit/tests/rls/pg-policy.test.ts | 122 +- drizzle-kit/tests/schemaDiffer.ts | 10 +- drizzle-kit/vitest.config.ts | 4 +- drizzle-orm/package.json | 4 +- drizzle-orm/src/pg-core/policies.ts | 8 +- integration-tests/package.json | 2 +- pnpm-lock.yaml | 68 +- 16 files changed, 1696 insertions(+), 145 deletions(-) diff --git a/drizzle-kit/package.json b/drizzle-kit/package.json index 9d9e1d227..d4650ed26 100644 --- a/drizzle-kit/package.json +++ b/drizzle-kit/package.json @@ -51,7 +51,7 @@ "@arethetypeswrong/cli": "^0.15.3", "@aws-sdk/client-rds-data": "^3.556.0", "@cloudflare/workers-types": "^4.20230518.0", - "@electric-sql/pglite": "^0.1.5", + "@electric-sql/pglite": "^0.2.4", "@hono/node-server": "^1.9.0", "@hono/zod-validator": "^0.2.1", "@libsql/client": "^0.4.2", diff --git a/drizzle-kit/src/api.ts b/drizzle-kit/src/api.ts index 8922fef0e..74659efb0 100644 --- a/drizzle-kit/src/api.ts +++ b/drizzle-kit/src/api.ts @@ -7,6 +7,7 @@ import { columnsResolver, enumsResolver, policyResolver, + roleResolver, schemasResolver, sequencesResolver, tablesResolver, @@ -74,6 +75,7 @@ export const generateMigration = async ( enumsResolver, sequencesResolver, policyResolver, + roleResolver, tablesResolver, columnsResolver, validatedPrev, @@ -118,6 +120,7 @@ export const pushSchema = async ( enumsResolver, sequencesResolver, policyResolver, + roleResolver, tablesResolver, columnsResolver, validatedPrev, diff --git a/drizzle-kit/src/introspect-pg.ts b/drizzle-kit/src/introspect-pg.ts index d9c909bb5..752ea0fba 100644 --- a/drizzle-kit/src/introspect-pg.ts +++ b/drizzle-kit/src/introspect-pg.ts @@ -357,6 +357,10 @@ export const schemaToTypeScript = ( (it) => 'unique', ); + const policiesImports = Object.values(it.policies).map( + (it) => 'pgPolicy', + ); + if (it.schema && it.schema !== 'public' && it.schema !== '') { res.pg.push('pgSchema'); } @@ -365,6 +369,7 @@ export const schemaToTypeScript = ( res.pg.push(...fkImpots); res.pg.push(...pkImports); res.pg.push(...uniqueImports); + res.pg.push(...policiesImports); const columnImports = Object.values(it.columns) .map((col) => { @@ -405,6 +410,10 @@ export const schemaToTypeScript = ( } }); + if (Object.keys(schema.roles).length > 0) { + imports.pg.push('pgRole'); + } + const enumStatements = Object.values(schema.enums) .map((it) => { const enumSchema = schemas[it.schema]; @@ -456,7 +465,7 @@ export const schemaToTypeScript = ( })\n`; }) .join('') - .concat('\n'); + .concat(''); const schemaStatements = Object.entries(schemas) // .filter((it) => it[0] !== "public") @@ -465,16 +474,21 @@ export const schemaToTypeScript = ( }) .join(''); + const rolesNameToTsKey: Record = {}; + const rolesStatements = Object.entries(schema.roles) .map((it) => { const fields = it[1]; - return `export const ${withCasing(it[0], casing)} = pgRole("${fields.name}"${ - !fields.createDb && !fields.createRole && !fields.inherit + rolesNameToTsKey[fields.name] = it[0]; + return `export const ${withCasing(it[0], casing)} = pgRole("${fields.name}", ${ + !fields.createDb && !fields.createRole && fields.inherit ? '' - : `, { ${fields.createDb ? `createDb: true, ` : ''}${fields.createRole ? `createRole: true, ` : ''}${ - fields.inherit ? `inherit: true, ` : '' - } }` - });\n`; + : `${ + `, { ${fields.createDb ? `createDb: true,` : ''}${fields.createRole ? ` createRole: true,` : ''}${ + !fields.inherit ? ` inherit: false ` : '' + }`.trimChar(',') + }}` + } );\n`; }) .join(''); @@ -504,10 +518,11 @@ export const schemaToTypeScript = ( if ( Object.keys(table.indexes).length > 0 || Object.values(table.foreignKeys).length > 0 + || Object.values(table.policies).length > 0 || Object.keys(table.compositePrimaryKeys).length > 0 || Object.keys(table.uniqueConstraints).length > 0 ) { - statement += ',\n'; + statement += ', '; statement += '(table) => {\n'; statement += '\treturn {\n'; statement += createTableIndexes( @@ -527,6 +542,7 @@ export const schemaToTypeScript = ( statement += createTablePolicies( Object.values(table.policies), casing, + rolesNameToTsKey, ); statement += '\t}\n'; statement += '}'; @@ -543,14 +559,14 @@ export const schemaToTypeScript = ( ', ', ) } } from "drizzle-orm/pg-core" - import { sql } from "drizzle-orm"\n\n`; +import { sql } from "drizzle-orm"\n\n`; let decalrations = schemaStatements; decalrations += rolesStatements; decalrations += enumStatements; decalrations += sequencesStatements; - decalrations += '\n'; - decalrations += tableStatements.join('\n\n'); + // decalrations += '\n'; + decalrations += tableStatements.join('\n'); const file = importsTs + decalrations; @@ -1267,22 +1283,30 @@ const createTablePKs = (pks: PrimaryKey[], casing: Casing): string => { return statement; }; +// get a map of db role name to ts key +// if to by key is in this map - no quotes, otherwise - quotes + const createTablePolicies = ( policies: Policy[], casing: Casing, + rolesNameToTsKey: Record = {}, ): string => { let statement = ''; policies.forEach((it) => { const idxKey = withCasing(it.name, casing); + const mappedItTo = it.to?.map((v) => { + return rolesNameToTsKey[v] ? withCasing(rolesNameToTsKey[v], casing) : `"${v}"`; + }); + statement += `\t\t${idxKey}: `; statement += 'pgPolicy('; statement += `"${it.name}", { `; - statement += `as: "${it.as}", for: "${it.for}", to: "${it.to?.join(',')}"${ + statement += `as: "${it.as?.toLowerCase()}", for: "${it.for?.toLowerCase()}", to: [${mappedItTo?.join(', ')}]${ it.using ? `, using: sql\`${it.using}\`` : '' - }${it.withCheck ? `, withCheck: sql\`${it.withCheck}\`` : ''}`; - statement += ` },\n`; + }${it.withCheck ? `, withCheck: sql\`${it.withCheck}\` ` : ''}`; + statement += ` }),\n`; }); return statement; diff --git a/drizzle-kit/src/serializer/pgSchema.ts b/drizzle-kit/src/serializer/pgSchema.ts index 494dc2a78..dd716c2fb 100644 --- a/drizzle-kit/src/serializer/pgSchema.ts +++ b/drizzle-kit/src/serializer/pgSchema.ts @@ -229,8 +229,8 @@ const uniqueConstraint = object({ const policy = object({ name: string(), - as: enumType(['permissive', 'restrictive']).optional(), - for: enumType(['all', 'select', 'insert', 'update', 'delete']).optional(), + as: enumType(['PERMISSIVE', 'RESTRICTIVE']).optional(), + for: enumType(['ALL', 'SELECT', 'INSERT', 'UPDATE', 'DELETE']).optional(), to: string().array().optional(), using: string().optional(), withCheck: string().optional(), @@ -585,6 +585,9 @@ export const PgSquasher = { withCheck: splitted[5] !== 'undefined' ? splitted[5] : undefined, }; }, + squashPolicyPush: (policy: Policy) => { + return `${policy.name}--${policy.as}--${policy.for}--${policy.to?.join(',')}`; + }, squashPK: (pk: PrimaryKey) => { return `${pk.columns.join(',')};${pk.name}`; }, @@ -708,7 +711,9 @@ export const squashPgScheme = ( ); const squashedPolicies = mapValues(it[1].policies, (policy) => { - return PgSquasher.squashPolicy(policy); + return action === 'push' + ? PgSquasher.squashPolicyPush(policy) + : PgSquasher.squashPolicy(policy); }); return [ diff --git a/drizzle-kit/src/serializer/pgSerializer.ts b/drizzle-kit/src/serializer/pgSerializer.ts index 8490ca6aa..7b84b32f4 100644 --- a/drizzle-kit/src/serializer/pgSerializer.ts +++ b/drizzle-kit/src/serializer/pgSerializer.ts @@ -504,7 +504,7 @@ export const generatePgSnapshot = ( const mappedTo = []; if (!policy.to) { - mappedTo.push('PUBLIC'); + mappedTo.push('public'); } else { if (policy.to && typeof policy.to === 'string') { mappedTo.push(policy.to); @@ -523,9 +523,9 @@ export const generatePgSnapshot = ( policiesObject[policy.name] = { name: policy.name, - as: policy.as ?? 'permissive', - for: policy.for ?? 'all', - to: mappedTo, + as: policy.as?.toUpperCase() as Policy['as'] ?? 'PERMISSIVE', + for: policy.for?.toUpperCase() as Policy['for'] ?? 'ALL', + to: mappedTo.sort(), using: is(policy.using, SQL) ? sqlToStr(policy.using) : undefined, withCheck: is(policy.withCheck, SQL) ? sqlToStr(policy.withCheck) : undefined, }; @@ -579,9 +579,9 @@ export const generatePgSnapshot = ( if (!(role as any)._existing) { rolesToReturn[role.name] = { name: role.name, - createDb: (role as any).createDb ?? false, - createRole: (role as any).createRole ?? false, - inherit: (role as any).inherit ?? true, + createDb: (role as any).createDb === undefined ? false : (role as any).createDb, + createRole: (role as any).createRole === undefined ? false : (role as any).createRole, + inherit: (role as any).inherit === undefined ? true : (role as any).inherit, }; } } @@ -838,7 +838,7 @@ export const fromDatabase = async ( rolesToReturn[dbRole.rolname] = { createDb: dbRole.rolcreatedb, - createRole: dbRole.rolcreatedb, + createRole: dbRole.rolcreaterole, inherit: dbRole.rolinherit, name: dbRole.rolname, }; @@ -863,21 +863,26 @@ export const fromDatabase = async ( using: string; withCheck: string; } - >(`SELECT schemaname, tablename, policyname as name, permissive as "as", roles as to, cmd as for, qual as using, "withCheck" FROM pg_policies${wherePolicies};`); + >(`SELECT schemaname, tablename, policyname as name, permissive as "as", roles as to, cmd as for, qual as using, with_check as "withCheck" FROM pg_policies${ + wherePolicies === '' ? '' : ` WHERE ${wherePolicies}` + };`); for (const dbPolicy of allPolicies) { - const { tablename, schemaname, to, ...rest } = dbPolicy; + const { tablename, schemaname, to, withCheck, using, ...rest } = dbPolicy; const tableForPolicy = policiesByTable[`${schemaname}.${tablename}`]; const parsedTo = to === '{}' ? [] - : to.substring(1, to.length - 1).split(/\s*,\s*/g); + : to.substring(1, to.length - 1).split(/\s*,\s*/g).sort(); + + const parsedWithCheck = withCheck === null ? undefined : withCheck; + const parsedUsing = using === null ? undefined : using; if (tableForPolicy) { tableForPolicy[dbPolicy.name] = { ...rest, to: parsedTo } as Policy; } else { policiesByTable[`${schemaname}.${tablename}`] = { - [dbPolicy.name]: { ...rest, to: parsedTo } as Policy, + [dbPolicy.name]: { ...rest, to: parsedTo, withCheck: parsedWithCheck, using: parsedUsing } as Policy, }; } } @@ -1359,7 +1364,7 @@ export const fromDatabase = async ( foreignKeys: foreignKeysToReturn, compositePrimaryKeys: primaryKeys, uniqueConstraints: uniqueConstrains, - policies: policiesByTable[`${tableSchema}.${tableName}`], + policies: policiesByTable[`${tableSchema}.${tableName}`] ?? {}, }; } catch (e) { rej(e); diff --git a/drizzle-kit/src/snapshotsDiffer.ts b/drizzle-kit/src/snapshotsDiffer.ts index cd7ab95dd..ffb6ba4a3 100644 --- a/drizzle-kit/src/snapshotsDiffer.ts +++ b/drizzle-kit/src/snapshotsDiffer.ts @@ -1453,6 +1453,11 @@ export const applyPgSnapshotsDiff = async ( jsonStatements.push(...renameSequences); jsonStatements.push(...jsonAlterSequences); + jsonStatements.push(...renameRoles); + jsonStatements.push(...dropRoles); + jsonStatements.push(...createRoles); + jsonStatements.push(...jsonAlterRoles); + jsonStatements.push(...createTables); jsonStatements.push(...jsonEnableRLSStatements); @@ -1494,11 +1499,6 @@ export const applyPgSnapshotsDiff = async ( jsonStatements.push(...jsonCreatePoliciesStatements); jsonStatements.push(...jsonAlterPoliciesStatements); - jsonStatements.push(...renameRoles); - jsonStatements.push(...dropRoles); - jsonStatements.push(...createRoles); - jsonStatements.push(...jsonAlterRoles); - jsonStatements.push(...dropEnums); jsonStatements.push(...dropSequences); jsonStatements.push(...dropSchemas); diff --git a/drizzle-kit/src/sqlgenerator.ts b/drizzle-kit/src/sqlgenerator.ts index a8d962c4d..d19619e78 100644 --- a/drizzle-kit/src/sqlgenerator.ts +++ b/drizzle-kit/src/sqlgenerator.ts @@ -202,9 +202,11 @@ class PgCreatePolicyConvertor extends Convertor { const withCheckPart = policy.withCheck ? ` WITH CHECK (${policy.withCheck})` : ''; - return `CREATE POLICY "${policy.name}" ON ${tableNameWithSchema} AS ${policy.as?.toUpperCase()} FOR ${policy.for?.toUpperCase()} TO ${ - policy.to?.join(', ') - }${usingPart}${withCheckPart};`; + const policyToPart = policy.to?.map((v) => + ['current_user', 'current_role', 'session_user', 'public'].includes(v) ? v : `"${v}"` + ).join(', '); + + return `CREATE POLICY "${policy.name}" ON ${tableNameWithSchema} AS ${policy.as?.toUpperCase()} FOR ${policy.for?.toUpperCase()} TO ${policyToPart}${usingPart}${withCheckPart};`; } } diff --git a/drizzle-kit/tests/introspect/pg.test.ts b/drizzle-kit/tests/introspect/pg.test.ts index e65c0f904..dd3317364 100644 --- a/drizzle-kit/tests/introspect/pg.test.ts +++ b/drizzle-kit/tests/introspect/pg.test.ts @@ -17,6 +17,8 @@ import { macaddr8, numeric, pgEnum, + pgPolicy, + pgRole, pgSchema, pgTable, real, @@ -405,3 +407,241 @@ test('introspect enum with similar name to native type', async () => { expect(statements.length).toBe(0); expect(sqlStatements.length).toBe(0); }); + +/// + +test('basic policy', async () => { + const client = new PGlite(); + + const schema = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test'), + })), + }; + + const { statements, sqlStatements } = await introspectPgToFile( + client, + schema, + 'basic-policy', + ); + + expect(statements.length).toBe(0); + expect(sqlStatements.length).toBe(0); +}); + +test('basic policy with "as"', async () => { + const client = new PGlite(); + + const schema = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const { statements, sqlStatements } = await introspectPgToFile( + client, + schema, + 'basic-policy-as', + ); + + expect(statements.length).toBe(0); + expect(sqlStatements.length).toBe(0); +}); + +test.todo('basic policy with CURRENT_USER role', async () => { + const client = new PGlite(); + + const schema = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { to: 'current_user' }), + })), + }; + + const { statements, sqlStatements } = await introspectPgToFile( + client, + schema, + 'basic-policy', + ); + + expect(statements.length).toBe(0); + expect(sqlStatements.length).toBe(0); +}); + +test('basic policy with all fields except "using" and "with"', async () => { + const client = new PGlite(); + + const schema = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive', for: 'all', to: ['postgres'] }), + })), + }; + + const { statements, sqlStatements } = await introspectPgToFile( + client, + schema, + 'basic-policy-all-fields', + ); + + expect(statements.length).toBe(0); + expect(sqlStatements.length).toBe(0); +}); + +test('basic policy with "using" and "with"', async () => { + const client = new PGlite(); + + const schema = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { using: sql`true`, withCheck: sql`true` }), + })), + }; + + const { statements, sqlStatements } = await introspectPgToFile( + client, + schema, + 'basic-policy-using-withcheck', + ); + + expect(statements.length).toBe(0); + expect(sqlStatements.length).toBe(0); +}); + +test('multiple policies', async () => { + const client = new PGlite(); + + const schema = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { using: sql`true`, withCheck: sql`true` }), + rlsPolicy: pgPolicy('newRls'), + })), + }; + + const { statements, sqlStatements } = await introspectPgToFile( + client, + schema, + 'multiple-policies', + ); + + expect(statements.length).toBe(0); + expect(sqlStatements.length).toBe(0); +}); + +test('multiple policies with roles', async () => { + const client = new PGlite(); + + client.query(`CREATE ROLE manager;`); + + const schema = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { using: sql`true`, withCheck: sql`true` }), + rlsPolicy: pgPolicy('newRls', { to: ['postgres', 'manager'] }), + })), + }; + + const { statements, sqlStatements } = await introspectPgToFile( + client, + schema, + 'multiple-policies-with-roles', + ); + + expect(statements.length).toBe(0); + expect(sqlStatements.length).toBe(0); +}); + +test('basic roles', async () => { + const client = new PGlite(); + + const schema = { + usersRole: pgRole('user'), + }; + + const { statements, sqlStatements } = await introspectPgToFile( + client, + schema, + 'basic-roles', + ['public'], + { roles: { include: ['user'] } }, + ); + + expect(statements.length).toBe(0); + expect(sqlStatements.length).toBe(0); +}); + +test('role with properties', async () => { + const client = new PGlite(); + + const schema = { + usersRole: pgRole('user', { inherit: false, createDb: true, createRole: true }), + }; + + const { statements, sqlStatements } = await introspectPgToFile( + client, + schema, + 'roles-with-properties', + ['public'], + { roles: { include: ['user'] } }, + ); + + expect(statements.length).toBe(0); + expect(sqlStatements.length).toBe(0); +}); + +test('role with a few properties', async () => { + const client = new PGlite(); + + const schema = { + usersRole: pgRole('user', { inherit: false, createRole: true }), + }; + + const { statements, sqlStatements } = await introspectPgToFile( + client, + schema, + 'roles-with-few-properties', + ['public'], + { roles: { include: ['user'] } }, + ); + + expect(statements.length).toBe(0); + expect(sqlStatements.length).toBe(0); +}); + +test('multiple policies with roles from schema', async () => { + const client = new PGlite(); + + const usersRole = pgRole('user_role', { inherit: false, createRole: true }); + + const schema = { + usersRole, + + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { using: sql`true`, withCheck: sql`true` }), + rlsPolicy: pgPolicy('newRls', { to: ['postgres', usersRole] }), + })), + }; + + const { statements, sqlStatements } = await introspectPgToFile( + client, + schema, + 'multiple-policies-with-roles-from-schema', + ['public'], + { roles: { include: ['user_role'] } }, + ); + + expect(statements.length).toBe(0); + expect(sqlStatements.length).toBe(0); +}); diff --git a/drizzle-kit/tests/push/pg.test.ts b/drizzle-kit/tests/push/pg.test.ts index d6d1212e6..88f02e8ae 100644 --- a/drizzle-kit/tests/push/pg.test.ts +++ b/drizzle-kit/tests/push/pg.test.ts @@ -13,6 +13,8 @@ import { jsonb, numeric, pgEnum, + pgPolicy, + pgRole, pgSchema, pgSequence, pgTable, @@ -2240,3 +2242,1267 @@ test('add array column - default', async () => { 'ALTER TABLE "test" ADD COLUMN "values" integer[] DEFAULT \'{1,2,3}\';', ]); }); + +// Policies and Roles push test +test('full policy: no changes', async () => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(statements.length).toBe(0); + expect(sqlStatements.length).toBe(0); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('add policy', async () => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(statements).toStrictEqual([ + { type: 'enable_rls', tableName: 'users', schema: '' }, + { + type: 'create_policy', + tableName: 'users', + data: { + name: 'test', + as: 'PERMISSIVE', + for: 'ALL', + to: ['public'], + using: undefined, + withCheck: undefined, + }, + schema: '', + }, + ]); + expect(sqlStatements).toStrictEqual([ + 'ALTER TABLE "users" ENABLE ROW LEVEL SECURITY;', + 'CREATE POLICY "test" ON "users" AS PERMISSIVE FOR ALL TO public;', + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('drop policy', async () => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(statements).toStrictEqual([ + { type: 'disable_rls', tableName: 'users', schema: '' }, + { + type: 'drop_policy', + tableName: 'users', + data: { + name: 'test', + as: 'PERMISSIVE', + for: 'ALL', + to: ['public'], + using: undefined, + withCheck: undefined, + }, + schema: '', + }, + ]); + expect(sqlStatements).toStrictEqual([ + 'ALTER TABLE "users" DISABLE ROW LEVEL SECURITY;', + 'DROP POLICY "test" ON "users" CASCADE;', + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('add policy without enable rls', async () => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + newrls: pgPolicy('newRls'), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(statements).toStrictEqual([ + { + type: 'create_policy', + tableName: 'users', + data: { + name: 'newRls', + as: 'PERMISSIVE', + for: 'ALL', + to: ['public'], + using: undefined, + withCheck: undefined, + }, + schema: '', + }, + ]); + expect(sqlStatements).toStrictEqual([ + 'CREATE POLICY "newRls" ON "users" AS PERMISSIVE FOR ALL TO public;', + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('drop policy without disable rls', async () => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + oldRls: pgPolicy('oldRls'), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(statements).toStrictEqual([ + { + type: 'drop_policy', + tableName: 'users', + data: { + name: 'oldRls', + as: 'PERMISSIVE', + for: 'ALL', + to: ['public'], + using: undefined, + withCheck: undefined, + }, + schema: '', + }, + ]); + expect(sqlStatements).toStrictEqual([ + 'DROP POLICY "oldRls" ON "users" CASCADE;', + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +//// + +test('alter policy without recreation: changing roles', async (t) => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive', to: 'current_role' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(sqlStatements).toStrictEqual([ + 'ALTER POLICY "test" ON "users" TO current_role;', + ]); + expect(statements).toStrictEqual([ + { + newData: 'test--PERMISSIVE--ALL--current_role', + oldData: 'test--PERMISSIVE--ALL--public', + schema: '', + tableName: 'users', + type: 'alter_policy', + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('alter policy without recreation: changing using', async (t) => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive', using: sql`true` }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(sqlStatements).toStrictEqual([]); + expect(statements).toStrictEqual([]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('alter policy without recreation: changing with check', async (t) => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive', withCheck: sql`true` }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(sqlStatements).toStrictEqual([]); + expect(statements).toStrictEqual([]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('alter policy with recreation: changing as', async (t) => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'restrictive' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(sqlStatements).toStrictEqual([ + 'DROP POLICY "test" ON "users" CASCADE;', + 'CREATE POLICY "test" ON "users" AS RESTRICTIVE FOR ALL TO public;', + ]); + expect(statements).toStrictEqual([ + { + data: { + as: 'PERMISSIVE', + for: 'ALL', + name: 'test', + to: ['public'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'drop_policy', + }, + { + data: { + as: 'RESTRICTIVE', + for: 'ALL', + name: 'test', + to: ['public'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'create_policy', + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('alter policy with recreation: changing for', async (t) => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive', for: 'delete' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(sqlStatements).toStrictEqual([ + 'DROP POLICY "test" ON "users" CASCADE;', + 'CREATE POLICY "test" ON "users" AS PERMISSIVE FOR DELETE TO public;', + ]); + expect(statements).toStrictEqual([ + { + data: { + as: 'PERMISSIVE', + for: 'ALL', + name: 'test', + to: ['public'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'drop_policy', + }, + { + data: { + as: 'PERMISSIVE', + for: 'DELETE', + name: 'test', + to: ['public'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'create_policy', + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('alter policy with recreation: changing both "as" and "for"', async (t) => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'restrictive', for: 'insert' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(sqlStatements).toStrictEqual([ + 'DROP POLICY "test" ON "users" CASCADE;', + 'CREATE POLICY "test" ON "users" AS RESTRICTIVE FOR INSERT TO public;', + ]); + expect(statements).toStrictEqual([ + { + data: { + as: 'PERMISSIVE', + for: 'ALL', + name: 'test', + to: ['public'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'drop_policy', + }, + { + data: { + as: 'RESTRICTIVE', + for: 'INSERT', + name: 'test', + to: ['public'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'create_policy', + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('alter policy with recreation: changing all fields', async (t) => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive', for: 'select', using: sql`true` }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'restrictive', to: 'current_role', withCheck: sql`true` }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(sqlStatements).toStrictEqual([ + 'DROP POLICY "test" ON "users" CASCADE;', + 'CREATE POLICY "test" ON "users" AS RESTRICTIVE FOR ALL TO current_role;', + ]); + expect(statements).toStrictEqual([ + { + data: { + as: 'PERMISSIVE', + for: 'SELECT', + name: 'test', + to: ['public'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'drop_policy', + }, + { + data: { + as: 'RESTRICTIVE', + for: 'ALL', + name: 'test', + to: ['current_role'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'create_policy', + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('rename policy', async (t) => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('newName', { as: 'permissive' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + ['public.users.test->public.users.newName'], + false, + ['public'], + ); + + expect(sqlStatements).toStrictEqual([ + 'ALTER POLICY "test" ON "users" RENAME TO "newName";', + ]); + expect(statements).toStrictEqual([ + { + newName: 'newName', + oldName: 'test', + schema: '', + tableName: 'users', + type: 'rename_policy', + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('rename policy in renamed table', async (t) => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = { + users: pgTable('users2', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('newName', { as: 'permissive' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [ + 'public.users->public.users2', + 'public.users2.test->public.users2.newName', + ], + false, + ['public'], + ); + + expect(sqlStatements).toStrictEqual([ + 'ALTER TABLE "users" RENAME TO "users2";', + 'ALTER POLICY "test" ON "users2" RENAME TO "newName";', + ]); + expect(statements).toStrictEqual([ + { + fromSchema: '', + tableNameFrom: 'users', + tableNameTo: 'users2', + toSchema: '', + type: 'rename_table', + }, + { + newName: 'newName', + oldName: 'test', + schema: '', + tableName: 'users2', + type: 'rename_policy', + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('create table with a policy', async (t) => { + const client = new PGlite(); + + const schema1 = {}; + + const schema2 = { + users: pgTable('users2', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(sqlStatements).toStrictEqual([ + 'CREATE TABLE IF NOT EXISTS "users2" (\n\t"id" integer PRIMARY KEY NOT NULL\n);\n', + 'ALTER TABLE "users2" ENABLE ROW LEVEL SECURITY;', + 'CREATE POLICY "test" ON "users2" AS PERMISSIVE FOR ALL TO public;', + ]); + expect(statements).toStrictEqual([ + { + columns: [ + { + name: 'id', + notNull: true, + primaryKey: true, + type: 'integer', + }, + ], + compositePKs: [], + compositePkName: '', + policies: [ + 'test--PERMISSIVE--ALL--public', + ], + schema: '', + tableName: 'users2', + type: 'create_table', + uniqueConstraints: [], + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('drop table with a policy', async (t) => { + const client = new PGlite(); + + const schema1 = { + users: pgTable('users2', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { as: 'permissive' }), + })), + }; + + const schema2 = {}; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(sqlStatements).toStrictEqual([ + 'DROP POLICY "test" ON "users2" CASCADE;', + 'DROP TABLE "users2";', + ]); + expect(statements).toStrictEqual([ + { + policies: [ + 'test--PERMISSIVE--ALL--public', + ], + schema: '', + tableName: 'users2', + type: 'drop_table', + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('add policy with multiple "to" roles', async (t) => { + const client = new PGlite(); + + client.query(`CREATE ROLE manager;`); + + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }), + }; + + const role = pgRole('manager').existing(); + + const schema2 = { + role, + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { to: ['current_role', role] }), + })), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + ); + + expect(sqlStatements).toStrictEqual([ + 'ALTER TABLE "users" ENABLE ROW LEVEL SECURITY;', + 'CREATE POLICY "test" ON "users" AS PERMISSIVE FOR ALL TO current_role, "manager";', + ]); + expect(statements).toStrictEqual([ + { + schema: '', + tableName: 'users', + type: 'enable_rls', + }, + { + data: { + as: 'PERMISSIVE', + for: 'ALL', + name: 'test', + to: ['current_role', 'manager'], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users', + type: 'create_policy', + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +//// + +test('create role', async (t) => { + const client = new PGlite(); + + const schema1 = {}; + + const schema2 = { + manager: pgRole('manager'), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + { roles: { include: ['manager'] } }, + ); + + expect(sqlStatements).toStrictEqual(['CREATE ROLE "manager";']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'create_role', + values: { + createDb: false, + createRole: false, + inherit: true, + }, + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('create role with properties', async (t) => { + const client = new PGlite(); + + const schema1 = {}; + + const schema2 = { + manager: pgRole('manager', { createDb: true, inherit: false, createRole: true }), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + { roles: { include: ['manager'] } }, + ); + + expect(sqlStatements).toStrictEqual(['CREATE ROLE "manager" WITH CREATEDB CREATEROLE NOINHERIT;']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'create_role', + values: { + createDb: true, + createRole: true, + inherit: false, + }, + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('create role with some properties', async (t) => { + const client = new PGlite(); + + const schema1 = {}; + + const schema2 = { + manager: pgRole('manager', { createDb: true, inherit: false }), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + { roles: { include: ['manager'] } }, + ); + + expect(sqlStatements).toStrictEqual(['CREATE ROLE "manager" WITH CREATEDB NOINHERIT;']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'create_role', + values: { + createDb: true, + createRole: false, + inherit: false, + }, + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('drop role', async (t) => { + const client = new PGlite(); + + const schema1 = { manager: pgRole('manager') }; + + const schema2 = {}; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + { roles: { include: ['manager'] } }, + ); + + expect(sqlStatements).toStrictEqual(['DROP ROLE "manager";']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'drop_role', + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('create and drop role', async (t) => { + const client = new PGlite(); + + const schema1 = { + manager: pgRole('manager'), + }; + + const schema2 = { + admin: pgRole('admin'), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + { roles: { include: ['manager', 'admin'] } }, + ); + + expect(sqlStatements).toStrictEqual(['DROP ROLE "manager";', 'CREATE ROLE "admin";']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'drop_role', + }, + { + name: 'admin', + type: 'create_role', + values: { + createDb: false, + createRole: false, + inherit: true, + }, + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('rename role', async (t) => { + const client = new PGlite(); + + const schema1 = { + manager: pgRole('manager'), + }; + + const schema2 = { + admin: pgRole('admin'), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + ['manager->admin'], + false, + ['public'], + { roles: { include: ['manager', 'admin'] } }, + ); + + expect(sqlStatements).toStrictEqual(['ALTER ROLE "manager" RENAME TO "admin";']); + expect(statements).toStrictEqual([ + { nameFrom: 'manager', nameTo: 'admin', type: 'rename_role' }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('alter all role field', async (t) => { + const client = new PGlite(); + + const schema1 = { + manager: pgRole('manager'), + }; + + const schema2 = { + manager: pgRole('manager', { createDb: true, createRole: true, inherit: false }), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + { roles: { include: ['manager'] } }, + ); + + expect(sqlStatements).toStrictEqual(['ALTER ROLE "manager" WITH CREATEDB CREATEROLE NOINHERIT;']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'alter_role', + values: { + createDb: true, + createRole: true, + inherit: false, + }, + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('alter createdb in role', async (t) => { + const client = new PGlite(); + + const schema1 = { + manager: pgRole('manager'), + }; + + const schema2 = { + manager: pgRole('manager', { createDb: true }), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + { roles: { include: ['manager'] } }, + ); + + expect(sqlStatements).toStrictEqual(['ALTER ROLE "manager" WITH CREATEDB NOCREATEROLE INHERIT;']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'alter_role', + values: { + createDb: true, + createRole: false, + inherit: true, + }, + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('alter createrole in role', async (t) => { + const client = new PGlite(); + + const schema1 = { + manager: pgRole('manager'), + }; + + const schema2 = { + manager: pgRole('manager', { createRole: true }), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + { roles: { include: ['manager'] } }, + ); + + expect(sqlStatements).toStrictEqual(['ALTER ROLE "manager" WITH NOCREATEDB CREATEROLE INHERIT;']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'alter_role', + values: { + createDb: false, + createRole: true, + inherit: true, + }, + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); + +test('alter inherit in role', async (t) => { + const client = new PGlite(); + + const schema1 = { + manager: pgRole('manager'), + }; + + const schema2 = { + manager: pgRole('manager', { inherit: false }), + }; + + const { statements, sqlStatements } = await diffTestSchemasPush( + client, + schema1, + schema2, + [], + false, + ['public'], + { roles: { include: ['manager'] } }, + ); + + expect(sqlStatements).toStrictEqual(['ALTER ROLE "manager" WITH NOCREATEDB NOCREATEROLE NOINHERIT;']); + expect(statements).toStrictEqual([ + { + name: 'manager', + type: 'alter_role', + values: { + createDb: false, + createRole: false, + inherit: false, + }, + }, + ]); + + for (const st of sqlStatements) { + await client.query(st); + } +}); diff --git a/drizzle-kit/tests/rls/pg-policy.test.ts b/drizzle-kit/tests/rls/pg-policy.test.ts index 3f12b3e64..17827f951 100644 --- a/drizzle-kit/tests/rls/pg-policy.test.ts +++ b/drizzle-kit/tests/rls/pg-policy.test.ts @@ -22,7 +22,7 @@ test('add policy + enable rls', async (t) => { expect(sqlStatements).toStrictEqual([ 'ALTER TABLE "users" ENABLE ROW LEVEL SECURITY;', - 'CREATE POLICY "test" ON "users" AS PERMISSIVE FOR ALL TO PUBLIC;', + 'CREATE POLICY "test" ON "users" AS PERMISSIVE FOR ALL TO public;', ]); expect(statements).toStrictEqual([ { @@ -32,10 +32,10 @@ test('add policy + enable rls', async (t) => { }, { data: { - as: 'permissive', - for: 'all', + as: 'PERMISSIVE', + for: 'ALL', name: 'test', - to: ['PUBLIC'], + to: ['public'], using: undefined, withCheck: undefined, }, @@ -75,10 +75,10 @@ test('drop policy + disable rls', async (t) => { }, { data: { - as: 'permissive', - for: 'all', + as: 'PERMISSIVE', + for: 'ALL', name: 'test', - to: ['PUBLIC'], + to: ['public'], using: undefined, withCheck: undefined, }, @@ -110,15 +110,15 @@ test('add policy without enable rls', async (t) => { const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); expect(sqlStatements).toStrictEqual([ - 'CREATE POLICY "newRls" ON "users" AS PERMISSIVE FOR ALL TO PUBLIC;', + 'CREATE POLICY "newRls" ON "users" AS PERMISSIVE FOR ALL TO public;', ]); expect(statements).toStrictEqual([ { data: { - as: 'permissive', - for: 'all', + as: 'PERMISSIVE', + for: 'ALL', name: 'newRls', - to: ['PUBLIC'], + to: ['public'], using: undefined, withCheck: undefined, }, @@ -155,10 +155,10 @@ test('drop policy without disable rls', async (t) => { expect(statements).toStrictEqual([ { data: { - as: 'permissive', - for: 'all', + as: 'PERMISSIVE', + for: 'ALL', name: 'oldRls', - to: ['PUBLIC'], + to: ['public'], using: undefined, withCheck: undefined, }, @@ -182,19 +182,19 @@ test('alter policy without recreation: changing roles', async (t) => { users: pgTable('users', { id: integer('id').primaryKey(), }, () => ({ - rls: pgPolicy('test', { as: 'permissive', to: 'CURRENT_ROLE' }), + rls: pgPolicy('test', { as: 'permissive', to: 'current_role' }), })), }; const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); expect(sqlStatements).toStrictEqual([ - 'ALTER POLICY "test" ON "users" TO CURRENT_ROLE;', + 'ALTER POLICY "test" ON "users" TO current_role;', ]); expect(statements).toStrictEqual([ { - newData: 'test--permissive--all--CURRENT_ROLE--undefined--undefined', - oldData: 'test--permissive--all--PUBLIC--undefined--undefined', + newData: 'test--PERMISSIVE--ALL--current_role--undefined--undefined', + oldData: 'test--PERMISSIVE--ALL--public--undefined--undefined', schema: '', tableName: 'users', type: 'alter_policy', @@ -222,12 +222,12 @@ test('alter policy without recreation: changing using', async (t) => { const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); expect(sqlStatements).toStrictEqual([ - 'ALTER POLICY "test" ON "users" TO PUBLIC USING (true);', + 'ALTER POLICY "test" ON "users" TO public USING (true);', ]); expect(statements).toStrictEqual([ { - newData: 'test--permissive--all--PUBLIC--true--undefined', - oldData: 'test--permissive--all--PUBLIC--undefined--undefined', + newData: 'test--PERMISSIVE--ALL--public--true--undefined', + oldData: 'test--PERMISSIVE--ALL--public--undefined--undefined', schema: '', tableName: 'users', type: 'alter_policy', @@ -255,12 +255,12 @@ test('alter policy without recreation: changing with check', async (t) => { const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); expect(sqlStatements).toStrictEqual([ - 'ALTER POLICY "test" ON "users" TO PUBLIC WITH CHECK (true);', + 'ALTER POLICY "test" ON "users" TO public WITH CHECK (true);', ]); expect(statements).toStrictEqual([ { - newData: 'test--permissive--all--PUBLIC--undefined--true', - oldData: 'test--permissive--all--PUBLIC--undefined--undefined', + newData: 'test--PERMISSIVE--ALL--public--undefined--true', + oldData: 'test--PERMISSIVE--ALL--public--undefined--undefined', schema: '', tableName: 'users', type: 'alter_policy', @@ -291,15 +291,15 @@ test('alter policy with recreation: changing as', async (t) => { expect(sqlStatements).toStrictEqual([ 'DROP POLICY "test" ON "users" CASCADE;', - 'CREATE POLICY "test" ON "users" AS RESTRICTIVE FOR ALL TO PUBLIC;', + 'CREATE POLICY "test" ON "users" AS RESTRICTIVE FOR ALL TO public;', ]); expect(statements).toStrictEqual([ { data: { - as: 'permissive', - for: 'all', + as: 'PERMISSIVE', + for: 'ALL', name: 'test', - to: ['PUBLIC'], + to: ['public'], using: undefined, withCheck: undefined, }, @@ -309,10 +309,10 @@ test('alter policy with recreation: changing as', async (t) => { }, { data: { - as: 'restrictive', - for: 'all', + as: 'RESTRICTIVE', + for: 'ALL', name: 'test', - to: ['PUBLIC'], + to: ['public'], using: undefined, withCheck: undefined, }, @@ -344,15 +344,15 @@ test('alter policy with recreation: changing for', async (t) => { expect(sqlStatements).toStrictEqual([ 'DROP POLICY "test" ON "users" CASCADE;', - 'CREATE POLICY "test" ON "users" AS PERMISSIVE FOR DELETE TO PUBLIC;', + 'CREATE POLICY "test" ON "users" AS PERMISSIVE FOR DELETE TO public;', ]); expect(statements).toStrictEqual([ { data: { - as: 'permissive', - for: 'all', + as: 'PERMISSIVE', + for: 'ALL', name: 'test', - to: ['PUBLIC'], + to: ['public'], using: undefined, withCheck: undefined, }, @@ -362,10 +362,10 @@ test('alter policy with recreation: changing for', async (t) => { }, { data: { - as: 'permissive', - for: 'delete', + as: 'PERMISSIVE', + for: 'DELETE', name: 'test', - to: ['PUBLIC'], + to: ['public'], using: undefined, withCheck: undefined, }, @@ -397,15 +397,15 @@ test('alter policy with recreation: changing both "as" and "for"', async (t) => expect(sqlStatements).toStrictEqual([ 'DROP POLICY "test" ON "users" CASCADE;', - 'CREATE POLICY "test" ON "users" AS RESTRICTIVE FOR INSERT TO PUBLIC;', + 'CREATE POLICY "test" ON "users" AS RESTRICTIVE FOR INSERT TO public;', ]); expect(statements).toStrictEqual([ { data: { - as: 'permissive', - for: 'all', + as: 'PERMISSIVE', + for: 'ALL', name: 'test', - to: ['PUBLIC'], + to: ['public'], using: undefined, withCheck: undefined, }, @@ -415,10 +415,10 @@ test('alter policy with recreation: changing both "as" and "for"', async (t) => }, { data: { - as: 'restrictive', - for: 'insert', + as: 'RESTRICTIVE', + for: 'INSERT', name: 'test', - to: ['PUBLIC'], + to: ['public'], using: undefined, withCheck: undefined, }, @@ -442,7 +442,7 @@ test('alter policy with recreation: changing all fields', async (t) => { users: pgTable('users', { id: integer('id').primaryKey(), }, () => ({ - rls: pgPolicy('test', { as: 'restrictive', to: 'CURRENT_ROLE', withCheck: sql`true` }), + rls: pgPolicy('test', { as: 'restrictive', to: 'current_role', withCheck: sql`true` }), })), }; @@ -450,15 +450,15 @@ test('alter policy with recreation: changing all fields', async (t) => { expect(sqlStatements).toStrictEqual([ 'DROP POLICY "test" ON "users" CASCADE;', - 'CREATE POLICY "test" ON "users" AS RESTRICTIVE FOR ALL TO CURRENT_ROLE WITH CHECK (true);', + 'CREATE POLICY "test" ON "users" AS RESTRICTIVE FOR ALL TO current_role WITH CHECK (true);', ]); expect(statements).toStrictEqual([ { data: { - as: 'permissive', - for: 'select', + as: 'PERMISSIVE', + for: 'SELECT', name: 'test', - to: ['PUBLIC'], + to: ['public'], using: 'true', withCheck: undefined, }, @@ -468,10 +468,10 @@ test('alter policy with recreation: changing all fields', async (t) => { }, { data: { - as: 'restrictive', - for: 'all', + as: 'RESTRICTIVE', + for: 'ALL', name: 'test', - to: ['CURRENT_ROLE'], + to: ['current_role'], using: undefined, withCheck: 'true', }, @@ -577,7 +577,7 @@ test('create table with a policy', async (t) => { expect(sqlStatements).toStrictEqual([ 'CREATE TABLE IF NOT EXISTS "users2" (\n\t"id" integer PRIMARY KEY NOT NULL\n);\n', 'ALTER TABLE "users2" ENABLE ROW LEVEL SECURITY;', - 'CREATE POLICY "test" ON "users2" AS PERMISSIVE FOR ALL TO PUBLIC;', + 'CREATE POLICY "test" ON "users2" AS PERMISSIVE FOR ALL TO public;', ]); expect(statements).toStrictEqual([ { @@ -592,7 +592,7 @@ test('create table with a policy', async (t) => { compositePKs: [], compositePkName: '', policies: [ - 'test--permissive--all--PUBLIC--undefined--undefined', + 'test--PERMISSIVE--ALL--public--undefined--undefined', ], schema: '', tableName: 'users2', @@ -622,7 +622,7 @@ test('drop table with a policy', async (t) => { expect(statements).toStrictEqual([ { policies: [ - 'test--permissive--all--PUBLIC--undefined--undefined', + 'test--PERMISSIVE--ALL--public--undefined--undefined', ], schema: '', tableName: 'users2', @@ -645,7 +645,7 @@ test('add policy with multiple "to" roles', async (t) => { users: pgTable('users', { id: integer('id').primaryKey(), }, () => ({ - rls: pgPolicy('test', { to: ['CURRENT_ROLE', role] }), + rls: pgPolicy('test', { to: ['current_role', role] }), })), }; @@ -653,7 +653,7 @@ test('add policy with multiple "to" roles', async (t) => { expect(sqlStatements).toStrictEqual([ 'ALTER TABLE "users" ENABLE ROW LEVEL SECURITY;', - 'CREATE POLICY "test" ON "users" AS PERMISSIVE FOR ALL TO CURRENT_ROLE, manager;', + 'CREATE POLICY "test" ON "users" AS PERMISSIVE FOR ALL TO current_role, "manager";', ]); expect(statements).toStrictEqual([ { @@ -663,10 +663,10 @@ test('add policy with multiple "to" roles', async (t) => { }, { data: { - as: 'permissive', - for: 'all', + as: 'PERMISSIVE', + for: 'ALL', name: 'test', - to: ['CURRENT_ROLE', 'manager'], + to: ['current_role', 'manager'], using: undefined, withCheck: undefined, }, diff --git a/drizzle-kit/tests/schemaDiffer.ts b/drizzle-kit/tests/schemaDiffer.ts index 092dc73a8..52504578b 100644 --- a/drizzle-kit/tests/schemaDiffer.ts +++ b/drizzle-kit/tests/schemaDiffer.ts @@ -536,6 +536,13 @@ export const diffTestSchemasPush = async ( renamesArr: string[], cli: boolean = false, schemas: string[] = ['public'], + entities?: { + roles: boolean | { + provider?: string | undefined; + include?: string[] | undefined; + exclude?: string[] | undefined; + }; + }, ) => { const { sqlStatements } = await applyPgDiffs(left); for (const st of sqlStatements) { @@ -552,6 +559,7 @@ export const diffTestSchemasPush = async ( }, undefined, schemas, + entities, ); const leftTables = Object.values(right).filter((it) => is(it, PgTable)) as PgTable[]; @@ -1192,7 +1200,7 @@ export const introspectPgToFile = async ( initSchema: PostgresSchema, testName: string, schemas: string[] = ['public'], - entities: Entities, + entities?: Entities, ) => { // put in db const { sqlStatements } = await applyPgDiffs(initSchema); diff --git a/drizzle-kit/vitest.config.ts b/drizzle-kit/vitest.config.ts index 9855269a0..602e96ede 100644 --- a/drizzle-kit/vitest.config.ts +++ b/drizzle-kit/vitest.config.ts @@ -4,9 +4,7 @@ import { defineConfig } from 'vitest/config'; export default defineConfig({ test: { include: [ - // 'tests/**/*.test.ts', - 'tests/rls/pg-role.test.ts', - 'tests/rls/pg-policy.test.ts', + 'tests/**/*.test.ts', ], typecheck: { diff --git a/drizzle-orm/package.json b/drizzle-orm/package.json index 888f7efcb..ac889df97 100644 --- a/drizzle-orm/package.json +++ b/drizzle-orm/package.json @@ -45,7 +45,7 @@ "peerDependencies": { "@aws-sdk/client-rds-data": ">=3", "@cloudflare/workers-types": ">=3", - "@electric-sql/pglite": ">=0.1.1", + "@electric-sql/pglite": ">=0.2.0", "@libsql/client": "*", "@neondatabase/serverless": ">=0.1", "@op-engineering/op-sqlite": ">=2", @@ -160,7 +160,7 @@ "devDependencies": { "@aws-sdk/client-rds-data": "^3.549.0", "@cloudflare/workers-types": "^4.20230904.0", - "@electric-sql/pglite": "^0.1.1", + "@electric-sql/pglite": "^0.2.4", "@libsql/client": "^0.5.6", "@neondatabase/serverless": "^0.9.0", "@op-engineering/op-sqlite": "^2.0.16", diff --git a/drizzle-orm/src/pg-core/policies.ts b/drizzle-orm/src/pg-core/policies.ts index 11e898eda..a61731578 100644 --- a/drizzle-orm/src/pg-core/policies.ts +++ b/drizzle-orm/src/pg-core/policies.ts @@ -3,10 +3,10 @@ import type { SQL } from '~/sql/sql.ts'; import type { PgRole } from './roles'; export type PgPolicyToOption = - | 'PUBLIC' - | 'CURRENT_ROLE' - | 'CURRENT_USER' - | 'SESSION_USER' + | 'public' + | 'current_role' + | 'current_user' + | 'session_user' | (string & {}) | PgPolicyToOption[] | PgRole; diff --git a/integration-tests/package.json b/integration-tests/package.json index a4fcab0b2..cd5e6d482 100644 --- a/integration-tests/package.json +++ b/integration-tests/package.json @@ -40,7 +40,7 @@ "dependencies": { "@aws-sdk/client-rds-data": "^3.549.0", "@aws-sdk/credential-providers": "^3.549.0", - "@electric-sql/pglite": "^0.1.1", + "@electric-sql/pglite": "^0.2.4", "@libsql/client": "^0.5.6", "@miniflare/d1": "^2.14.2", "@miniflare/shared": "^2.14.2", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index d2d091ad6..7d6120238 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -114,8 +114,8 @@ importers: specifier: ^4.20230518.0 version: 4.20240524.0 '@electric-sql/pglite': - specifier: ^0.1.5 - version: 0.1.5 + specifier: ^0.2.4 + version: 0.2.4 '@hono/node-server': specifier: ^1.9.0 version: 1.12.0 @@ -300,8 +300,8 @@ importers: specifier: ^4.20230904.0 version: 4.20240512.0 '@electric-sql/pglite': - specifier: ^0.1.1 - version: 0.1.5 + specifier: ^0.2.4 + version: 0.2.4 '@libsql/client': specifier: ^0.5.6 version: 0.5.6(bufferutil@4.0.8)(encoding@0.1.13)(utf-8-validate@6.0.3) @@ -547,10 +547,10 @@ importers: version: 3.583.0 '@aws-sdk/credential-providers': specifier: ^3.549.0 - version: 3.569.0(@aws-sdk/client-sso-oidc@3.583.0) + version: 3.569.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0)) '@electric-sql/pglite': - specifier: ^0.1.1 - version: 0.1.5 + specifier: ^0.2.4 + version: 0.2.4 '@libsql/client': specifier: ^0.5.6 version: 0.5.6(bufferutil@4.0.8)(encoding@0.1.13)(utf-8-validate@6.0.3) @@ -1980,8 +1980,8 @@ packages: '@drizzle-team/studio@0.0.5': resolution: {integrity: sha512-ps5qF0tMxWRVu+V5gvCRrQNqlY92aTnIKdq27gm9LZMSdaKYZt6AVvSK1dlUMzs6Rt0Jm80b+eWct6xShBKhIw==} - '@electric-sql/pglite@0.1.5': - resolution: {integrity: sha512-eymv4ONNvoPZQTvOQIi5dbpR+J5HzEv0qQH9o/y3gvNheJV/P/NFcrbsfJZYTsDKoq7DKrTiFNexsRkJKy8x9Q==} + '@electric-sql/pglite@0.2.4': + resolution: {integrity: sha512-NN1ATH9aYTCD4257wZH1CjAKdro8jDd5r/BumxZtEVTDNDztBmPYYtBd4TEIHWdi0+QuM7SWk2JRCLvAYCSEWg==} '@esbuild-kit/core-utils@3.1.0': resolution: {integrity: sha512-Uuk8RpCg/7fdHSceR1M6XbSZFSuMrxcePFuGgyvsBn+u339dk5OeL4jv2EojwTN2st/unJGsVm4qHWjWNmJ/tw==} @@ -10210,7 +10210,7 @@ snapshots: '@aws-sdk/client-sso-oidc': 3.583.0(@aws-sdk/client-sts@3.583.0) '@aws-sdk/client-sts': 3.583.0 '@aws-sdk/core': 3.582.0 - '@aws-sdk/credential-provider-node': 3.583.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))(@aws-sdk/client-sts@3.583.0) + '@aws-sdk/credential-provider-node': 3.583.0(@aws-sdk/client-sso-oidc@3.583.0)(@aws-sdk/client-sts@3.583.0) '@aws-sdk/middleware-host-header': 3.577.0 '@aws-sdk/middleware-logger': 3.577.0 '@aws-sdk/middleware-recursion-detection': 3.577.0 @@ -10300,7 +10300,7 @@ snapshots: '@aws-crypto/sha256-js': 3.0.0 '@aws-sdk/client-sts': 3.583.0 '@aws-sdk/core': 3.582.0 - '@aws-sdk/credential-provider-node': 3.583.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))(@aws-sdk/client-sts@3.583.0) + '@aws-sdk/credential-provider-node': 3.583.0(@aws-sdk/client-sso-oidc@3.583.0)(@aws-sdk/client-sts@3.583.0) '@aws-sdk/middleware-host-header': 3.577.0 '@aws-sdk/middleware-logger': 3.577.0 '@aws-sdk/middleware-recursion-detection': 3.577.0 @@ -10610,7 +10610,7 @@ snapshots: '@aws-crypto/sha256-js': 3.0.0 '@aws-sdk/client-sso-oidc': 3.583.0(@aws-sdk/client-sts@3.583.0) '@aws-sdk/core': 3.582.0 - '@aws-sdk/credential-provider-node': 3.583.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))(@aws-sdk/client-sts@3.583.0) + '@aws-sdk/credential-provider-node': 3.583.0(@aws-sdk/client-sso-oidc@3.583.0)(@aws-sdk/client-sts@3.583.0) '@aws-sdk/middleware-host-header': 3.577.0 '@aws-sdk/middleware-logger': 3.577.0 '@aws-sdk/middleware-recursion-detection': 3.577.0 @@ -10782,12 +10782,12 @@ snapshots: - '@aws-sdk/client-sso-oidc' - aws-crt - '@aws-sdk/credential-provider-ini@3.568.0(@aws-sdk/client-sso-oidc@3.583.0)(@aws-sdk/client-sts@3.569.0(@aws-sdk/client-sso-oidc@3.569.0))': + '@aws-sdk/credential-provider-ini@3.568.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))(@aws-sdk/client-sts@3.569.0(@aws-sdk/client-sso-oidc@3.569.0))': dependencies: '@aws-sdk/client-sts': 3.569.0(@aws-sdk/client-sso-oidc@3.569.0) '@aws-sdk/credential-provider-env': 3.568.0 '@aws-sdk/credential-provider-process': 3.568.0 - '@aws-sdk/credential-provider-sso': 3.568.0(@aws-sdk/client-sso-oidc@3.583.0) + '@aws-sdk/credential-provider-sso': 3.568.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0)) '@aws-sdk/credential-provider-web-identity': 3.568.0(@aws-sdk/client-sts@3.569.0(@aws-sdk/client-sso-oidc@3.569.0)) '@aws-sdk/types': 3.567.0 '@smithy/credential-provider-imds': 2.3.0 @@ -10799,12 +10799,12 @@ snapshots: - '@aws-sdk/client-sso-oidc' - aws-crt - '@aws-sdk/credential-provider-ini@3.583.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))(@aws-sdk/client-sts@3.583.0)': + '@aws-sdk/credential-provider-ini@3.583.0(@aws-sdk/client-sso-oidc@3.583.0)(@aws-sdk/client-sts@3.583.0)': dependencies: '@aws-sdk/client-sts': 3.583.0 '@aws-sdk/credential-provider-env': 3.577.0 '@aws-sdk/credential-provider-process': 3.577.0 - '@aws-sdk/credential-provider-sso': 3.583.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0)) + '@aws-sdk/credential-provider-sso': 3.583.0(@aws-sdk/client-sso-oidc@3.583.0) '@aws-sdk/credential-provider-web-identity': 3.577.0(@aws-sdk/client-sts@3.583.0) '@aws-sdk/types': 3.577.0 '@smithy/credential-provider-imds': 3.0.0 @@ -10870,13 +10870,13 @@ snapshots: - '@aws-sdk/client-sts' - aws-crt - '@aws-sdk/credential-provider-node@3.569.0(@aws-sdk/client-sso-oidc@3.583.0)(@aws-sdk/client-sts@3.569.0(@aws-sdk/client-sso-oidc@3.569.0))': + '@aws-sdk/credential-provider-node@3.569.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))(@aws-sdk/client-sts@3.569.0(@aws-sdk/client-sso-oidc@3.569.0))': dependencies: '@aws-sdk/credential-provider-env': 3.568.0 '@aws-sdk/credential-provider-http': 3.568.0 - '@aws-sdk/credential-provider-ini': 3.568.0(@aws-sdk/client-sso-oidc@3.583.0)(@aws-sdk/client-sts@3.569.0(@aws-sdk/client-sso-oidc@3.569.0)) + '@aws-sdk/credential-provider-ini': 3.568.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))(@aws-sdk/client-sts@3.569.0(@aws-sdk/client-sso-oidc@3.569.0)) '@aws-sdk/credential-provider-process': 3.568.0 - '@aws-sdk/credential-provider-sso': 3.568.0(@aws-sdk/client-sso-oidc@3.583.0) + '@aws-sdk/credential-provider-sso': 3.568.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0)) '@aws-sdk/credential-provider-web-identity': 3.568.0(@aws-sdk/client-sts@3.569.0(@aws-sdk/client-sso-oidc@3.569.0)) '@aws-sdk/types': 3.567.0 '@smithy/credential-provider-imds': 2.3.0 @@ -10889,13 +10889,13 @@ snapshots: - '@aws-sdk/client-sts' - aws-crt - '@aws-sdk/credential-provider-node@3.583.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))(@aws-sdk/client-sts@3.583.0)': + '@aws-sdk/credential-provider-node@3.583.0(@aws-sdk/client-sso-oidc@3.583.0)(@aws-sdk/client-sts@3.583.0)': dependencies: '@aws-sdk/credential-provider-env': 3.577.0 '@aws-sdk/credential-provider-http': 3.582.0 - '@aws-sdk/credential-provider-ini': 3.583.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))(@aws-sdk/client-sts@3.583.0) + '@aws-sdk/credential-provider-ini': 3.583.0(@aws-sdk/client-sso-oidc@3.583.0)(@aws-sdk/client-sts@3.583.0) '@aws-sdk/credential-provider-process': 3.577.0 - '@aws-sdk/credential-provider-sso': 3.583.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0)) + '@aws-sdk/credential-provider-sso': 3.583.0(@aws-sdk/client-sso-oidc@3.583.0) '@aws-sdk/credential-provider-web-identity': 3.577.0(@aws-sdk/client-sts@3.583.0) '@aws-sdk/types': 3.577.0 '@smithy/credential-provider-imds': 3.0.0 @@ -10957,10 +10957,10 @@ snapshots: - '@aws-sdk/client-sso-oidc' - aws-crt - '@aws-sdk/credential-provider-sso@3.568.0(@aws-sdk/client-sso-oidc@3.583.0)': + '@aws-sdk/credential-provider-sso@3.568.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))': dependencies: '@aws-sdk/client-sso': 3.568.0 - '@aws-sdk/token-providers': 3.568.0(@aws-sdk/client-sso-oidc@3.583.0) + '@aws-sdk/token-providers': 3.568.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0)) '@aws-sdk/types': 3.567.0 '@smithy/property-provider': 2.2.0 '@smithy/shared-ini-file-loader': 2.4.0 @@ -10970,10 +10970,10 @@ snapshots: - '@aws-sdk/client-sso-oidc' - aws-crt - '@aws-sdk/credential-provider-sso@3.583.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))': + '@aws-sdk/credential-provider-sso@3.583.0(@aws-sdk/client-sso-oidc@3.583.0)': dependencies: '@aws-sdk/client-sso': 3.583.0 - '@aws-sdk/token-providers': 3.577.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0)) + '@aws-sdk/token-providers': 3.577.0(@aws-sdk/client-sso-oidc@3.583.0) '@aws-sdk/types': 3.577.0 '@smithy/property-provider': 3.0.0 '@smithy/shared-ini-file-loader': 3.0.0 @@ -11014,7 +11014,7 @@ snapshots: '@smithy/types': 3.0.0 tslib: 2.6.2 - '@aws-sdk/credential-providers@3.569.0(@aws-sdk/client-sso-oidc@3.583.0)': + '@aws-sdk/credential-providers@3.569.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))': dependencies: '@aws-sdk/client-cognito-identity': 3.569.0 '@aws-sdk/client-sso': 3.568.0 @@ -11022,10 +11022,10 @@ snapshots: '@aws-sdk/credential-provider-cognito-identity': 3.569.0 '@aws-sdk/credential-provider-env': 3.568.0 '@aws-sdk/credential-provider-http': 3.568.0 - '@aws-sdk/credential-provider-ini': 3.568.0(@aws-sdk/client-sso-oidc@3.583.0)(@aws-sdk/client-sts@3.569.0(@aws-sdk/client-sso-oidc@3.569.0)) - '@aws-sdk/credential-provider-node': 3.569.0(@aws-sdk/client-sso-oidc@3.583.0)(@aws-sdk/client-sts@3.569.0(@aws-sdk/client-sso-oidc@3.569.0)) + '@aws-sdk/credential-provider-ini': 3.568.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))(@aws-sdk/client-sts@3.569.0(@aws-sdk/client-sso-oidc@3.569.0)) + '@aws-sdk/credential-provider-node': 3.569.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))(@aws-sdk/client-sts@3.569.0(@aws-sdk/client-sso-oidc@3.569.0)) '@aws-sdk/credential-provider-process': 3.568.0 - '@aws-sdk/credential-provider-sso': 3.568.0(@aws-sdk/client-sso-oidc@3.583.0) + '@aws-sdk/credential-provider-sso': 3.568.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0)) '@aws-sdk/credential-provider-web-identity': 3.568.0(@aws-sdk/client-sts@3.569.0(@aws-sdk/client-sso-oidc@3.569.0)) '@aws-sdk/types': 3.567.0 '@smithy/credential-provider-imds': 2.3.0 @@ -11207,7 +11207,7 @@ snapshots: '@smithy/types': 2.12.0 tslib: 2.6.2 - '@aws-sdk/token-providers@3.568.0(@aws-sdk/client-sso-oidc@3.583.0)': + '@aws-sdk/token-providers@3.568.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))': dependencies: '@aws-sdk/client-sso-oidc': 3.583.0(@aws-sdk/client-sts@3.583.0) '@aws-sdk/types': 3.567.0 @@ -11216,7 +11216,7 @@ snapshots: '@smithy/types': 2.12.0 tslib: 2.6.2 - '@aws-sdk/token-providers@3.577.0(@aws-sdk/client-sso-oidc@3.583.0(@aws-sdk/client-sts@3.583.0))': + '@aws-sdk/token-providers@3.577.0(@aws-sdk/client-sso-oidc@3.583.0)': dependencies: '@aws-sdk/client-sso-oidc': 3.583.0(@aws-sdk/client-sts@3.583.0) '@aws-sdk/types': 3.577.0 @@ -12362,7 +12362,7 @@ snapshots: '@drizzle-team/studio@0.0.5': {} - '@electric-sql/pglite@0.1.5': {} + '@electric-sql/pglite@0.2.4': {} '@esbuild-kit/core-utils@3.1.0': dependencies: @@ -15124,7 +15124,7 @@ snapshots: pathe: 1.1.2 picocolors: 1.0.1 sirv: 2.0.4 - vitest: 1.6.0(@types/node@18.19.33)(@vitest/ui@1.6.0)(lightningcss@1.25.1)(terser@5.31.0) + vitest: 1.6.0(@types/node@20.12.12)(@vitest/ui@1.6.0)(lightningcss@1.25.1)(terser@5.31.0) '@vitest/utils@1.6.0': dependencies: From e591c5d51fa07e4b574301d9ea17d043c2508000 Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Thu, 5 Sep 2024 17:41:00 +0300 Subject: [PATCH 06/24] Add neon import --- drizzle-kit/package.json | 2 +- drizzle-orm/package.json | 2 +- drizzle-orm/src/neon/index.ts | 1 + drizzle-orm/src/neon/rls.ts | 46 ++++++++++++++++++++++++ drizzle-orm/src/pg-core/utils.ts | 20 ++++++++++- integration-tests/tests/pg/pg-common.ts | 48 +++++++++++++++++++++++-- pnpm-lock.yaml | 15 +++----- 7 files changed, 117 insertions(+), 17 deletions(-) create mode 100644 drizzle-orm/src/neon/index.ts create mode 100644 drizzle-orm/src/neon/rls.ts diff --git a/drizzle-kit/package.json b/drizzle-kit/package.json index 752b5f7ed..341ee21a6 100644 --- a/drizzle-kit/package.json +++ b/drizzle-kit/package.json @@ -51,7 +51,7 @@ "@arethetypeswrong/cli": "^0.15.3", "@aws-sdk/client-rds-data": "^3.556.0", "@cloudflare/workers-types": "^4.20230518.0", - "@electric-sql/pglite": "^0.2.4", + "@electric-sql/pglite": "^0.2.5", "@hono/node-server": "^1.9.0", "@hono/zod-validator": "^0.2.1", "@libsql/client": "^0.10.0", diff --git a/drizzle-orm/package.json b/drizzle-orm/package.json index 4ae5d5e5a..707641969 100644 --- a/drizzle-orm/package.json +++ b/drizzle-orm/package.json @@ -160,7 +160,7 @@ "devDependencies": { "@aws-sdk/client-rds-data": "^3.549.0", "@cloudflare/workers-types": "^4.20230904.0", - "@electric-sql/pglite": "^0.2.4", + "@electric-sql/pglite": "^0.2.5", "@libsql/client": "^0.10.0", "@neondatabase/serverless": "^0.9.0", "@op-engineering/op-sqlite": "^2.0.16", diff --git a/drizzle-orm/src/neon/index.ts b/drizzle-orm/src/neon/index.ts new file mode 100644 index 000000000..ee201ff1c --- /dev/null +++ b/drizzle-orm/src/neon/index.ts @@ -0,0 +1 @@ +export * from './rls.ts'; diff --git a/drizzle-orm/src/neon/rls.ts b/drizzle-orm/src/neon/rls.ts new file mode 100644 index 000000000..de5c49be2 --- /dev/null +++ b/drizzle-orm/src/neon/rls.ts @@ -0,0 +1,46 @@ +import { pgPolicy, pgRole } from '~/pg-core/index.ts'; +import type { AnyPgColumn, PgPolicyToOption } from '~/pg-core/index.ts'; +import { type SQL, sql } from '~/sql/sql.ts'; + +export const crudPolicy = ( + options: { + role: PgPolicyToOption; + read: SQL | boolean; + modify: SQL | boolean; + }, +) => { + const read: SQL = options.read === true + ? sql`select true` + : options.read === false + ? sql`select false` + : options.read; + + const modify: SQL = options.modify === true + ? sql`select true` + : options.modify === false + ? sql`select false` + : options.modify; + + // Return the modify policy, followed by the read policy. + return { + // Important to have "_drizzle_internal" prefix for any key here. Right after we will make + // 3rd param in table as an array - we will move it to array and use ... operator + [`_drizzle_internal-crud-policy-modify`]: pgPolicy(`crud-policy-modify`, { + for: 'insert', + to: options.role, + using: modify, + withCheck: modify, + }), + [`_drizzle_internal-crud-policy-read`]: pgPolicy(`crud-policy-read`, { + for: 'select', + to: options.role, + using: read, + }), + }; +}; + +// These are default roles that Neon will set up. +export const authenticatedRole = pgRole('authenticated').existing(); +export const anonymousRole = pgRole('anonymous').existing(); + +export const authUid = (userIdColumn: AnyPgColumn) => sql`select auth.user_id() = ${userIdColumn}`; diff --git a/drizzle-orm/src/pg-core/utils.ts b/drizzle-orm/src/pg-core/utils.ts index 43978ae58..3e7c153af 100644 --- a/drizzle-orm/src/pg-core/utils.ts +++ b/drizzle-orm/src/pg-core/utils.ts @@ -1,4 +1,5 @@ import { is } from '~/entity.ts'; +import type { PgTableExtraConfig } from '~/pg-core/table.ts'; import { PgTable } from '~/pg-core/table.ts'; import { Table } from '~/table.ts'; import { ViewBaseConfig } from '~/view-common.ts'; @@ -13,6 +14,23 @@ import { type UniqueConstraint, UniqueConstraintBuilder } from './unique-constra import { PgViewConfig } from './view-common.ts'; import { type PgMaterializedView, PgMaterializedViewConfig, type PgView } from './view.ts'; +function flattenValues(obj: Record): PgTableExtraConfig[string][] { + const values = []; + for (const key in obj) { + if (typeof obj[key] === 'object' && obj[key] !== null) { + const internalNestedKeys = Object.keys(obj[key]).filter((it) => it.startsWith('_drizzle_internal')); + if (internalNestedKeys.length > 0) { + for (const nestedKey of internalNestedKeys) { + values.push(obj[key][nestedKey]); + } + } else { + values.push(obj[key]); + } + } + } + return values; +} + export function getTableConfig(table: TTable) { const columns = Object.values(table[Table.Symbol.Columns]); const indexes: Index[] = []; @@ -28,7 +46,7 @@ export function getTableConfig(table: TTable) { if (extraConfigBuilder !== undefined) { const extraConfig = extraConfigBuilder(table[Table.Symbol.ExtraConfigColumns]); - for (const builder of Object.values(extraConfig)) { + for (const builder of flattenValues(extraConfig)) { if (is(builder, IndexBuilder)) { indexes.push(builder.build(table)); } else if (is(builder, CheckBuilder)) { diff --git a/integration-tests/tests/pg/pg-common.ts b/integration-tests/tests/pg/pg-common.ts index f098d3b38..03343847d 100644 --- a/integration-tests/tests/pg/pg-common.ts +++ b/integration-tests/tests/pg/pg-common.ts @@ -31,6 +31,7 @@ import { sumDistinct, TransactionRollbackError, } from 'drizzle-orm'; +import { authenticatedRole, crudPolicy } from 'drizzle-orm/neon'; import type { NeonHttpDatabase } from 'drizzle-orm/neon-http'; import type { PgColumn, PgDatabase, PgQueryResultHKT } from 'drizzle-orm/pg-core'; import { @@ -45,6 +46,7 @@ import { getMaterializedViewConfig, getTableConfig, getViewConfig, + index, inet, integer, intersect, @@ -4677,7 +4679,7 @@ export function tests() { const policy = pgPolicy('test policy', { as: 'permissive', for: 'all', - to: 'PUBLIC', + to: 'public', using: sql`1=1`, withCheck: sql`1=1`, }); @@ -4686,7 +4688,7 @@ export function tests() { expect(policy.name).toBe('test policy'); expect(policy.as).toBe('permissive'); expect(policy.for).toBe('all'); - expect(policy.to).toBe('PUBLIC'); + expect(policy.to).toBe('public'); const dialect = new PgDialect(); expect(is(policy.using, SQL)).toBe(true); expect(dialect.sqlToQuery(policy.using!).sql).toBe('1=1'); @@ -4707,7 +4709,7 @@ export function tests() { const p2 = pgPolicy('test policy 2', { as: 'permissive', for: 'all', - to: 'PUBLIC', + to: 'public', using: sql`1=1`, withCheck: sql`1=1`, }); @@ -4725,6 +4727,46 @@ export function tests() { } }); + test.only('neon: policy', () => { + { + const policy = crudPolicy({ + read: true, + modify: true, + role: authenticatedRole, + }); + + for (const it of Object.values(policy)) { + expect(is(it, PgPolicy)).toBe(true); + expect(it.to).toStrictEqual(authenticatedRole); + expect(it.using).toStrictEqual(sql`select true`); + it.withCheck ? expect(it.withCheck).toStrictEqual(sql`select true`) : ''; + } + } + + { + const table = pgTable('name', { + id: integer('id'), + }, (t) => ({ + in: index('name').on(t.id), + neonPolicy: crudPolicy({ + read: true, + modify: true, + role: authenticatedRole, + }), + pk: primaryKey({ columns: [t.id], name: 'custom' }), + })); + + const { policies, indexes, primaryKeys } = getTableConfig(table); + + expect(policies.length).toBe(2); + expect(indexes.length).toBe(1); + expect(primaryKeys.length).toBe(1); + + expect(policies[0]?.name === 'crud-policy-modify'); + expect(policies[1]?.name === 'crud-policy-read'); + } + }); + test('$count separate', async (ctx) => { const { db } = ctx.pg; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 222526efd..75a9d50d4 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -114,8 +114,8 @@ importers: specifier: ^4.20230518.0 version: 4.20240524.0 '@electric-sql/pglite': - specifier: ^0.2.4 - version: 0.2.4 + specifier: ^0.2.5 + version: 0.2.5 '@hono/node-server': specifier: ^1.9.0 version: 1.12.0 @@ -303,8 +303,8 @@ importers: specifier: ^4.20230904.0 version: 4.20240512.0 '@electric-sql/pglite': - specifier: ^0.2.4 - version: 0.2.4 + specifier: ^0.2.5 + version: 0.2.5 '@libsql/client': specifier: ^0.10.0 version: 0.10.0(bufferutil@4.0.8)(utf-8-validate@6.0.3) @@ -1983,9 +1983,6 @@ packages: '@drizzle-team/studio@0.0.5': resolution: {integrity: sha512-ps5qF0tMxWRVu+V5gvCRrQNqlY92aTnIKdq27gm9LZMSdaKYZt6AVvSK1dlUMzs6Rt0Jm80b+eWct6xShBKhIw==} - '@electric-sql/pglite@0.2.4': - resolution: {integrity: sha512-NN1ATH9aYTCD4257wZH1CjAKdro8jDd5r/BumxZtEVTDNDztBmPYYtBd4TEIHWdi0+QuM7SWk2JRCLvAYCSEWg==} - '@electric-sql/pglite@0.2.5': resolution: {integrity: sha512-LrMX2kX0mCVN4xkhIDv1KBVukWtoOI/+P9MDQgHX5QEeZCi5S60LZOa0VWXjufPEz7mJtbuXWJRujD++t0gsHA==} @@ -7092,12 +7089,10 @@ packages: libsql@0.3.19: resolution: {integrity: sha512-Aj5cQ5uk/6fHdmeW0TiXK42FqUlwx7ytmMLPSaUQPin5HKKKuUPD62MAbN4OEweGBBI7q1BekoEN4gPUEL6MZA==} - cpu: [x64, arm64, wasm32] os: [darwin, linux, win32] libsql@0.4.1: resolution: {integrity: sha512-qZlR9Yu1zMBeLChzkE/cKfoKV3Esp9cn9Vx5Zirn4AVhDWPcjYhKwbtJcMuHehgk3mH+fJr9qW+3vesBWbQpBg==} - cpu: [x64, arm64, wasm32] os: [darwin, linux, win32] lighthouse-logger@1.4.2: @@ -12364,8 +12359,6 @@ snapshots: '@drizzle-team/studio@0.0.5': {} - '@electric-sql/pglite@0.2.4': {} - '@electric-sql/pglite@0.2.5': {} '@esbuild-kit/core-utils@3.1.0': From dc783a22f0852eea2eacb63fa1934ff75d648382 Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Fri, 6 Sep 2024 09:44:28 +0300 Subject: [PATCH 07/24] Fix tests for kit and orm --- drizzle-kit/tests/push/libsql.test.ts | 1 + integration-tests/tests/pg/pg-common.ts | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/drizzle-kit/tests/push/libsql.test.ts b/drizzle-kit/tests/push/libsql.test.ts index 89ec008ca..5dc77abf7 100644 --- a/drizzle-kit/tests/push/libsql.test.ts +++ b/drizzle-kit/tests/push/libsql.test.ts @@ -651,6 +651,7 @@ test('drop table with data', async (t) => { expect(statements!.length).toBe(1); expect(statements![0]).toStrictEqual({ + policies: [], schema: undefined, tableName: 'users', type: 'drop_table', diff --git a/integration-tests/tests/pg/pg-common.ts b/integration-tests/tests/pg/pg-common.ts index 45223e516..3f191e618 100644 --- a/integration-tests/tests/pg/pg-common.ts +++ b/integration-tests/tests/pg/pg-common.ts @@ -4727,7 +4727,7 @@ export function tests() { } }); - test.only('neon: policy', () => { + test('neon: policy', () => { { const policy = crudPolicy({ read: true, From 109837d10a0b4f40e494917e1825fc10a430833e Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Fri, 6 Sep 2024 10:11:33 +0300 Subject: [PATCH 08/24] Update pglite tests --- integration-tests/tests/pg/pglite.test.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/integration-tests/tests/pg/pglite.test.ts b/integration-tests/tests/pg/pglite.test.ts index 37cd3fe62..20b1b61b3 100644 --- a/integration-tests/tests/pg/pglite.test.ts +++ b/integration-tests/tests/pg/pglite.test.ts @@ -86,6 +86,9 @@ skipTests([ 'subquery with view', 'mySchema :: materialized view', 'select count()', + 'mySchema :: select with group by as column + sql', + 'select with group by as column + sql', + 'select with group by as sql + column', ]); tests(); From 7bea25c6d9ddb6bd5651b30987930be1f1083a0c Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Fri, 6 Sep 2024 11:00:54 +0300 Subject: [PATCH 09/24] Fix role import --- drizzle-orm/src/pg-core/policies.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/drizzle-orm/src/pg-core/policies.ts b/drizzle-orm/src/pg-core/policies.ts index a61731578..74a3ea85c 100644 --- a/drizzle-orm/src/pg-core/policies.ts +++ b/drizzle-orm/src/pg-core/policies.ts @@ -1,6 +1,6 @@ import { entityKind } from '~/entity.ts'; import type { SQL } from '~/sql/sql.ts'; -import type { PgRole } from './roles'; +import type { PgRole } from './roles.ts'; export type PgPolicyToOption = | 'public' From 7469abe066a1e135cd04416be74cf97a38f27878 Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Fri, 6 Sep 2024 12:50:59 +0300 Subject: [PATCH 10/24] Update naming for crudPolicy keys --- drizzle-orm/src/neon/rls.ts | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/drizzle-orm/src/neon/rls.ts b/drizzle-orm/src/neon/rls.ts index de5c49be2..85d02bec0 100644 --- a/drizzle-orm/src/neon/rls.ts +++ b/drizzle-orm/src/neon/rls.ts @@ -1,37 +1,47 @@ -import { pgPolicy, pgRole } from '~/pg-core/index.ts'; +import { is } from '~/entity.ts'; +import { pgPolicy, PgRole, pgRole } from '~/pg-core/index.ts'; import type { AnyPgColumn, PgPolicyToOption } from '~/pg-core/index.ts'; import { type SQL, sql } from '~/sql/sql.ts'; export const crudPolicy = ( options: { role: PgPolicyToOption; - read: SQL | boolean; - modify: SQL | boolean; + read?: SQL | boolean; + modify?: SQL | boolean; }, ) => { const read: SQL = options.read === true ? sql`select true` - : options.read === false + : options.read === false || options.read === undefined ? sql`select false` : options.read; const modify: SQL = options.modify === true ? sql`select true` - : options.modify === false + : options.modify === false || options.modify === undefined ? sql`select false` : options.modify; + let rolesName = ''; + if (Array.isArray(options.role)) { + rolesName = options.role.map((it) => { + return is(it, PgRole) ? it.name : it as string; + }).join('-'); + } else { + rolesName = is(options.role, PgRole) ? options.role.name : options.role as string; + } + // Return the modify policy, followed by the read policy. return { // Important to have "_drizzle_internal" prefix for any key here. Right after we will make // 3rd param in table as an array - we will move it to array and use ... operator - [`_drizzle_internal-crud-policy-modify`]: pgPolicy(`crud-policy-modify`, { + [`_drizzle_internal-${rolesName}-crud-policy-modify`]: pgPolicy(`crud-policy-modify`, { for: 'insert', to: options.role, using: modify, withCheck: modify, }), - [`_drizzle_internal-crud-policy-read`]: pgPolicy(`crud-policy-read`, { + [`_drizzle_internal-${rolesName}-crud-policy-read`]: pgPolicy(`crud-policy-read`, { for: 'select', to: options.role, using: read, From 277e6de7ad371876924de3b72bf684167fe5b85e Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Fri, 6 Sep 2024 16:04:01 +0300 Subject: [PATCH 11/24] Fix sql to query mapping in policy --- drizzle-kit/src/serializer/pgSerializer.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/drizzle-kit/src/serializer/pgSerializer.ts b/drizzle-kit/src/serializer/pgSerializer.ts index 7b84b32f4..12cae6b08 100644 --- a/drizzle-kit/src/serializer/pgSerializer.ts +++ b/drizzle-kit/src/serializer/pgSerializer.ts @@ -526,8 +526,8 @@ export const generatePgSnapshot = ( as: policy.as?.toUpperCase() as Policy['as'] ?? 'PERMISSIVE', for: policy.for?.toUpperCase() as Policy['for'] ?? 'ALL', to: mappedTo.sort(), - using: is(policy.using, SQL) ? sqlToStr(policy.using) : undefined, - withCheck: is(policy.withCheck, SQL) ? sqlToStr(policy.withCheck) : undefined, + using: is(policy.using, SQL) ? dialect.sqlToQuery(policy.using).sql : undefined, + withCheck: is(policy.withCheck, SQL) ? dialect.sqlToQuery(policy.withCheck).sql : undefined, }; }); From d496e6f851040d37e73eef16d21325ee06713417 Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Fri, 6 Sep 2024 18:55:18 +0300 Subject: [PATCH 12/24] Create policies separately --- drizzle-kit/src/snapshotsDiffer.ts | 6 ++++++ drizzle-kit/src/sqlgenerator.ts | 11 +---------- drizzle-kit/tests/push/pg.test.ts | 15 +++++++++++++++ drizzle-kit/tests/rls/pg-policy.test.ts | 15 +++++++++++++++ 4 files changed, 37 insertions(+), 10 deletions(-) diff --git a/drizzle-kit/src/snapshotsDiffer.ts b/drizzle-kit/src/snapshotsDiffer.ts index d0f04cdfb..007a7d10a 100644 --- a/drizzle-kit/src/snapshotsDiffer.ts +++ b/drizzle-kit/src/snapshotsDiffer.ts @@ -1443,6 +1443,12 @@ export const applyPgSnapshotsDiff = async ( return preparePgCreateTableJson(it, curFull); }); + jsonCreatePoliciesStatements.push(...([] as JsonCreatePolicyStatement[]).concat( + ...(createdTables.map((it) => + prepareCreatePolicyJsons(it.name, it.schema, Object.values(it.policies).map(PgSquasher.unsquashPolicy)) + )), + )); + jsonStatements.push(...createSchemas); jsonStatements.push(...renameSchemas); jsonStatements.push(...createEnums); diff --git a/drizzle-kit/src/sqlgenerator.ts b/drizzle-kit/src/sqlgenerator.ts index a317e7331..7d87319b4 100644 --- a/drizzle-kit/src/sqlgenerator.ts +++ b/drizzle-kit/src/sqlgenerator.ts @@ -394,22 +394,13 @@ class PgCreateTableConvertor extends Convertor { statement += `\n);`; statement += `\n`; - const createPolicyConvertor = new PgCreatePolicyConvertor(); - const createPolicies = policies?.map((p) => { - return createPolicyConvertor.convert({ - type: 'create_policy', - tableName, - data: PgSquasher.unsquashPolicy(p), - schema, - }) as string; - }) ?? []; const enableRls = new PgEnableRlsConvertor().convert({ type: 'enable_rls', tableName, schema, }); - return [statement, ...(createPolicies.length > 0 ? [enableRls] : []), ...createPolicies]; + return [statement, ...(policies && policies.length > 0 ? [enableRls] : [])]; } } diff --git a/drizzle-kit/tests/push/pg.test.ts b/drizzle-kit/tests/push/pg.test.ts index 88f02e8ae..e3ed739f0 100644 --- a/drizzle-kit/tests/push/pg.test.ts +++ b/drizzle-kit/tests/push/pg.test.ts @@ -3024,6 +3024,21 @@ test('create table with a policy', async (t) => { type: 'create_table', uniqueConstraints: [], }, + { + data: { + as: 'PERMISSIVE', + for: 'ALL', + name: 'test', + to: [ + 'public', + ], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users2', + type: 'create_policy', + }, ]); for (const st of sqlStatements) { diff --git a/drizzle-kit/tests/rls/pg-policy.test.ts b/drizzle-kit/tests/rls/pg-policy.test.ts index 17827f951..f6fec3f60 100644 --- a/drizzle-kit/tests/rls/pg-policy.test.ts +++ b/drizzle-kit/tests/rls/pg-policy.test.ts @@ -599,6 +599,21 @@ test('create table with a policy', async (t) => { type: 'create_table', uniqueConstraints: [], }, + { + data: { + as: 'PERMISSIVE', + for: 'ALL', + name: 'test', + to: [ + 'public', + ], + using: undefined, + withCheck: undefined, + }, + schema: '', + tableName: 'users2', + type: 'create_policy', + }, ]); }); From 2475c4fd4df2922b369cb1ee0fe1ade684eab1ba Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Fri, 6 Sep 2024 19:27:35 +0300 Subject: [PATCH 13/24] Add roles to policy name --- drizzle-orm/src/neon/rls.ts | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/drizzle-orm/src/neon/rls.ts b/drizzle-orm/src/neon/rls.ts index 85d02bec0..6f3a2f870 100644 --- a/drizzle-orm/src/neon/rls.ts +++ b/drizzle-orm/src/neon/rls.ts @@ -35,13 +35,28 @@ export const crudPolicy = ( return { // Important to have "_drizzle_internal" prefix for any key here. Right after we will make // 3rd param in table as an array - we will move it to array and use ... operator - [`_drizzle_internal-${rolesName}-crud-policy-modify`]: pgPolicy(`crud-policy-modify`, { + + // We can't use table name here, because in examples you can specify several crudPolicies on one table + // So we need some other way to have a unique name + [`_drizzle_internal-${rolesName}-crud-policy-insert`]: pgPolicy(`crud-${rolesName}-policy-insert`, { for: 'insert', to: options.role, using: modify, withCheck: modify, }), - [`_drizzle_internal-${rolesName}-crud-policy-read`]: pgPolicy(`crud-policy-read`, { + [`_drizzle_internal-${rolesName}-crud-policy-update`]: pgPolicy(`crud-${rolesName}-policy-update`, { + for: 'update', + to: options.role, + using: modify, + withCheck: modify, + }), + [`_drizzle_internal-${rolesName}-crud-policy-delete`]: pgPolicy(`crud-${rolesName}-policy-delete`, { + for: 'delete', + to: options.role, + using: modify, + withCheck: modify, + }), + [`_drizzle_internal-${rolesName}-crud-policy-select`]: pgPolicy(`crud-${rolesName}-policy-select`, { for: 'select', to: options.role, using: read, From 8944ef11c733af9226b574716802c7c4e5a411c9 Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Fri, 6 Sep 2024 19:43:14 +0300 Subject: [PATCH 14/24] Fix tests after policy update --- integration-tests/tests/pg/pg-common.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/integration-tests/tests/pg/pg-common.ts b/integration-tests/tests/pg/pg-common.ts index 3f191e618..ab44b64b5 100644 --- a/integration-tests/tests/pg/pg-common.ts +++ b/integration-tests/tests/pg/pg-common.ts @@ -4758,7 +4758,7 @@ export function tests() { const { policies, indexes, primaryKeys } = getTableConfig(table); - expect(policies.length).toBe(2); + expect(policies.length).toBe(4); expect(indexes.length).toBe(1); expect(primaryKeys.length).toBe(1); From b16ebacee2b2977cd3eaa7f45b57b2dfb9baacbb Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Wed, 9 Oct 2024 17:01:28 +0300 Subject: [PATCH 15/24] Add rls fixes and 3 param as array --- drizzle-orm/src/neon/rls.ts | 26 ++++----- drizzle-orm/src/pg-core/table.ts | 55 +++++++++++++++++-- drizzle-orm/src/pg-core/utils.ts | 21 +------ drizzle-orm/src/table.ts | 2 +- drizzle-orm/type-tests/pg/tables.ts | 22 ++++---- .../tests/pg/rls/rls.definition.test.ts | 16 ++++++ 6 files changed, 88 insertions(+), 54 deletions(-) create mode 100644 integration-tests/tests/pg/rls/rls.definition.test.ts diff --git a/drizzle-orm/src/neon/rls.ts b/drizzle-orm/src/neon/rls.ts index 6f3a2f870..9bb421434 100644 --- a/drizzle-orm/src/neon/rls.ts +++ b/drizzle-orm/src/neon/rls.ts @@ -11,15 +11,15 @@ export const crudPolicy = ( }, ) => { const read: SQL = options.read === true - ? sql`select true` + ? sql`true` : options.read === false || options.read === undefined - ? sql`select false` + ? sql`false` : options.read; const modify: SQL = options.modify === true - ? sql`select true` + ? sql`true` : options.modify === false || options.modify === undefined - ? sql`select false` + ? sql`false` : options.modify; let rolesName = ''; @@ -31,37 +31,31 @@ export const crudPolicy = ( rolesName = is(options.role, PgRole) ? options.role.name : options.role as string; } - // Return the modify policy, followed by the read policy. - return { - // Important to have "_drizzle_internal" prefix for any key here. Right after we will make - // 3rd param in table as an array - we will move it to array and use ... operator - - // We can't use table name here, because in examples you can specify several crudPolicies on one table - // So we need some other way to have a unique name - [`_drizzle_internal-${rolesName}-crud-policy-insert`]: pgPolicy(`crud-${rolesName}-policy-insert`, { + return [ + pgPolicy(`crud-${rolesName}-policy-insert`, { for: 'insert', to: options.role, using: modify, withCheck: modify, }), - [`_drizzle_internal-${rolesName}-crud-policy-update`]: pgPolicy(`crud-${rolesName}-policy-update`, { + pgPolicy(`crud-${rolesName}-policy-update`, { for: 'update', to: options.role, using: modify, withCheck: modify, }), - [`_drizzle_internal-${rolesName}-crud-policy-delete`]: pgPolicy(`crud-${rolesName}-policy-delete`, { + pgPolicy(`crud-${rolesName}-policy-delete`, { for: 'delete', to: options.role, using: modify, withCheck: modify, }), - [`_drizzle_internal-${rolesName}-crud-policy-select`]: pgPolicy(`crud-${rolesName}-policy-select`, { + pgPolicy(`crud-${rolesName}-policy-select`, { for: 'select', to: options.role, using: read, }), - }; + ]; }; // These are default roles that Neon will set up. diff --git a/drizzle-orm/src/pg-core/table.ts b/drizzle-orm/src/pg-core/table.ts index b7804372b..63e82cfd0 100644 --- a/drizzle-orm/src/pg-core/table.ts +++ b/drizzle-orm/src/pg-core/table.ts @@ -10,14 +10,17 @@ import type { PgPolicy } from './policies.ts'; import type { PrimaryKeyBuilder } from './primary-keys.ts'; import type { UniqueConstraintBuilder } from './unique-constraint.ts'; -export type PgTableExtraConfig = Record< - string, +export type PgTableExtraConfigValue = | AnyIndexBuilder | CheckBuilder | ForeignKeyBuilder | PrimaryKeyBuilder | UniqueConstraintBuilder - | PgPolicy + | PgPolicy; + +export type PgTableExtraConfig = Record< + string, + PgTableExtraConfigValue >; export type TableConfig = TableConfigBase; @@ -57,7 +60,9 @@ export function pgTableWithSchema< >( name: TTableName, columns: TColumnsMap | ((columnTypes: PgColumnsBuilders) => TColumnsMap), - extraConfig: ((self: BuildExtraConfigColumns) => PgTableExtraConfig) | undefined, + extraConfig: + | ((self: BuildExtraConfigColumns) => PgTableExtraConfig | PgTableExtraConfigValue[]) + | undefined, schema: TSchemaName, baseName = name, ): PgTableWithColumns<{ @@ -107,13 +112,51 @@ export function pgTableWithSchema< } export interface PgTableFn { + /** + * @deprecated This overload is deprecated. Use the other method overload instead. + */ + < + TTableName extends string, + TColumnsMap extends Record, + >( + name: TTableName, + columns: TColumnsMap, + extraConfig: ( + self: BuildExtraConfigColumns, + ) => PgTableExtraConfig, + ): PgTableWithColumns<{ + name: TTableName; + schema: TSchema; + columns: BuildColumns; + dialect: 'pg'; + }>; + + /** + * @deprecated This overload is deprecated. Use the other method overload instead. + */ + < + TTableName extends string, + TColumnsMap extends Record, + >( + name: TTableName, + columns: (columnTypes: PgColumnsBuilders) => TColumnsMap, + extraConfig: (self: BuildExtraConfigColumns) => PgTableExtraConfig, + ): PgTableWithColumns<{ + name: TTableName; + schema: TSchema; + columns: BuildColumns; + dialect: 'pg'; + }>; + < TTableName extends string, TColumnsMap extends Record, >( name: TTableName, columns: TColumnsMap, - extraConfig?: (self: BuildExtraConfigColumns) => PgTableExtraConfig, + extraConfig?: ( + self: BuildExtraConfigColumns, + ) => PgTableExtraConfigValue[], ): PgTableWithColumns<{ name: TTableName; schema: TSchema; @@ -127,7 +170,7 @@ export interface PgTableFn { >( name: TTableName, columns: (columnTypes: PgColumnsBuilders) => TColumnsMap, - extraConfig?: (self: BuildExtraConfigColumns) => PgTableExtraConfig, + extraConfig?: (self: BuildExtraConfigColumns) => PgTableExtraConfigValue[], ): PgTableWithColumns<{ name: TTableName; schema: TSchema; diff --git a/drizzle-orm/src/pg-core/utils.ts b/drizzle-orm/src/pg-core/utils.ts index 3e7c153af..d7e6c5d8f 100644 --- a/drizzle-orm/src/pg-core/utils.ts +++ b/drizzle-orm/src/pg-core/utils.ts @@ -1,5 +1,4 @@ import { is } from '~/entity.ts'; -import type { PgTableExtraConfig } from '~/pg-core/table.ts'; import { PgTable } from '~/pg-core/table.ts'; import { Table } from '~/table.ts'; import { ViewBaseConfig } from '~/view-common.ts'; @@ -14,23 +13,6 @@ import { type UniqueConstraint, UniqueConstraintBuilder } from './unique-constra import { PgViewConfig } from './view-common.ts'; import { type PgMaterializedView, PgMaterializedViewConfig, type PgView } from './view.ts'; -function flattenValues(obj: Record): PgTableExtraConfig[string][] { - const values = []; - for (const key in obj) { - if (typeof obj[key] === 'object' && obj[key] !== null) { - const internalNestedKeys = Object.keys(obj[key]).filter((it) => it.startsWith('_drizzle_internal')); - if (internalNestedKeys.length > 0) { - for (const nestedKey of internalNestedKeys) { - values.push(obj[key][nestedKey]); - } - } else { - values.push(obj[key]); - } - } - } - return values; -} - export function getTableConfig(table: TTable) { const columns = Object.values(table[Table.Symbol.Columns]); const indexes: Index[] = []; @@ -46,7 +28,8 @@ export function getTableConfig(table: TTable) { if (extraConfigBuilder !== undefined) { const extraConfig = extraConfigBuilder(table[Table.Symbol.ExtraConfigColumns]); - for (const builder of flattenValues(extraConfig)) { + const extraValues = Array.isArray(extraConfig) ? extraConfig as any[] : Object.values(extraConfig); + for (const builder of extraValues) { if (is(builder, IndexBuilder)) { indexes.push(builder.build(table)); } else if (is(builder, CheckBuilder)) { diff --git a/drizzle-orm/src/table.ts b/drizzle-orm/src/table.ts index 0bf08fb3b..6bacfc207 100644 --- a/drizzle-orm/src/table.ts +++ b/drizzle-orm/src/table.ts @@ -109,7 +109,7 @@ export class Table implements SQLWrapper { [IsDrizzleTable] = true; /** @internal */ - [ExtraConfigBuilder]: ((self: any) => Record) | undefined = undefined; + [ExtraConfigBuilder]: ((self: any) => Record | unknown[]) | undefined = undefined; constructor(name: string, schema: string | undefined, baseName: string) { this[TableName] = this[OriginalName] = name; diff --git a/drizzle-orm/type-tests/pg/tables.ts b/drizzle-orm/type-tests/pg/tables.ts index 068f8fcf6..decd61d92 100644 --- a/drizzle-orm/type-tests/pg/tables.ts +++ b/drizzle-orm/type-tests/pg/tables.ts @@ -78,23 +78,23 @@ export const users = pgTable( enumCol: myEnum('enum_col').notNull(), arrayCol: text('array_col').array().notNull(), }, - (users) => ({ - usersAge1Idx: uniqueIndex('usersAge1Idx').on(users.class.asc().nullsFirst(), sql``), - usersAge2Idx: index('usersAge2Idx').on(sql``), - uniqueClass: uniqueIndex('uniqueClass') + (users) => [ + uniqueIndex('usersAge1Idx').on(users.class.asc().nullsFirst(), sql``), + index('usersAge2Idx').on(sql``), + uniqueIndex('uniqueClass') .using('btree', users.class.desc().op('text_ops'), users.subClass.nullsLast()) .where(sql`${users.class} is not null`) .concurrently(), - legalAge: check('legalAge', sql`${users.age1} > 18`), - usersClassFK: foreignKey({ columns: [users.subClass], foreignColumns: [classes.subClass] }) + check('legalAge', sql`${users.age1} > 18`), + foreignKey({ columns: [users.subClass], foreignColumns: [classes.subClass] }) .onUpdate('cascade') .onDelete('cascade'), - usersClassComplexFK: foreignKey({ + foreignKey({ columns: [users.class, users.subClass], foreignColumns: [classes.class, classes.subClass], }), - pk: primaryKey(users.age1, users.class), - }), + primaryKey(users.age1, users.class), + ], ); Expect, typeof users['$inferSelect']>>; @@ -172,9 +172,7 @@ export const citiesCustom = customSchema.table('cities_table', { id: serial('id').primaryKey(), name: text('name').notNull(), population: integer('population').default(0), -}, (cities) => ({ - citiesNameIdx: index().on(cities.id), -})); +}, (cities) => [index().on(cities.id)]); export const newYorkers = pgView('new_yorkers') .with({ diff --git a/integration-tests/tests/pg/rls/rls.definition.test.ts b/integration-tests/tests/pg/rls/rls.definition.test.ts new file mode 100644 index 000000000..611116aae --- /dev/null +++ b/integration-tests/tests/pg/rls/rls.definition.test.ts @@ -0,0 +1,16 @@ +import { crudPolicy } from 'drizzle-orm/neon'; +import { getTableConfig, integer, pgPolicy, pgRole, pgTable } from 'drizzle-orm/pg-core'; +import { test } from 'vitest'; + +test('getTableConfig: policies', async () => { + const schema = pgTable('hhh', { + id: integer(), + }, () => [ + pgPolicy('name'), + crudPolicy({ role: pgRole('users') }), + ]); + + const tc = getTableConfig(schema); + + console.log(tc.policies); +}); From 0986503b0c6dc07bf0967907cafaf26e22fcc319 Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Wed, 9 Oct 2024 17:21:51 +0300 Subject: [PATCH 16/24] Add policies to counter in drizzle pull --- drizzle-kit/src/cli/views.ts | 9 ++++++++- drizzle-kit/src/serializer/pgSerializer.ts | 10 ++++++++++ drizzle-kit/src/snapshotsDiffer.ts | 16 +++++++++------- 3 files changed, 27 insertions(+), 8 deletions(-) diff --git a/drizzle-kit/src/cli/views.ts b/drizzle-kit/src/cli/views.ts index 5cae2ac3f..041239099 100644 --- a/drizzle-kit/src/cli/views.ts +++ b/drizzle-kit/src/cli/views.ts @@ -401,7 +401,8 @@ export type IntrospectStage = | 'columns' | 'enums' | 'indexes' - | 'fks'; + | 'fks' + | 'policies'; type IntrospectState = { [key in IntrospectStage]: { count: number; @@ -440,6 +441,11 @@ export class IntrospectProgress extends TaskView { name: 'foreign keys', status: 'fetching', }, + policies: { + count: 0, + name: 'policies', + status: 'fetching', + }, }; constructor(private readonly hasEnums: boolean = false) { @@ -493,6 +499,7 @@ export class IntrospectProgress extends TaskView { info += this.hasEnums ? this.statusText(spin, this.state.enums) : ''; info += this.statusText(spin, this.state.indexes); info += this.statusText(spin, this.state.fks); + info += this.statusText(spin, this.state.policies); return info; } } diff --git a/drizzle-kit/src/serializer/pgSerializer.ts b/drizzle-kit/src/serializer/pgSerializer.ts index 3bad54100..7132013e8 100644 --- a/drizzle-kit/src/serializer/pgSerializer.ts +++ b/drizzle-kit/src/serializer/pgSerializer.ts @@ -909,6 +909,16 @@ export const fromDatabase = async ( } } + if (progressCallback) { + progressCallback( + 'policies', + Object.values(policiesByTable).reduce((total, innerRecord) => { + return total + Object.keys(innerRecord).length; + }, 0), + 'done', + ); + } + const sequencesInColumns: string[] = []; const all = allTables.map((row) => { diff --git a/drizzle-kit/src/snapshotsDiffer.ts b/drizzle-kit/src/snapshotsDiffer.ts index 007a7d10a..e75df6584 100644 --- a/drizzle-kit/src/snapshotsDiffer.ts +++ b/drizzle-kit/src/snapshotsDiffer.ts @@ -988,13 +988,15 @@ export const applyPgSnapshotsDiff = async ( const jsonSetTableSchemas: JsonAlterTableSetSchema[] = []; - for (let it of movedTables) { - jsonSetTableSchemas.push({ - type: 'alter_table_set_schema', - tableName: it.name, - schemaFrom: it.schemaFrom || 'public', - schemaTo: it.schemaTo || 'public', - }); + if (movedTables) { + for (let it of movedTables) { + jsonSetTableSchemas.push({ + type: 'alter_table_set_schema', + tableName: it.name, + schemaFrom: it.schemaFrom || 'public', + schemaTo: it.schemaTo || 'public', + }); + } } for (let it of alteredTables) { From 225f91d26f93b24efd518866ec998275087effdf Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Thu, 10 Oct 2024 10:38:29 +0300 Subject: [PATCH 17/24] Fix tests --- drizzle-kit/tests/push/pg.test.ts | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/drizzle-kit/tests/push/pg.test.ts b/drizzle-kit/tests/push/pg.test.ts index e3ed739f0..fdcad88cb 100644 --- a/drizzle-kit/tests/push/pg.test.ts +++ b/drizzle-kit/tests/push/pg.test.ts @@ -3167,6 +3167,7 @@ test('create role', async (t) => { [], false, ['public'], + undefined, { roles: { include: ['manager'] } }, ); @@ -3204,6 +3205,7 @@ test('create role with properties', async (t) => { [], false, ['public'], + undefined, { roles: { include: ['manager'] } }, ); @@ -3241,6 +3243,7 @@ test('create role with some properties', async (t) => { [], false, ['public'], + undefined, { roles: { include: ['manager'] } }, ); @@ -3276,6 +3279,7 @@ test('drop role', async (t) => { [], false, ['public'], + undefined, { roles: { include: ['manager'] } }, ); @@ -3310,6 +3314,7 @@ test('create and drop role', async (t) => { [], false, ['public'], + undefined, { roles: { include: ['manager', 'admin'] } }, ); @@ -3353,6 +3358,7 @@ test('rename role', async (t) => { ['manager->admin'], false, ['public'], + undefined, { roles: { include: ['manager', 'admin'] } }, ); @@ -3384,6 +3390,7 @@ test('alter all role field', async (t) => { [], false, ['public'], + undefined, { roles: { include: ['manager'] } }, ); @@ -3423,6 +3430,7 @@ test('alter createdb in role', async (t) => { [], false, ['public'], + undefined, { roles: { include: ['manager'] } }, ); @@ -3462,6 +3470,7 @@ test('alter createrole in role', async (t) => { [], false, ['public'], + undefined, { roles: { include: ['manager'] } }, ); @@ -3501,6 +3510,7 @@ test('alter inherit in role', async (t) => { [], false, ['public'], + undefined, { roles: { include: ['manager'] } }, ); From 662aa3587527b55c093fbd2a331a9077da581260 Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Fri, 11 Oct 2024 13:02:48 +0300 Subject: [PATCH 18/24] Fix serializer and tests --- drizzle-kit/src/serializer/pgSerializer.ts | 6 ++---- drizzle-orm/src/pg-core/utils.ts | 2 +- integration-tests/tests/pg/pg-common.ts | 18 +++++++++--------- 3 files changed, 12 insertions(+), 14 deletions(-) diff --git a/drizzle-kit/src/serializer/pgSerializer.ts b/drizzle-kit/src/serializer/pgSerializer.ts index 7132013e8..06897d2d0 100644 --- a/drizzle-kit/src/serializer/pgSerializer.ts +++ b/drizzle-kit/src/serializer/pgSerializer.ts @@ -880,7 +880,7 @@ export const fromDatabase = async ( tablename: string; name: string; as: string; - to: string; + to: string[]; for: string; using: string; withCheck: string; @@ -893,9 +893,7 @@ export const fromDatabase = async ( const { tablename, schemaname, to, withCheck, using, ...rest } = dbPolicy; const tableForPolicy = policiesByTable[`${schemaname}.${tablename}`]; - const parsedTo = to === '{}' - ? [] - : to.substring(1, to.length - 1).split(/\s*,\s*/g).sort(); + const parsedTo = to; const parsedWithCheck = withCheck === null ? undefined : withCheck; const parsedUsing = using === null ? undefined : using; diff --git a/drizzle-orm/src/pg-core/utils.ts b/drizzle-orm/src/pg-core/utils.ts index d7e6c5d8f..8ab5e868f 100644 --- a/drizzle-orm/src/pg-core/utils.ts +++ b/drizzle-orm/src/pg-core/utils.ts @@ -28,7 +28,7 @@ export function getTableConfig(table: TTable) { if (extraConfigBuilder !== undefined) { const extraConfig = extraConfigBuilder(table[Table.Symbol.ExtraConfigColumns]); - const extraValues = Array.isArray(extraConfig) ? extraConfig as any[] : Object.values(extraConfig); + const extraValues = Array.isArray(extraConfig) ? extraConfig.flat(1) as any[] : Object.values(extraConfig); for (const builder of extraValues) { if (is(builder, IndexBuilder)) { indexes.push(builder.build(table)); diff --git a/integration-tests/tests/pg/pg-common.ts b/integration-tests/tests/pg/pg-common.ts index 65671b1ba..0b558d49b 100644 --- a/integration-tests/tests/pg/pg-common.ts +++ b/integration-tests/tests/pg/pg-common.ts @@ -4788,23 +4788,23 @@ export function tests() { for (const it of Object.values(policy)) { expect(is(it, PgPolicy)).toBe(true); expect(it.to).toStrictEqual(authenticatedRole); - expect(it.using).toStrictEqual(sql`select true`); - it.withCheck ? expect(it.withCheck).toStrictEqual(sql`select true`) : ''; + expect(it.using).toStrictEqual(sql`true`); + it.withCheck ? expect(it.withCheck).toStrictEqual(sql`true`) : ''; } } { const table = pgTable('name', { id: integer('id'), - }, (t) => ({ - in: index('name').on(t.id), - neonPolicy: crudPolicy({ + }, (t) => [ + index('name').on(t.id), + crudPolicy({ read: true, modify: true, role: authenticatedRole, }), - pk: primaryKey({ columns: [t.id], name: 'custom' }), - })); + primaryKey({ columns: [t.id], name: 'custom' }), + ]); const { policies, indexes, primaryKeys } = getTableConfig(table); @@ -4812,8 +4812,8 @@ export function tests() { expect(indexes.length).toBe(1); expect(primaryKeys.length).toBe(1); - expect(policies[0]?.name === 'crud-policy-modify'); - expect(policies[1]?.name === 'crud-policy-read'); + expect(policies[0]?.name === 'crud-custom-policy-modify'); + expect(policies[1]?.name === 'crud-custom-policy-read'); } }); From ddba14b8be84fbbc33358146e5377524d0fb1c6c Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Fri, 11 Oct 2024 14:53:41 +0300 Subject: [PATCH 19/24] Make pglite 0.2.5 for now --- integration-tests/package.json | 2 +- pnpm-lock.yaml | 9 +++++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/integration-tests/package.json b/integration-tests/package.json index d648d89a2..4d185fb36 100644 --- a/integration-tests/package.json +++ b/integration-tests/package.json @@ -43,7 +43,7 @@ "dependencies": { "@aws-sdk/client-rds-data": "^3.549.0", "@aws-sdk/credential-providers": "^3.549.0", - "@electric-sql/pglite": "^0.2.5", + "@electric-sql/pglite": "0.2.5", "@libsql/client": "^0.10.0", "@miniflare/d1": "^2.14.4", "@miniflare/shared": "^2.14.4", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 6ba658f6a..6ef07b775 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -550,8 +550,8 @@ importers: specifier: ^3.549.0 version: 3.569.0(@aws-sdk/client-sso-oidc@3.583.0) '@electric-sql/pglite': - specifier: ^0.2.5 - version: 0.2.11 + specifier: 0.2.5 + version: 0.2.5 '@libsql/client': specifier: ^0.10.0 version: 0.10.0(bufferutil@4.0.8)(utf-8-validate@6.0.3) @@ -1996,6 +1996,9 @@ packages: '@electric-sql/pglite@0.2.11': resolution: {integrity: sha512-CeCsbtVqccz0MDlEeT6p4xfgKMpq3GvGTKkswuZ2Vv4r40ZQlE0byl3QNpXcpiTxfye+VFt9RZ+lSL+Wq+jhLQ==} + '@electric-sql/pglite@0.2.5': + resolution: {integrity: sha512-LrMX2kX0mCVN4xkhIDv1KBVukWtoOI/+P9MDQgHX5QEeZCi5S60LZOa0VWXjufPEz7mJtbuXWJRujD++t0gsHA==} + '@esbuild-kit/core-utils@3.1.0': resolution: {integrity: sha512-Uuk8RpCg/7fdHSceR1M6XbSZFSuMrxcePFuGgyvsBn+u339dk5OeL4jv2EojwTN2st/unJGsVm4qHWjWNmJ/tw==} @@ -12495,6 +12498,8 @@ snapshots: '@electric-sql/pglite@0.2.11': {} + '@electric-sql/pglite@0.2.5': {} + '@esbuild-kit/core-utils@3.1.0': dependencies: esbuild: 0.17.19 From a5ec472eb3046af512d5ee6b980fdf5630ff19dd Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Fri, 11 Oct 2024 15:43:24 +0300 Subject: [PATCH 20/24] Update deps --- drizzle-orm/package.json | 2 +- pnpm-lock.yaml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/drizzle-orm/package.json b/drizzle-orm/package.json index c448b1396..52923bd6e 100644 --- a/drizzle-orm/package.json +++ b/drizzle-orm/package.json @@ -160,7 +160,7 @@ "devDependencies": { "@aws-sdk/client-rds-data": "^3.549.0", "@cloudflare/workers-types": "^4.20230904.0", - "@electric-sql/pglite": "^0.2.5", + "@electric-sql/pglite": "0.2.5", "@libsql/client": "^0.10.0", "@miniflare/d1": "^2.14.4", "@neondatabase/serverless": "^0.9.0", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 6ef07b775..9b5332d6a 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -298,8 +298,8 @@ importers: specifier: ^4.20230904.0 version: 4.20240512.0 '@electric-sql/pglite': - specifier: ^0.2.5 - version: 0.2.11 + specifier: 0.2.5 + version: 0.2.5 '@libsql/client': specifier: ^0.10.0 version: 0.10.0(bufferutil@4.0.8)(utf-8-validate@6.0.3) From 35471bd8b38458e30de8f5721fe718f04a6af463 Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Mon, 21 Oct 2024 19:40:49 +0300 Subject: [PATCH 21/24] Add enableRLS() --- drizzle-kit/src/cli/commands/pgPushUtils.ts | 8 +----- drizzle-kit/src/serializer/pgSchema.ts | 1 + drizzle-kit/src/serializer/pgSerializer.ts | 21 +++++++++++++--- drizzle-kit/src/snapshotsDiffer.ts | 6 ++--- drizzle-kit/src/sqlgenerator.ts | 2 +- drizzle-kit/tests/pg-checks.test.ts | 1 + drizzle-kit/tests/pg-views.test.ts | 9 +++++++ drizzle-kit/tests/push/pg.test.ts | 6 ++++- drizzle-kit/tests/rls/pg-policy.test.ts | 3 ++- drizzle-kit/tests/schemaDiffer.ts | 2 ++ drizzle-orm/src/neon/rls.ts | 28 ++++++++++----------- drizzle-orm/src/pg-core/table.ts | 24 +++++++++++++++++- drizzle-orm/src/pg-core/utils.ts | 2 ++ integration-tests/tests/pg/pg-common.ts | 17 +++++++++++++ 14 files changed, 98 insertions(+), 32 deletions(-) diff --git a/drizzle-kit/src/cli/commands/pgPushUtils.ts b/drizzle-kit/src/cli/commands/pgPushUtils.ts index a8e2570df..fcc1ad7f3 100644 --- a/drizzle-kit/src/cli/commands/pgPushUtils.ts +++ b/drizzle-kit/src/cli/commands/pgPushUtils.ts @@ -252,13 +252,7 @@ export const pgSuggestions = async (db: DB, statements: JsonStatement[]) => { } const stmnt = fromJson([statement], 'postgresql'); if (typeof stmnt !== 'undefined') { - if (statement.type === 'drop_table') { - statementsToExecute.push( - `DROP TABLE ${concatSchemaAndTableName(statement.schema, statement.tableName)} CASCADE;`, - ); - } else { - statementsToExecute.push(...stmnt); - } + statementsToExecute.push(...stmnt); } } diff --git a/drizzle-kit/src/serializer/pgSchema.ts b/drizzle-kit/src/serializer/pgSchema.ts index c342e1eb3..677aae629 100644 --- a/drizzle-kit/src/serializer/pgSchema.ts +++ b/drizzle-kit/src/serializer/pgSchema.ts @@ -331,6 +331,7 @@ const table = object({ uniqueConstraints: record(string(), uniqueConstraint).default({}), policies: record(string(), policy).default({}), checkConstraints: record(string(), checkConstraint).default({}), + isRLSEnabled: boolean().default(false), }).strict(); const schemaHash = object({ diff --git a/drizzle-kit/src/serializer/pgSerializer.ts b/drizzle-kit/src/serializer/pgSerializer.ts index a6eda1bcf..98848918a 100644 --- a/drizzle-kit/src/serializer/pgSerializer.ts +++ b/drizzle-kit/src/serializer/pgSerializer.ts @@ -124,8 +124,18 @@ export const generatePgSnapshot = ( // within the same PostgreSQL table const checksInTable: Record = {}; - const { name: tableName, columns, indexes, foreignKeys, checks, schema, primaryKeys, uniqueConstraints, policies } = - getTableConfig(table); + const { + name: tableName, + columns, + indexes, + foreignKeys, + checks, + schema, + primaryKeys, + uniqueConstraints, + policies, + enableRLS, + } = getTableConfig(table); if (schemaFilter && !schemaFilter.includes(schema ?? 'public')) { continue; @@ -545,6 +555,7 @@ export const generatePgSnapshot = ( uniqueConstraints: uniqueConstraintObject, policies: policiesObject, checkConstraints: checksObject, + isRLSEnabled: enableRLS, }; } @@ -869,7 +880,7 @@ export const fromDatabase = async ( const where = schemaFilters.map((t) => `n.nspname = '${t}'`).join(' or '); - const allTables = await db.query<{ table_schema: string; table_name: string; type: string }>( + const allTables = await db.query<{ table_schema: string; table_name: string; type: string; rls_enabled: boolean }>( `SELECT n.nspname AS table_schema, c.relname AS table_name, @@ -877,7 +888,8 @@ export const fromDatabase = async ( WHEN c.relkind = 'r' THEN 'table' WHEN c.relkind = 'v' THEN 'view' WHEN c.relkind = 'm' THEN 'materialized_view' - END AS type + END AS type, + c.relrowsecurity AS rls_enabled FROM pg_catalog.pg_class c JOIN @@ -1532,6 +1544,7 @@ WHERE uniqueConstraints: uniqueConstrains, checkConstraints: checkConstraints, policies: policiesByTable[`${tableSchema}.${tableName}`] ?? {}, + isRLSEnabled: row.rls_enabled, }; } catch (e) { rej(e); diff --git a/drizzle-kit/src/snapshotsDiffer.ts b/drizzle-kit/src/snapshotsDiffer.ts index 2a6d98d6c..69c998c53 100644 --- a/drizzle-kit/src/snapshotsDiffer.ts +++ b/drizzle-kit/src/snapshotsDiffer.ts @@ -923,7 +923,7 @@ export const applyPgSnapshotsDiff = async ( //// Policies - const policyRes = diffPolicies(columnsPatchedSnap1.tables, json2.tables); + const policyRes = diffPolicies(tablesPatchedSnap1.tables, json2.tables); const policyRenames = [] as { table: string; @@ -990,7 +990,7 @@ export const applyPgSnapshotsDiff = async ( >, ); - const policyPatchedSnap1 = copy(columnsPatchedSnap1); + const policyPatchedSnap1 = copy(tablesPatchedSnap1); policyPatchedSnap1.tables = mapEntries( policyPatchedSnap1.tables, (tableKey, tableValue) => { @@ -1015,7 +1015,7 @@ export const applyPgSnapshotsDiff = async ( ); //// - const viewsDiff = diffSchemasOrTables(json1.views, json2.views); + const viewsDiff = diffSchemasOrTables(policyPatchedSnap1.views, json2.views); const { created: createdViews, diff --git a/drizzle-kit/src/sqlgenerator.ts b/drizzle-kit/src/sqlgenerator.ts index 6c023325d..5033c143a 100644 --- a/drizzle-kit/src/sqlgenerator.ts +++ b/drizzle-kit/src/sqlgenerator.ts @@ -1355,7 +1355,7 @@ class PgDropTableConvertor extends Convertor { return [ ...droppedPolicies, - `DROP TABLE ${tableNameWithSchema};`, + `DROP TABLE ${tableNameWithSchema} CASCADE;`, ]; } } diff --git a/drizzle-kit/tests/pg-checks.test.ts b/drizzle-kit/tests/pg-checks.test.ts index 1f5e5e1c5..a87b5d115 100644 --- a/drizzle-kit/tests/pg-checks.test.ts +++ b/drizzle-kit/tests/pg-checks.test.ts @@ -39,6 +39,7 @@ test('create table with check', async (t) => { checkConstraints: ['some_check_name;"users"."age" > 21'], compositePkName: '', uniqueConstraints: [], + policies: [], } as JsonCreateTableStatement); expect(sqlStatements.length).toBe(1); diff --git a/drizzle-kit/tests/pg-views.test.ts b/drizzle-kit/tests/pg-views.test.ts index 2874caf5b..39b4ff5be 100644 --- a/drizzle-kit/tests/pg-views.test.ts +++ b/drizzle-kit/tests/pg-views.test.ts @@ -29,6 +29,7 @@ test('create table and view #1', async () => { uniqueConstraints: [], compositePkName: '', checkConstraints: [], + policies: [], }); expect(statements[1]).toStrictEqual({ type: 'create_view', @@ -74,6 +75,7 @@ test('create table and view #2', async () => { compositePKs: [], uniqueConstraints: [], compositePkName: '', + policies: [], checkConstraints: [], }); expect(statements[1]).toStrictEqual({ @@ -130,6 +132,7 @@ test('create table and view #3', async () => { uniqueConstraints: [], compositePkName: '', checkConstraints: [], + policies: [], }); expect(statements[1]).toStrictEqual({ type: 'create_view', @@ -215,6 +218,7 @@ test('create table and view #4', async () => { compositePKs: [], uniqueConstraints: [], compositePkName: '', + policies: [], checkConstraints: [], }); expect(statements[2]).toStrictEqual({ @@ -302,6 +306,7 @@ test('create table and view #6', async () => { type: 'create_table', uniqueConstraints: [], checkConstraints: [], + policies: [], }); expect(statements[1]).toStrictEqual({ definition: 'SELECT * FROM "users"', @@ -370,6 +375,7 @@ test('create table and materialized view #1', async () => { }], compositePKs: [], uniqueConstraints: [], + policies: [], compositePkName: '', checkConstraints: [], }); @@ -417,6 +423,7 @@ test('create table and materialized view #2', async () => { compositePKs: [], uniqueConstraints: [], compositePkName: '', + policies: [], checkConstraints: [], }); expect(statements[1]).toStrictEqual({ @@ -483,6 +490,7 @@ test('create table and materialized view #3', async () => { compositePKs: [], uniqueConstraints: [], compositePkName: '', + policies: [], checkConstraints: [], }); expect(statements[1]).toStrictEqual({ @@ -582,6 +590,7 @@ test('create table and materialized view #5', async () => { tableName: 'users', type: 'create_table', uniqueConstraints: [], + policies: [], checkConstraints: [], }); expect(statements[1]).toEqual({ diff --git a/drizzle-kit/tests/push/pg.test.ts b/drizzle-kit/tests/push/pg.test.ts index d01b435ac..1be6cd666 100644 --- a/drizzle-kit/tests/push/pg.test.ts +++ b/drizzle-kit/tests/push/pg.test.ts @@ -2363,6 +2363,7 @@ test('drop mat view with data', async () => { false, ['public'], undefined, + undefined, { after: seedStatements }, ); @@ -2470,6 +2471,7 @@ test('drop view with data', async () => { false, ['public'], undefined, + undefined, { after: seedStatements }, ); @@ -2551,6 +2553,7 @@ test('enums ordering', async () => { false, ['public'], undefined, + undefined, { before: [...createEnum, ...addedValueSql], runApply: false }, ); @@ -3446,6 +3449,7 @@ test('create table with a policy', async (t) => { type: 'integer', }, ], + checkConstraints: [], compositePKs: [], compositePkName: '', policies: [ @@ -3502,7 +3506,7 @@ test('drop table with a policy', async (t) => { expect(sqlStatements).toStrictEqual([ 'DROP POLICY "test" ON "users2" CASCADE;', - 'DROP TABLE "users2";', + 'DROP TABLE "users2" CASCADE;', ]); expect(statements).toStrictEqual([ { diff --git a/drizzle-kit/tests/rls/pg-policy.test.ts b/drizzle-kit/tests/rls/pg-policy.test.ts index f6fec3f60..46f05a77f 100644 --- a/drizzle-kit/tests/rls/pg-policy.test.ts +++ b/drizzle-kit/tests/rls/pg-policy.test.ts @@ -590,6 +590,7 @@ test('create table with a policy', async (t) => { }, ], compositePKs: [], + checkConstraints: [], compositePkName: '', policies: [ 'test--PERMISSIVE--ALL--public--undefined--undefined', @@ -632,7 +633,7 @@ test('drop table with a policy', async (t) => { expect(sqlStatements).toStrictEqual([ 'DROP POLICY "test" ON "users2" CASCADE;', - 'DROP TABLE "users2";', + 'DROP TABLE "users2" CASCADE;', ]); expect(statements).toStrictEqual([ { diff --git a/drizzle-kit/tests/schemaDiffer.ts b/drizzle-kit/tests/schemaDiffer.ts index e22a4e6b5..d8f6efc51 100644 --- a/drizzle-kit/tests/schemaDiffer.ts +++ b/drizzle-kit/tests/schemaDiffer.ts @@ -11,6 +11,7 @@ import { isPgView, PgEnum, PgMaterializedView, + PgRole, PgSchema, PgSequence, PgTable, @@ -794,6 +795,7 @@ export const diffTestSchemasPush = async ( const leftSequences = Object.values(right).filter((it) => isPgSequence(it)) as PgSequence[]; const leftRoles = Object.values(right).filter((it) => is(it, PgRole)) as PgRole[]; + const leftViews = Object.values(right).filter((it) => isPgView(it)) as PgView[]; const leftMaterializedViews = Object.values(right).filter((it) => isPgMaterializedView(it)) as PgMaterializedView[]; diff --git a/drizzle-orm/src/neon/rls.ts b/drizzle-orm/src/neon/rls.ts index 9bb421434..0645b1d01 100644 --- a/drizzle-orm/src/neon/rls.ts +++ b/drizzle-orm/src/neon/rls.ts @@ -3,13 +3,11 @@ import { pgPolicy, PgRole, pgRole } from '~/pg-core/index.ts'; import type { AnyPgColumn, PgPolicyToOption } from '~/pg-core/index.ts'; import { type SQL, sql } from '~/sql/sql.ts'; -export const crudPolicy = ( - options: { - role: PgPolicyToOption; - read?: SQL | boolean; - modify?: SQL | boolean; - }, -) => { +export const crudPolicy = (options: { + role: PgPolicyToOption; + read?: SQL | boolean; + modify?: SQL | boolean; +}) => { const read: SQL = options.read === true ? sql`true` : options.read === false || options.read === undefined @@ -24,18 +22,21 @@ export const crudPolicy = ( let rolesName = ''; if (Array.isArray(options.role)) { - rolesName = options.role.map((it) => { - return is(it, PgRole) ? it.name : it as string; - }).join('-'); + rolesName = options.role + .map((it) => { + return is(it, PgRole) ? it.name : (it as string); + }) + .join('-'); } else { - rolesName = is(options.role, PgRole) ? options.role.name : options.role as string; + rolesName = is(options.role, PgRole) + ? options.role.name + : (options.role as string); } return [ pgPolicy(`crud-${rolesName}-policy-insert`, { for: 'insert', to: options.role, - using: modify, withCheck: modify, }), pgPolicy(`crud-${rolesName}-policy-update`, { @@ -48,7 +49,6 @@ export const crudPolicy = ( for: 'delete', to: options.role, using: modify, - withCheck: modify, }), pgPolicy(`crud-${rolesName}-policy-select`, { for: 'select', @@ -62,4 +62,4 @@ export const crudPolicy = ( export const authenticatedRole = pgRole('authenticated').existing(); export const anonymousRole = pgRole('anonymous').existing(); -export const authUid = (userIdColumn: AnyPgColumn) => sql`select auth.user_id() = ${userIdColumn}`; +export const authUid = (userIdColumn: AnyPgColumn) => sql`(select auth.user_id() = ${userIdColumn})`; diff --git a/drizzle-orm/src/pg-core/table.ts b/drizzle-orm/src/pg-core/table.ts index 5b463a681..c3c34c577 100644 --- a/drizzle-orm/src/pg-core/table.ts +++ b/drizzle-orm/src/pg-core/table.ts @@ -27,6 +27,8 @@ export type TableConfig = TableConfigBase; /** @internal */ export const InlineForeignKeys = Symbol.for('drizzle:PgInlineForeignKeys'); +/** @internal */ +export const EnableRLS = Symbol.for('drizzle:EnableRLS'); export class PgTable extends Table { static override readonly [entityKind]: string = 'PgTable'; @@ -34,11 +36,15 @@ export class PgTable extends Table { /** @internal */ static override readonly Symbol = Object.assign({}, Table.Symbol, { InlineForeignKeys: InlineForeignKeys as typeof InlineForeignKeys, + EnableRLS: EnableRLS as typeof EnableRLS, }); /**@internal */ [InlineForeignKeys]: ForeignKey[] = []; + /** @internal */ + [EnableRLS]: boolean = false; + /** @internal */ override [Table.Symbol.ExtraConfigBuilder]: ((self: Record) => PgTableExtraConfig) | undefined = undefined; @@ -50,6 +56,12 @@ export type PgTableWithColumns = & PgTable & { [Key in keyof T['columns']]: T['columns'][Key]; + } + & { + enableRLS: () => Omit< + PgTableWithColumns, + 'enableRLS' + >; }; /** @internal */ @@ -108,7 +120,17 @@ export function pgTableWithSchema< table[PgTable.Symbol.ExtraConfigBuilder] = extraConfig as any; } - return table; + return Object.assign(table, { + enableRLS: () => { + table[PgTable.Symbol.EnableRLS] = true; + return table as PgTableWithColumns<{ + name: TTableName; + schema: TSchemaName; + columns: BuildColumns; + dialect: 'pg'; + }>; + }, + }); } export interface PgTableFn { diff --git a/drizzle-orm/src/pg-core/utils.ts b/drizzle-orm/src/pg-core/utils.ts index 8ab5e868f..0191c2439 100644 --- a/drizzle-orm/src/pg-core/utils.ts +++ b/drizzle-orm/src/pg-core/utils.ts @@ -23,6 +23,7 @@ export function getTableConfig(table: TTable) { const name = table[Table.Symbol.Name]; const schema = table[Table.Symbol.Schema]; const policies: PgPolicy[] = []; + const enableRLS: boolean = table[PgTable.Symbol.EnableRLS]; const extraConfigBuilder = table[PgTable.Symbol.ExtraConfigBuilder]; @@ -56,6 +57,7 @@ export function getTableConfig(table: TTable) { name, schema, policies, + enableRLS, }; } diff --git a/integration-tests/tests/pg/pg-common.ts b/integration-tests/tests/pg/pg-common.ts index 0b558d49b..1ee566477 100644 --- a/integration-tests/tests/pg/pg-common.ts +++ b/integration-tests/tests/pg/pg-common.ts @@ -4817,6 +4817,23 @@ export function tests() { } }); + test('Enable RLS function', () => { + const usersWithRLS = pgTable('users', { + id: integer(), + }).enableRLS(); + + const config1 = getTableConfig(usersWithRLS); + + const usersNoRLS = pgTable('users', { + id: integer(), + }); + + const config2 = getTableConfig(usersNoRLS); + + expect(config1.enableRLS).toBeTruthy(); + expect(config2.enableRLS).toBeFalsy(); + }); + test('$count separate', async (ctx) => { const { db } = ctx.pg; From 31c804126ec3a7dbeec0005f5e918d11a1935375 Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Wed, 23 Oct 2024 14:10:11 +0300 Subject: [PATCH 22/24] Add enableRLS function for pg tables --- drizzle-kit/src/cli/commands/pgPushUtils.ts | 2 +- drizzle-kit/src/jsonStatements.ts | 5 +- drizzle-kit/src/serializer/pgSchema.ts | 2 + drizzle-kit/src/snapshotsDiffer.ts | 31 +++--- drizzle-kit/src/sqlgenerator.ts | 5 +- drizzle-kit/tests/pg-checks.test.ts | 1 + drizzle-kit/tests/pg-identity.test.ts | 3 + drizzle-kit/tests/pg-tables.test.ts | 10 ++ drizzle-kit/tests/pg-views.test.ts | 9 ++ drizzle-kit/tests/push/pg.test.ts | 10 ++ drizzle-kit/tests/rls/pg-policy.test.ts | 107 ++++++++++++++++++++ 11 files changed, 165 insertions(+), 20 deletions(-) diff --git a/drizzle-kit/src/cli/commands/pgPushUtils.ts b/drizzle-kit/src/cli/commands/pgPushUtils.ts index fcc1ad7f3..b53fec3e7 100644 --- a/drizzle-kit/src/cli/commands/pgPushUtils.ts +++ b/drizzle-kit/src/cli/commands/pgPushUtils.ts @@ -257,7 +257,7 @@ export const pgSuggestions = async (db: DB, statements: JsonStatement[]) => { } return { - statementsToExecute, + statementsToExecute: [...new Set(statementsToExecute)], shouldAskForApprove, infoToPrint, matViewsToRemove: [...new Set(matViewsToRemove)], diff --git a/drizzle-kit/src/jsonStatements.ts b/drizzle-kit/src/jsonStatements.ts index b36e7d490..0d248422d 100644 --- a/drizzle-kit/src/jsonStatements.ts +++ b/drizzle-kit/src/jsonStatements.ts @@ -51,6 +51,7 @@ export interface JsonCreateTableStatement { policies?: string[]; checkConstraints?: string[]; internals?: MySqlKitInternals; + isRLSEnabled?: boolean; } export interface JsonRecreateTableStatement { @@ -822,7 +823,8 @@ export const preparePgCreateTableJson = ( // TODO: remove? json2: PgSchema, ): JsonCreateTableStatement => { - const { name, schema, columns, compositePrimaryKeys, uniqueConstraints, checkConstraints, policies } = table; + const { name, schema, columns, compositePrimaryKeys, uniqueConstraints, checkConstraints, policies, isRLSEnabled } = + table; const tableKey = `${schema || 'public'}.${name}`; // TODO: @AndriiSherman. We need this, will add test cases @@ -842,6 +844,7 @@ export const preparePgCreateTableJson = ( uniqueConstraints: Object.values(uniqueConstraints), policies: Object.values(policies), checkConstraints: Object.values(checkConstraints), + isRLSEnabled: isRLSEnabled ?? false, }; }; diff --git a/drizzle-kit/src/serializer/pgSchema.ts b/drizzle-kit/src/serializer/pgSchema.ts index 677aae629..52c5a192a 100644 --- a/drizzle-kit/src/serializer/pgSchema.ts +++ b/drizzle-kit/src/serializer/pgSchema.ts @@ -455,6 +455,7 @@ const tableSquashed = object({ uniqueConstraints: record(string(), string()), policies: record(string(), string()), checkConstraints: record(string(), string()), + isRLSEnabled: boolean().default(false), }).strict(); const tableSquashedV4 = object({ @@ -802,6 +803,7 @@ export const squashPgScheme = ( uniqueConstraints: squashedUniqueConstraints, policies: squashedPolicies, checkConstraints: squashedChecksContraints, + isRLSEnabled: it[1].isRLSEnabled ?? false, }, ]; }), diff --git a/drizzle-kit/src/snapshotsDiffer.ts b/drizzle-kit/src/snapshotsDiffer.ts index 69c998c53..3818d6ac9 100644 --- a/drizzle-kit/src/snapshotsDiffer.ts +++ b/drizzle-kit/src/snapshotsDiffer.ts @@ -256,6 +256,7 @@ const tableScheme = object({ uniqueConstraints: record(string(), string()).default({}), policies: record(string(), string()).default({}), checkConstraints: record(string(), string()).default({}), + isRLSEnabled: boolean().default(false), }).strict(); export const alteredTableScheme = object({ @@ -1370,9 +1371,6 @@ export const applyPgSnapshotsDiff = async ( ); }); - // if there were no tables in json1 and is a table in json2 and it has policies -> enable - // if there was a table in json1 and is no table in json2 and it had policies -> disable - // Handle enabling and disabling RLS for (const table of Object.values(json2.tables)) { const policiesInCurrentState = Object.keys(table.policies); @@ -1380,33 +1378,34 @@ export const applyPgSnapshotsDiff = async ( columnsPatchedSnap1.tables[`${table.schema === '' ? 'public' : table.schema}.${table.name}`]; const policiesInPreviousState = tableInPreviousState ? Object.keys(tableInPreviousState.policies) : []; - if (policiesInPreviousState.length === 0 && policiesInCurrentState.length > 0) { + if (policiesInPreviousState.length === 0 && policiesInCurrentState.length > 0 && !table.isRLSEnabled) { jsonEnableRLSStatements.push({ type: 'enable_rls', tableName: table.name, schema: table.schema }); } - if (policiesInPreviousState.length > 0 && policiesInCurrentState.length === 0) { + if (policiesInPreviousState.length > 0 && policiesInCurrentState.length === 0 && !table.isRLSEnabled) { jsonDisableRLSStatements.push({ type: 'disable_rls', tableName: table.name, schema: table.schema }); } + + // handle table.isRLSEnabled + if (table.isRLSEnabled !== tableInPreviousState.isRLSEnabled) { + if (table.isRLSEnabled) { + // was force enabled + jsonEnableRLSStatements.push({ type: 'enable_rls', tableName: table.name, schema: table.schema }); + } else if (!table.isRLSEnabled && policiesInCurrentState.length === 0) { + // was force disabled + jsonDisableRLSStatements.push({ type: 'disable_rls', tableName: table.name, schema: table.schema }); + } + } } for (const table of Object.values(columnsPatchedSnap1.tables)) { const tableInCurrentState = json2.tables[`${table.schema === '' ? 'public' : table.schema}.${table.name}`]; - if (tableInCurrentState === undefined) { + if (tableInCurrentState === undefined && !table.isRLSEnabled) { jsonDisableRLSStatements.push({ type: 'disable_rls', tableName: table.name, schema: table.schema }); } } - // TODO - // Add to sql generators - // Test generate logic manually - // Add tests - // Add introspect - // add introspect-pg.ts - think about introspecting from supabase and neon - // add push logic - think about using with neon and supabase - // add push and introspect tests - // beta release - // handle indexes const droppedIndexes = Object.keys(it.alteredIndexes).reduce( (current, item: string) => { diff --git a/drizzle-kit/src/sqlgenerator.ts b/drizzle-kit/src/sqlgenerator.ts index 5033c143a..4b75f4286 100644 --- a/drizzle-kit/src/sqlgenerator.ts +++ b/drizzle-kit/src/sqlgenerator.ts @@ -322,7 +322,8 @@ class PgCreateTableConvertor extends Convertor { } convert(st: JsonCreateTableStatement) { - const { tableName, schema, columns, compositePKs, uniqueConstraints, checkConstraints, policies } = st; + const { tableName, schema, columns, compositePKs, uniqueConstraints, checkConstraints, policies, isRLSEnabled } = + st; let statement = ''; const name = schema ? `"${schema}"."${tableName}"` : `"${tableName}"`; @@ -425,7 +426,7 @@ class PgCreateTableConvertor extends Convertor { schema, }); - return [statement, ...(policies && policies.length > 0 ? [enableRls] : [])]; + return [statement, ...(policies && policies.length > 0 || isRLSEnabled ? [enableRls] : [])]; } } diff --git a/drizzle-kit/tests/pg-checks.test.ts b/drizzle-kit/tests/pg-checks.test.ts index a87b5d115..50a01a6c1 100644 --- a/drizzle-kit/tests/pg-checks.test.ts +++ b/drizzle-kit/tests/pg-checks.test.ts @@ -39,6 +39,7 @@ test('create table with check', async (t) => { checkConstraints: ['some_check_name;"users"."age" > 21'], compositePkName: '', uniqueConstraints: [], + isRLSEnabled: false, policies: [], } as JsonCreateTableStatement); diff --git a/drizzle-kit/tests/pg-identity.test.ts b/drizzle-kit/tests/pg-identity.test.ts index 13816fbf3..9f6ce8ba7 100644 --- a/drizzle-kit/tests/pg-identity.test.ts +++ b/drizzle-kit/tests/pg-identity.test.ts @@ -46,6 +46,7 @@ test('create table: identity always/by default - no params', async () => { compositePkName: '', schema: '', policies: [], + isRLSEnabled: false, tableName: 'users', type: 'create_table', uniqueConstraints: [], @@ -86,6 +87,7 @@ test('create table: identity always/by default - few params', async () => { compositePkName: '', policies: [], schema: '', + isRLSEnabled: false, tableName: 'users', type: 'create_table', uniqueConstraints: [], @@ -129,6 +131,7 @@ test('create table: identity always/by default - all params', async () => { compositePKs: [], compositePkName: '', policies: [], + isRLSEnabled: false, schema: '', tableName: 'users', type: 'create_table', diff --git a/drizzle-kit/tests/pg-tables.test.ts b/drizzle-kit/tests/pg-tables.test.ts index f4816c823..1f2885f92 100644 --- a/drizzle-kit/tests/pg-tables.test.ts +++ b/drizzle-kit/tests/pg-tables.test.ts @@ -37,6 +37,7 @@ test('add table #1', async () => { policies: [], uniqueConstraints: [], checkConstraints: [], + isRLSEnabled: false, compositePkName: '', }); }); @@ -64,6 +65,7 @@ test('add table #2', async () => { }, ], compositePKs: [], + isRLSEnabled: false, policies: [], uniqueConstraints: [], checkConstraints: [], @@ -107,6 +109,7 @@ test('add table #3', async () => { compositePKs: ['id;users_pk'], policies: [], uniqueConstraints: [], + isRLSEnabled: false, checkConstraints: [], compositePkName: 'users_pk', }); @@ -130,6 +133,7 @@ test('add table #4', async () => { policies: [], uniqueConstraints: [], checkConstraints: [], + isRLSEnabled: false, compositePkName: '', }); expect(statements[1]).toStrictEqual({ @@ -139,6 +143,7 @@ test('add table #4', async () => { schema: '', columns: [], compositePKs: [], + isRLSEnabled: false, uniqueConstraints: [], checkConstraints: [], compositePkName: '', @@ -169,6 +174,7 @@ test('add table #5', async () => { uniqueConstraints: [], compositePkName: '', checkConstraints: [], + isRLSEnabled: false, }); }); @@ -194,6 +200,7 @@ test('add table #6', async () => { policies: [], compositePkName: '', checkConstraints: [], + isRLSEnabled: false, }); expect(statements[1]).toStrictEqual({ type: 'drop_table', @@ -227,6 +234,7 @@ test('add table #7', async () => { policies: [], uniqueConstraints: [], compositePkName: '', + isRLSEnabled: false, checkConstraints: [], }); expect(statements[1]).toStrictEqual({ @@ -284,6 +292,7 @@ test('multiproject schema add table #1', async () => { compositePKs: [], policies: [], compositePkName: '', + isRLSEnabled: false, uniqueConstraints: [], checkConstraints: [], }); @@ -379,6 +388,7 @@ test('add schema + table #1', async () => { policies: [], columns: [], compositePKs: [], + isRLSEnabled: false, uniqueConstraints: [], compositePkName: '', checkConstraints: [], diff --git a/drizzle-kit/tests/pg-views.test.ts b/drizzle-kit/tests/pg-views.test.ts index 39b4ff5be..002004c47 100644 --- a/drizzle-kit/tests/pg-views.test.ts +++ b/drizzle-kit/tests/pg-views.test.ts @@ -27,6 +27,7 @@ test('create table and view #1', async () => { }], compositePKs: [], uniqueConstraints: [], + isRLSEnabled: false, compositePkName: '', checkConstraints: [], policies: [], @@ -74,6 +75,7 @@ test('create table and view #2', async () => { }], compositePKs: [], uniqueConstraints: [], + isRLSEnabled: false, compositePkName: '', policies: [], checkConstraints: [], @@ -132,6 +134,7 @@ test('create table and view #3', async () => { uniqueConstraints: [], compositePkName: '', checkConstraints: [], + isRLSEnabled: false, policies: [], }); expect(statements[1]).toStrictEqual({ @@ -218,6 +221,7 @@ test('create table and view #4', async () => { compositePKs: [], uniqueConstraints: [], compositePkName: '', + isRLSEnabled: false, policies: [], checkConstraints: [], }); @@ -306,6 +310,7 @@ test('create table and view #6', async () => { type: 'create_table', uniqueConstraints: [], checkConstraints: [], + isRLSEnabled: false, policies: [], }); expect(statements[1]).toStrictEqual({ @@ -375,6 +380,7 @@ test('create table and materialized view #1', async () => { }], compositePKs: [], uniqueConstraints: [], + isRLSEnabled: false, policies: [], compositePkName: '', checkConstraints: [], @@ -423,6 +429,7 @@ test('create table and materialized view #2', async () => { compositePKs: [], uniqueConstraints: [], compositePkName: '', + isRLSEnabled: false, policies: [], checkConstraints: [], }); @@ -489,6 +496,7 @@ test('create table and materialized view #3', async () => { }], compositePKs: [], uniqueConstraints: [], + isRLSEnabled: false, compositePkName: '', policies: [], checkConstraints: [], @@ -590,6 +598,7 @@ test('create table and materialized view #5', async () => { tableName: 'users', type: 'create_table', uniqueConstraints: [], + isRLSEnabled: false, policies: [], checkConstraints: [], }); diff --git a/drizzle-kit/tests/push/pg.test.ts b/drizzle-kit/tests/push/pg.test.ts index 1be6cd666..bd4b2840d 100644 --- a/drizzle-kit/tests/push/pg.test.ts +++ b/drizzle-kit/tests/push/pg.test.ts @@ -536,6 +536,7 @@ const pgSuite: DialectSuite = { ], compositePKs: [], compositePkName: '', + isRLSEnabled: false, schema: '', tableName: 'users', policies: [], @@ -1214,6 +1215,7 @@ test('create table: identity always/by default - no params', async () => { policies: [], type: 'create_table', uniqueConstraints: [], + isRLSEnabled: false, checkConstraints: [], }, ]); @@ -1275,6 +1277,7 @@ test('create table: identity always/by default - few params', async () => { schema: '', tableName: 'users', type: 'create_table', + isRLSEnabled: false, uniqueConstraints: [], checkConstraints: [], }, @@ -1343,6 +1346,7 @@ test('create table: identity always/by default - all params', async () => { tableName: 'users', type: 'create_table', policies: [], + isRLSEnabled: false, uniqueConstraints: [], checkConstraints: [], }, @@ -2795,6 +2799,11 @@ test('drop policy', async () => { expect(statements).toStrictEqual([ { type: 'disable_rls', tableName: 'users', schema: '' }, + { + schema: '', + tableName: 'users', + type: 'disable_rls', + }, { type: 'drop_policy', tableName: 'users', @@ -3451,6 +3460,7 @@ test('create table with a policy', async (t) => { ], checkConstraints: [], compositePKs: [], + isRLSEnabled: false, compositePkName: '', policies: [ 'test--PERMISSIVE--ALL--public', diff --git a/drizzle-kit/tests/rls/pg-policy.test.ts b/drizzle-kit/tests/rls/pg-policy.test.ts index 46f05a77f..5ca667faa 100644 --- a/drizzle-kit/tests/rls/pg-policy.test.ts +++ b/drizzle-kit/tests/rls/pg-policy.test.ts @@ -597,6 +597,7 @@ test('create table with a policy', async (t) => { ], schema: '', tableName: 'users2', + isRLSEnabled: false, type: 'create_table', uniqueConstraints: [], }, @@ -692,3 +693,109 @@ test('add policy with multiple "to" roles', async (t) => { }, ]); }); + +test('create table with rls enabled', async (t) => { + const schema1 = {}; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }).enableRLS(), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + `CREATE TABLE IF NOT EXISTS "users" (\n\t"id" integer PRIMARY KEY NOT NULL\n); +`, + 'ALTER TABLE "users" ENABLE ROW LEVEL SECURITY;', + ]); + + console.log(statements); +}); + +test('enable rls force', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }).enableRLS(), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual(['ALTER TABLE "users" ENABLE ROW LEVEL SECURITY;']); +}); + +test('disable rls force', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }).enableRLS(), + }; + + const schema2 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual(['ALTER TABLE "users" DISABLE ROW LEVEL SECURITY;']); +}); + +test('drop policy with enabled rls', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { to: ['current_role', role] }), + })).enableRLS(), + }; + + const role = pgRole('manager').existing(); + + const schema2 = { + role, + users: pgTable('users', { + id: integer('id').primaryKey(), + }).enableRLS(), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'DROP POLICY "test" ON "users" CASCADE;', + ]); +}); + +test('add policy with enabled rls', async (t) => { + const schema1 = { + users: pgTable('users', { + id: integer('id').primaryKey(), + }).enableRLS(), + }; + + const role = pgRole('manager').existing(); + + const schema2 = { + role, + users: pgTable('users', { + id: integer('id').primaryKey(), + }, () => ({ + rls: pgPolicy('test', { to: ['current_role', role] }), + })).enableRLS(), + }; + + const { statements, sqlStatements } = await diffTestSchemas(schema1, schema2, []); + + expect(sqlStatements).toStrictEqual([ + 'CREATE POLICY "test" ON "users" AS PERMISSIVE FOR ALL TO current_role, "manager";', + ]); +}); From 4cb1bdb6e5c528e43fbd7750179231b971792d99 Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Wed, 23 Oct 2024 15:50:28 +0300 Subject: [PATCH 23/24] Fix tests --- integration-tests/tests/pg/pg-common.ts | 2 +- integration-tests/tests/pg/rls/rls.definition.test.ts | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/integration-tests/tests/pg/pg-common.ts b/integration-tests/tests/pg/pg-common.ts index 1ee566477..a079db973 100644 --- a/integration-tests/tests/pg/pg-common.ts +++ b/integration-tests/tests/pg/pg-common.ts @@ -4788,7 +4788,7 @@ export function tests() { for (const it of Object.values(policy)) { expect(is(it, PgPolicy)).toBe(true); expect(it.to).toStrictEqual(authenticatedRole); - expect(it.using).toStrictEqual(sql`true`); + it.using ? expect(it.using).toStrictEqual(sql`true`) : ''; it.withCheck ? expect(it.withCheck).toStrictEqual(sql`true`) : ''; } } diff --git a/integration-tests/tests/pg/rls/rls.definition.test.ts b/integration-tests/tests/pg/rls/rls.definition.test.ts index 611116aae..81090633e 100644 --- a/integration-tests/tests/pg/rls/rls.definition.test.ts +++ b/integration-tests/tests/pg/rls/rls.definition.test.ts @@ -2,7 +2,7 @@ import { crudPolicy } from 'drizzle-orm/neon'; import { getTableConfig, integer, pgPolicy, pgRole, pgTable } from 'drizzle-orm/pg-core'; import { test } from 'vitest'; -test('getTableConfig: policies', async () => { +test.skip('getTableConfig: policies', async () => { const schema = pgTable('hhh', { id: integer(), }, () => [ From 6a162f5882ed690970d08d6de4e357e67ff44836 Mon Sep 17 00:00:00 2001 From: AndriiSherman Date: Thu, 24 Oct 2024 16:53:27 +0300 Subject: [PATCH 24/24] Add rls for supabase --- drizzle-kit/src/index.ts | 2 +- drizzle-orm/src/neon/rls.ts | 4 ++-- drizzle-orm/src/supabase/index.ts | 1 + drizzle-orm/src/supabase/rls.ts | 7 +++++++ 4 files changed, 11 insertions(+), 3 deletions(-) create mode 100644 drizzle-orm/src/supabase/index.ts create mode 100644 drizzle-orm/src/supabase/rls.ts diff --git a/drizzle-kit/src/index.ts b/drizzle-kit/src/index.ts index da8d1036b..4a57e59e3 100644 --- a/drizzle-kit/src/index.ts +++ b/drizzle-kit/src/index.ts @@ -127,7 +127,7 @@ export type Config = casing: 'camel' | 'preserve'; }; entities?: { - roles?: boolean | { provider?: string; exclude?: string[]; include?: string[] }[]; + roles?: boolean | { provider?: 'supabase' | 'neon' | string & {}; exclude?: string[]; include?: string[] }; }; } & ( diff --git a/drizzle-orm/src/neon/rls.ts b/drizzle-orm/src/neon/rls.ts index 0645b1d01..85f6ba220 100644 --- a/drizzle-orm/src/neon/rls.ts +++ b/drizzle-orm/src/neon/rls.ts @@ -1,6 +1,6 @@ import { is } from '~/entity.ts'; -import { pgPolicy, PgRole, pgRole } from '~/pg-core/index.ts'; -import type { AnyPgColumn, PgPolicyToOption } from '~/pg-core/index.ts'; +import { type AnyPgColumn, pgPolicy, type PgPolicyToOption } from '~/pg-core/index.ts'; +import { PgRole, pgRole } from '~/pg-core/roles.ts'; import { type SQL, sql } from '~/sql/sql.ts'; export const crudPolicy = (options: { diff --git a/drizzle-orm/src/supabase/index.ts b/drizzle-orm/src/supabase/index.ts new file mode 100644 index 000000000..ee201ff1c --- /dev/null +++ b/drizzle-orm/src/supabase/index.ts @@ -0,0 +1 @@ +export * from './rls.ts'; diff --git a/drizzle-orm/src/supabase/rls.ts b/drizzle-orm/src/supabase/rls.ts new file mode 100644 index 000000000..f6919a980 --- /dev/null +++ b/drizzle-orm/src/supabase/rls.ts @@ -0,0 +1,7 @@ +import { pgRole } from '~/pg-core/roles.ts'; + +// These are default roles that Supabase will set up. +export const anonRole = pgRole('anon').existing(); +export const authenticatedRole = pgRole('authenticated').existing(); +export const serviceRole = pgRole('service_role').existing(); +export const postgresRole = pgRole('postgres_role').existing();