Skip to content

Commit

Permalink
Merge pull request #3076 from drizzle-team/drizzle-kit/enums
Browse files Browse the repository at this point in the history 
[Drizzle-kit]: enums updates for pg
  • Loading branch information
@AndriiSherman
AndriiSherman authored Oct 10, 2024
2 parents a9aca5c + 98c1480 commit 9937f59
Show file tree
Hide file tree
Showing 6 changed files with 606 additions and 80 deletions.
42 changes: 41 additions & 1 deletion drizzle-kit/src/jsonStatements.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,15 @@ export interface JsonAddValueToEnumStatement {
before: string;
}

export interface JsonDropValueFromEnumStatement {
type: 'alter_type_drop_value';
name: string;
schema: string;
deletedValues: string[];
newValues: string[];
columnsWithEnum: { schema: string; table: string; column: string }[];
}

export interface JsonCreateSequenceStatement {
type: 'create_sequence';
name: string;
Expand Down Expand Up @@ -582,7 +591,8 @@ export type JsonStatement =
| JsonDropSequenceStatement
| JsonCreateSequenceStatement
| JsonMoveSequenceStatement
| JsonRenameSequenceStatement;
| JsonRenameSequenceStatement
| JsonDropValueFromEnumStatement;

export const preparePgCreateTableJson = (
table: Table,
Expand Down Expand Up @@ -717,6 +727,36 @@ export const prepareAddValuesToEnumJson = (
});
};

export const prepareDropEnumValues = (
name: string,
schema: string,
removedValues: string[],
json2: PgSchema,
): JsonDropValueFromEnumStatement[] => {
if (!removedValues.length) return [];

const affectedColumns: { schema: string; table: string; column: string }[] = [];

for (const tableKey in json2.tables) {
const table = json2.tables[tableKey];
for (const columnKey in table.columns) {
const column = table.columns[columnKey];
if (column.type === name && column.typeSchema === schema) {
affectedColumns.push({ schema: table.schema || 'public', table: table.name, column: column.name });
}
}
}

return [{
type: 'alter_type_drop_value',
name: name,
schema: schema,
deletedValues: removedValues,
newValues: json2.enums[`${schema}.${name}`].values,
columnsWithEnum: affectedColumns,
}];
};

