Skip to content

Commit 007d1eb

Browse files
AleksandrShermanAndriiSherman
AleksandrSherman
and
AndriiSherman
authored
Drizzle-kit: fix recreate enums + altering data type to enums, from enums in pg (#4330)
Co-authored-by: AndriiSherman <andreysherman11@gmail.com>
1 parent f1c2dd6 commit 007d1eb

7 files changed

Lines changed: 2012 additions & 44 deletions

File tree

‎drizzle-kit/src/jsonStatements.ts‎

Lines changed: 55 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import {
77
Index,
88
MatViewWithOption,
99
PgSchema,
10+
PgSchemaSquashed,
1011
PgSquasher,
1112
Policy,
1213
Role,
@@ -168,10 +169,10 @@ export interface JsonAlterRoleStatement {
168169
export interface JsonDropValueFromEnumStatement {
169170
type: 'alter_type_drop_value';
170171
name: string;
171-
schema: string;
172+
enumSchema: string;
172173
deletedValues: string[];
173174
newValues: string[];
174-
columnsWithEnum: { schema: string; table: string; column: string }[];
175+
columnsWithEnum: { tableSchema: string; table: string; column: string; default?: string; columnType: string }[];
175176
}
176177

177178
export interface JsonCreateSequenceStatement {
@@ -472,6 +473,22 @@ export interface JsonAlterColumnTypeStatement {
472473
columnGenerated?: { as: string; type: 'stored' | 'virtual' };
473474
}
474475

476+
export interface JsonAlterColumnPgTypeStatement {
477+
type: 'pg_alter_table_alter_column_set_type';
478+
tableName: string;
479+
columnName: string;
480+
typeSchema: string | undefined;
481+
newDataType: { name: string; isEnum: boolean };
482+
oldDataType: { name: string; isEnum: boolean };
483+
schema: string;
484+
columnDefault: string;
485+
columnOnUpdate: boolean;
486+
columnNotNull: boolean;
487+
columnAutoIncrement: boolean;
488+
columnPk: boolean;
489+
columnGenerated?: { as: string; type: 'stored' | 'virtual' };
490+
}
491+
475492
export interface JsonAlterColumnSetPrimaryKeyStatement {
476493
type: 'alter_table_alter_column_set_pk';
477494
tableName: string;
@@ -784,6 +801,7 @@ export type JsonAlterViewStatement =
784801
export type JsonAlterColumnStatement =
785802
| JsonRenameColumnStatement
786803
| JsonAlterColumnTypeStatement
804+
| JsonAlterColumnPgTypeStatement
787805
| JsonAlterColumnSetDefaultStatement
788806
| JsonAlterColumnDropDefaultStatement
789807
| JsonAlterColumnSetNotNullStatement
@@ -1044,22 +1062,32 @@ export const prepareDropEnumValues = (
10441062
): JsonDropValueFromEnumStatement[] => {
10451063
if (!removedValues.length) return [];
10461064

1047-
const affectedColumns: { schema: string; table: string; column: string }[] = [];
1065+
const affectedColumns: JsonDropValueFromEnumStatement['columnsWithEnum'] = [];
10481066

10491067
for (const tableKey in json2.tables) {
10501068
const table = json2.tables[tableKey];
10511069
for (const columnKey in table.columns) {
10521070
const column = table.columns[columnKey];
1053-
if (column.type === name && column.typeSchema === schema) {
1054-
affectedColumns.push({ schema: table.schema || 'public', table: table.name, column: column.name });
1071+
1072+
const arrayDefinitionRegex = /\[\d*(?:\[\d*\])*\]/g;
1073+
const parsedColumnType = column.type.replace(arrayDefinitionRegex, '');
1074+
1075+
if (parsedColumnType === name && column.typeSchema === schema) {
1076+
affectedColumns.push({
1077+
tableSchema: table.schema,
1078+
table: table.name,
1079+
column: column.name,
1080+
columnType: column.type,
1081+
default: column.default,
1082+
});
10551083
}
10561084
}
10571085
}
10581086

10591087
return [{
10601088
type: 'alter_type_drop_value',
10611089
name: name,
1062-
schema: schema,
1090+
enumSchema: schema,
10631091
deletedValues: removedValues,
10641092
newValues: json2.enums[`${schema}.${name}`].values,
10651093
columnsWithEnum: affectedColumns,
@@ -2048,7 +2076,8 @@ export const preparePgAlterColumns = (
20482076
schema: string,
20492077
columns: AlteredColumn[],
20502078
// TODO: remove?
2051-
json2: CommonSquashedSchema,
2079+
json2: PgSchemaSquashed,
2080+
json1: PgSchemaSquashed,
20522081
action?: 'push' | undefined,
20532082
): JsonAlterColumnStatement[] => {
20542083
const tableKey = `${schema || 'public'}.${_tableName}`;
@@ -2074,6 +2103,8 @@ export const preparePgAlterColumns = (
20742103
).autoincrement;
20752104
const columnPk = (json2.tables[tableKey].columns[columnName] as any)
20762105
.primaryKey;
2106+
const typeSchema = json2.tables[tableKey].columns[columnName].typeSchema;
2107+
const json1ColumnTypeSchema = json1.tables[tableKey].columns[columnName].typeSchema;
20772108

20782109
const compositePk = json2.tables[tableKey].compositePrimaryKeys[`${tableName}_${columnName}`];
20792110

@@ -2088,12 +2119,26 @@ export const preparePgAlterColumns = (
20882119
}
20892120

20902121
if (column.type?.type === 'changed') {
2122+
const arrayDefinitionRegex = /\[\d*(?:\[\d*\])*\]/g;
2123+
const parsedNewColumnType = column.type.new.replace(arrayDefinitionRegex, '');
2124+
const parsedOldColumnType = column.type.old.replace(arrayDefinitionRegex, '');
2125+
2126+
const isNewTypeIsEnum = json2.enums[`${typeSchema}.${parsedNewColumnType}`];
2127+
const isOldTypeIsEnum = json1.enums[`${json1ColumnTypeSchema}.${parsedOldColumnType}`];
2128+
20912129
statements.push({
2092-
type: 'alter_table_alter_column_set_type',
2130+
type: 'pg_alter_table_alter_column_set_type',
20932131
tableName,
20942132
columnName,
2095-
newDataType: column.type.new,
2096-
oldDataType: column.type.old,
2133+
typeSchema: typeSchema,
2134+
newDataType: {
2135+
name: column.type.new,
2136+
isEnum: isNewTypeIsEnum ? true : false,
2137+
},
2138+
oldDataType: {
2139+
name: column.type.old,
2140+
isEnum: isOldTypeIsEnum ? true : false,
2141+
},
20972142
schema,
20982143
columnDefault,
20992144
columnOnUpdate,

‎drizzle-kit/src/serializer/pgSerializer.ts‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import {
66
getTableConfig,
77
getViewConfig,
88
IndexedColumn,
9+
PgArray,
910
PgColumn,
1011
PgDialect,
1112
PgEnum,
@@ -158,7 +159,14 @@ export const generatePgSnapshot = (
158159
const primaryKey: boolean = column.primary;
159160
const sqlTypeLowered = column.getSQLType().toLowerCase();
160161

161-
const typeSchema = is(column, PgEnumColumn) ? column.enum.schema || 'public' : undefined;
162+
const getEnumSchema = (column: PgColumn) => {
163+
while (is(column, PgArray)) {
164+
column = column.baseColumn;
165+
}
166+
return is(column, PgEnumColumn) ? column.enum.schema || 'public' : undefined;
167+
};
168+
const typeSchema: string | undefined = getEnumSchema(column);
169+
162170
const generated = column.generated;
163171
const identity = column.generatedIdentity;
164172

‎drizzle-kit/src/snapshotsDiffer.ts‎

Lines changed: 41 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1372,6 +1372,7 @@ export const applyPgSnapshotsDiff = async (
13721372
it.schema,
13731373
it.altered,
13741374
json2,
1375+
json1,
13751376
action,
13761377
);
13771378
})
@@ -1979,6 +1980,8 @@ export const applyPgSnapshotsDiff = async (
19791980

19801981
jsonStatements.push(...jsonDroppedReferencesForAlteredTables);
19811982

1983+
jsonStatements.push(...jsonAlterEnumsWithDroppedValues);
1984+
19821985
// Will need to drop indexes before changing any columns in table
19831986
// Then should go column alternations and then index creation
19841987
jsonStatements.push(...jsonDropIndexesForAllAlteredTables);
@@ -2001,7 +2004,6 @@ export const applyPgSnapshotsDiff = async (
20012004
jsonStatements.push(...jsonCreatedCheckConstraints);
20022005

20032006
jsonStatements.push(...jsonAlteredUniqueConstraints);
2004-
jsonStatements.push(...jsonAlterEnumsWithDroppedValues);
20052007

20062008
jsonStatements.push(...createViews);
20072009

@@ -2053,20 +2055,55 @@ export const applyPgSnapshotsDiff = async (
20532055
const filteredEnumsJsonStatements = filteredJsonStatements.filter((st) => {
20542056
if (st.type === 'alter_type_add_value') {
20552057
if (
2056-
jsonStatements.find(
2058+
filteredJsonStatements.find(
20572059
(it) =>
20582060
it.type === 'alter_type_drop_value'
20592061
&& it.name === st.name
2062+
&& it.enumSchema === st.schema,
2063+
)
2064+
) {
2065+
return false;
2066+
}
2067+
}
2068+
return true;
2069+
});
2070+
2071+
// This is needed because in sql generator on type pg_alter_table_alter_column_set_type and alter_type_drop_value
2072+
// drizzle kit checks whether column has defaults to cast them to new types properly
2073+
const filteredEnums2JsonStatements = filteredEnumsJsonStatements.filter((st) => {
2074+
if (st.type === 'alter_table_alter_column_set_default') {
2075+
if (
2076+
filteredEnumsJsonStatements.find(
2077+
(it) =>
2078+
it.type === 'pg_alter_table_alter_column_set_type'
2079+
&& it.columnDefault === st.newDefaultValue
2080+
&& it.columnName === st.columnName
2081+
&& it.tableName === st.tableName
20602082
&& it.schema === st.schema,
20612083
)
20622084
) {
20632085
return false;
20642086
}
2087+
2088+
if (
2089+
filteredEnumsJsonStatements.find(
2090+
(it) =>
2091+
it.type === 'alter_type_drop_value'
2092+
&& it.columnsWithEnum.find((column) =>
2093+
column.default === st.newDefaultValue
2094+
&& column.column === st.columnName
2095+
&& column.table === st.tableName
2096+
&& column.tableSchema === st.schema
2097+
),
2098+
)
2099+
) {
2100+
return false;
2101+
}
20652102
}
20662103
return true;
20672104
});
20682105

2069-
const sqlStatements = fromJson(filteredEnumsJsonStatements, 'postgresql', action);
2106+
const sqlStatements = fromJson(filteredEnums2JsonStatements, 'postgresql', action);
20702107

20712108
const uniqueSqlStatements: string[] = [];
20722109
sqlStatements.forEach((ss) => {
@@ -2087,7 +2124,7 @@ export const applyPgSnapshotsDiff = async (
20872124
const _meta = prepareMigrationMeta(rSchemas, rTables, rColumns);
20882125

20892126
return {
2090-
statements: filteredEnumsJsonStatements,
2127+
statements: filteredEnums2JsonStatements,
20912128
sqlStatements: uniqueSqlStatements,
20922129
_meta,
20932130
};

‎drizzle-kit/src/sqlgenerator.ts‎

Lines changed: 83 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
JsonAlterColumnDropNotNullStatement,
1212
JsonAlterColumnDropOnUpdateStatement,
1313
JsonAlterColumnDropPrimaryKeyStatement,
14+
JsonAlterColumnPgTypeStatement,
1415
JsonAlterColumnSetAutoincrementStatement,
1516
JsonAlterColumnSetDefaultStatement,
1617
JsonAlterColumnSetGeneratedStatement,
@@ -1456,33 +1457,53 @@ class AlterRenameTypeConvertor extends Convertor {
14561457
}
14571458

14581459
class AlterTypeDropValueConvertor extends Convertor {
1459-
can(statement: JsonStatement): boolean {
1460+
can(statement: JsonDropValueFromEnumStatement): boolean {
14601461
return statement.type === 'alter_type_drop_value';
14611462
}
14621463

14631464
convert(st: JsonDropValueFromEnumStatement) {
1464-
const { columnsWithEnum, name, newValues, schema } = st;
1465+
const { columnsWithEnum, name, newValues, enumSchema } = st;
14651466

14661467
const statements: string[] = [];
14671468

14681469
for (const withEnum of columnsWithEnum) {
1470+
const tableNameWithSchema = withEnum.tableSchema
1471+
? `"${withEnum.tableSchema}"."${withEnum.table}"`
1472+
: `"${withEnum.table}"`;
1473+
14691474
statements.push(
1470-
`ALTER TABLE "${withEnum.schema}"."${withEnum.table}" ALTER COLUMN "${withEnum.column}" SET DATA TYPE text;`,
1475+
`ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${withEnum.column}" SET DATA TYPE text;`,
14711476
);
1477+
if (withEnum.default) {
1478+
statements.push(
1479+
`ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${withEnum.column}" SET DEFAULT ${withEnum.default}::text;`,
1480+
);
1481+
}
14721482
}
14731483

1474-
statements.push(new DropTypeEnumConvertor().convert({ name: name, schema, type: 'drop_type_enum' }));
1484+
statements.push(new DropTypeEnumConvertor().convert({ name: name, schema: enumSchema, type: 'drop_type_enum' }));
14751485

14761486
statements.push(new CreateTypeEnumConvertor().convert({
14771487
name: name,
1478-
schema: schema,
1488+
schema: enumSchema,
14791489
values: newValues,
14801490
type: 'create_type_enum',
14811491
}));
14821492

14831493
for (const withEnum of columnsWithEnum) {
1494+
const tableNameWithSchema = withEnum.tableSchema
1495+
? `"${withEnum.tableSchema}"."${withEnum.table}"`
1496+
: `"${withEnum.table}"`;
1497+
1498+
const parsedType = parseType(`"${enumSchema}".`, withEnum.columnType);
1499+
if (withEnum.default) {
1500+
statements.push(
1501+
`ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${withEnum.column}" SET DEFAULT ${withEnum.default}::${parsedType};`,
1502+
);
1503+
}
1504+
14841505
statements.push(
1485-
`ALTER TABLE "${withEnum.schema}"."${withEnum.table}" ALTER COLUMN "${withEnum.column}" SET DATA TYPE "${schema}"."${name}" USING "${withEnum.column}"::"${schema}"."${name}";`,
1506+
`ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${withEnum.column}" SET DATA TYPE ${parsedType} USING "${withEnum.column}"::${parsedType};`,
14861507
);
14871508
}
14881509

@@ -1871,19 +1892,71 @@ export class SQLiteAlterTableAddColumnConvertor extends Convertor {
18711892
class PgAlterTableAlterColumnSetTypeConvertor extends Convertor {
18721893
can(statement: JsonStatement, dialect: Dialect): boolean {
18731894
return (
1874-
statement.type === 'alter_table_alter_column_set_type'
1895+
statement.type === 'pg_alter_table_alter_column_set_type'
18751896
&& dialect === 'postgresql'
18761897
);
18771898
}
18781899

1879-
convert(statement: JsonAlterColumnTypeStatement) {
1880-
const { tableName, columnName, newDataType, schema } = statement;
1900+
convert(statement: JsonAlterColumnPgTypeStatement) {
1901+
const { tableName, columnName, newDataType, schema, oldDataType, columnDefault, typeSchema } = statement;
18811902

18821903
const tableNameWithSchema = schema
18831904
? `"${schema}"."${tableName}"`
18841905
: `"${tableName}"`;
18851906

1886-
return `ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${columnName}" SET DATA TYPE ${newDataType};`;
1907+
const statements: string[] = [];
1908+
1909+
const type = parseType(`"${typeSchema}".`, newDataType.name);
1910+
1911+
if (!oldDataType.isEnum && !newDataType.isEnum) {
1912+
statements.push(
1913+
`ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${columnName}" SET DATA TYPE ${type};`,
1914+
);
1915+
if (columnDefault) {
1916+
statements.push(
1917+
`ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${columnName}" SET DEFAULT ${columnDefault};`,
1918+
);
1919+
}
1920+
}
1921+
1922+
if (oldDataType.isEnum && !newDataType.isEnum) {
1923+
statements.push(
1924+
`ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${columnName}" SET DATA TYPE ${type};`,
1925+
);
1926+
if (columnDefault) {
1927+
statements.push(
1928+
`ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${columnName}" SET DEFAULT ${columnDefault};`,
1929+
);
1930+
}
1931+
}
1932+
1933+
if (!oldDataType.isEnum && newDataType.isEnum) {
1934+
if (columnDefault) {
1935+
statements.push(
1936+
`ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${columnName}" SET DEFAULT ${columnDefault}::${type};`,
1937+
);
1938+
}
1939+
statements.push(
1940+
`ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${columnName}" SET DATA TYPE ${type} USING "${columnName}"::${type};`,
1941+
);
1942+
}
1943+
1944+
if (oldDataType.isEnum && newDataType.isEnum) {
1945+
const alterType =
1946+
`ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${columnName}" SET DATA TYPE ${type} USING "${columnName}"::text::${type};`;
1947+
1948+
if (newDataType.name !== oldDataType.name && columnDefault) {
1949+
statements.push(
1950+
`ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${columnName}" DROP DEFAULT;`,
1951+
alterType,
1952+
`ALTER TABLE ${tableNameWithSchema} ALTER COLUMN "${columnName}" SET DEFAULT ${columnDefault};`,
1953+
);
1954+
} else {
1955+
statements.push(alterType);
1956+
}
1957+
}
1958+
1959+
return statements;
18871960
}
18881961
}
18891962

0 commit comments

Comments
 (0)