export const prepareDropEnumJson = (
name: string,
schema: string,
Expand Down
55 changes: 27 additions & 28 deletions drizzle-kit/src/snapshotsDiffer.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ import {
prepareDeleteSchemasJson as prepareDropSchemasJson,
prepareDeleteUniqueConstraintPg as prepareDeleteUniqueConstraint,
prepareDropEnumJson,
prepareDropEnumValues,
prepareDropIndexesJson,
prepareDropReferencesJson,
prepareDropSequenceJson,
Expand Down Expand Up @@ -1002,30 +1003,6 @@ export const applyPgSnapshotsDiff = async (
// - create table with generated
// - alter - should be not triggered, but should get warning

// TODO:
// let hasEnumValuesDeletions = false;
// let enumValuesDeletions: { name: string; schema: string; values: string[] }[] =
// [];
// for (let alteredEnum of typedResult.alteredEnums) {
// if (alteredEnum.deletedValues.length > 0) {
// hasEnumValuesDeletions = true;
// enumValuesDeletions.push({
// name: alteredEnum.name,
// schema: alteredEnum.schema,
// values: alteredEnum.deletedValues,
// });
// }
// }
// if (hasEnumValuesDeletions) {
// console.log(error("Deletion of enum values is prohibited in Postgres - see here"));
// for(let entry of enumValuesDeletions){
// console.log(error(`You're trying to delete ${chalk.blue(`[${entry.values.join(", ")}]`)} values from ${chalk.blue(`${entry.schema}.${entry.name}`)}`))
// }
// }
// if (hasEnumValuesDeletions && action === "push") {
// process.exit(1);
// }

const createEnums = createdEnums.map((it) => {
return prepareCreateEnumJson(it.name, it.schema, it.values);
}) ?? [];
Expand All @@ -1042,14 +1019,17 @@ export const applyPgSnapshotsDiff = async (
return prepareRenameEnumJson(it.from.name, it.to.name, it.to.schema);
});

// todo: block enum rename, enum value rename and enun deletion for now
const jsonAlterEnumsWithAddedValues = typedResult.alteredEnums
.map((it) => {
return prepareAddValuesToEnumJson(it.name, it.schema, it.addedValues);
})
.flat() ?? [];

///////////
const jsonAlterEnumsWithDroppedValues = typedResult.alteredEnums
.map((it) => {
return prepareDropEnumValues(it.name, it.schema, it.deletedValues, curFull);
})
.flat() ?? [];

const createSequences = createdSequences.map((it) => {
return prepareCreateSequenceJson(it);
Expand Down Expand Up @@ -1135,6 +1115,7 @@ export const applyPgSnapshotsDiff = async (
jsonStatements.push(...jsonAddedUniqueConstraints);

jsonStatements.push(...jsonAlteredUniqueConstraints);
jsonStatements.push(...jsonAlterEnumsWithDroppedValues);

jsonStatements.push(...dropEnums);
jsonStatements.push(...dropSequences);
Expand Down Expand Up @@ -1169,7 +1150,25 @@ export const applyPgSnapshotsDiff = async (
return true;
});

const sqlStatements = fromJson(filteredJsonStatements, 'postgresql');
// enum filters
// Need to find add and drop enum values in same enum and remove add values
const filteredEnumsJsonStatements = filteredJsonStatements.filter((st) => {
if (st.type === 'alter_type_add_value') {
if (
jsonStatements.find(
(it) =>
it.type === 'alter_type_drop_value'
&& it.name === st.name
&& it.schema === st.schema,
)
) {
return false;
}
}
return true;
});

const sqlStatements = fromJson(filteredEnumsJsonStatements, 'postgresql');

const uniqueSqlStatements: string[] = [];
sqlStatements.forEach((ss) => {
Expand All @@ -1190,7 +1189,7 @@ export const applyPgSnapshotsDiff = async (
const _meta = prepareMigrationMeta(rSchemas, rTables, rColumns);

return {
statements: filteredJsonStatements,
statements: filteredEnumsJsonStatements,
sqlStatements: uniqueSqlStatements,
_meta,
};
Expand Down
121 changes: 105 additions & 16 deletions drizzle-kit/src/sqlgenerator.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,13 +38,17 @@ import {
JsonDeleteReferenceStatement,
JsonDeleteUniqueConstraint,
JsonDropColumnStatement,
JsonDropEnumStatement,
JsonDropIndexStatement,
JsonDropSequenceStatement,
JsonDropTableStatement,
JsonDropValueFromEnumStatement,
JsonMoveEnumStatement,
JsonMoveSequenceStatement,
JsonPgCreateIndexStatement,
JsonRecreateTableStatement,
JsonRenameColumnStatement,
JsonRenameEnumStatement,
JsonRenameSchema,
JsonRenameSequenceStatement,
JsonRenameTableStatement,
Expand Down Expand Up @@ -694,22 +698,39 @@ class CreateTypeEnumConvertor extends Convertor {
convert(st: JsonCreateEnumStatement) {
const { name, values, schema } = st;

const tableNameWithSchema = schema ? `"${schema}"."${name}"` : `"${name}"`;
const enumNameWithSchema = schema ? `"${schema}"."${name}"` : `"${name}"`;

let valuesStatement = '(';
valuesStatement += values.map((it) => `'${it}'`).join(', ');
valuesStatement += ')';

let statement = 'DO $$ BEGIN';
statement += '\n';
statement += ` CREATE TYPE ${tableNameWithSchema} AS ENUM${valuesStatement};`;
statement += '\n';
statement += 'EXCEPTION';
statement += '\n';
statement += ' WHEN duplicate_object THEN null;';
statement += '\n';
statement += 'END $$;';
statement += '\n';
// TODO do we need this?
// let statement = 'DO $$ BEGIN';
// statement += '\n';
let statement = `CREATE TYPE ${enumNameWithSchema} AS ENUM${valuesStatement};`;
// statement += '\n';
// statement += 'EXCEPTION';
// statement += '\n';
// statement += ' WHEN duplicate_object THEN null;';
// statement += '\n';
// statement += 'END $$;';
// statement += '\n';
return statement;
}
}

class DropTypeEnumConvertor extends Convertor {
can(statement: JsonStatement): boolean {
return statement.type === 'drop_type_enum';
}

convert(st: JsonDropEnumStatement) {
const { name, schema } = st;

const enumNameWithSchema = schema ? `"${schema}"."${name}"` : `"${name}"`;

let statement = `DROP TYPE ${enumNameWithSchema};`;

return statement;
}
}
Expand All @@ -720,9 +741,74 @@ class AlterTypeAddValueConvertor extends Convertor {
}

convert(st: JsonAddValueToEnumStatement) {
const { name, schema, value } = st;
const schemaPrefix = schema && schema !== 'public' ? `"${schema}".` : '';
return `ALTER TYPE ${schemaPrefix}"${name}" ADD VALUE '${value}';`;
const { name, schema, value, before } = st;

const enumNameWithSchema = schema ? `"${schema}"."${name}"` : `"${name}"`;

return `ALTER TYPE ${enumNameWithSchema} ADD VALUE '${value}'${before.length ? ` BEFORE '${before}'` : ''};`;
}
}

class AlterTypeSetSchemaConvertor extends Convertor {
can(statement: JsonStatement): boolean {
return statement.type === 'move_type_enum';
}

convert(st: JsonMoveEnumStatement) {
const { name, schemaFrom, schemaTo } = st;

const enumNameWithSchema = schemaFrom ? `"${schemaFrom}"."${name}"` : `"${name}"`;

return `ALTER TYPE ${enumNameWithSchema} SET SCHEMA "${schemaTo}";`;
}
}

class AlterRenameTypeConvertor extends Convertor {
can(statement: JsonStatement): boolean {
return statement.type === 'rename_type_enum';
}

convert(st: JsonRenameEnumStatement) {
const { nameTo, nameFrom, schema } = st;

const enumNameWithSchema = schema ? `"${schema}"."${nameFrom}"` : `"${nameFrom}"`;

return `ALTER TYPE ${enumNameWithSchema} RENAME TO "${nameTo}";`;
}
}

class AlterTypeDropValueConvertor extends Convertor {
can(statement: JsonStatement): boolean {
return statement.type === 'alter_type_drop_value';
}

convert(st: JsonDropValueFromEnumStatement) {
const { columnsWithEnum, name, newValues, schema } = st;

const statements: string[] = [];

for (const withEnum of columnsWithEnum) {
statements.push(
`ALTER TABLE "${withEnum.schema}"."${withEnum.table}" ALTER COLUMN "${withEnum.column}" SET DATA TYPE text;`,
);
}

statements.push(new DropTypeEnumConvertor().convert({ name: name, schema, type: 'drop_type_enum' }));

statements.push(new CreateTypeEnumConvertor().convert({
name: name,
schema: schema,
values: newValues,
type: 'create_type_enum',
}));

for (const withEnum of columnsWithEnum) {
statements.push(
`ALTER TABLE "${withEnum.schema}"."${withEnum.table}" ALTER COLUMN "${withEnum.column}" SET DATA TYPE "${schema}"."${name}" USING "${withEnum.column}"::"${schema}"."${name}";`,
);
}

return statements;
}
}

Expand Down Expand Up @@ -2487,6 +2573,11 @@ convertors.push(new SQLiteRecreateTableConvertor());
convertors.push(new LibSQLRecreateTableConvertor());

convertors.push(new CreateTypeEnumConvertor());
convertors.push(new DropTypeEnumConvertor());
convertors.push(new AlterTypeAddValueConvertor());
convertors.push(new AlterTypeSetSchemaConvertor());
convertors.push(new AlterRenameTypeConvertor());
convertors.push(new AlterTypeDropValueConvertor());

convertors.push(new CreatePgSequenceConvertor());
convertors.push(new DropPgSequenceConvertor());
Expand Down Expand Up @@ -2530,8 +2621,6 @@ convertors.push(new PgDropIndexConvertor());
convertors.push(new SqliteDropIndexConvertor());
convertors.push(new MySqlDropIndexConvertor());

convertors.push(new AlterTypeAddValueConvertor());

convertors.push(new PgAlterTableAlterColumnSetPrimaryKeyConvertor());
convertors.push(new PgAlterTableAlterColumnDropPrimaryKeyConvertor());
convertors.push(new PgAlterTableAlterColumnSetNotNullConvertor());
Expand Down
Loading

0 comments on commit 9937f59

Please sign in to comment.