[psql]: kit pk tests

Author: AleksandrSherman

drizzle-kit/src/dialects/postgres/convertor.ts

@@ -544,9 +544,10 @@ const dropPrimaryKeyConvertor = convertor('drop_pk', (st) => {
 });
 
 const recreatePrimaryKeyConvertor = convertor('alter_pk', (it) => {
-	const drop = dropPrimaryKeyConvertor.convert({ pk: it.pk }) as string;
-	const create = addPrimaryKeyConvertor.convert({ pk: it.pk }) as string;
-	return [drop, create];
+	const st: string[] = [];
+	if (!it.deleted) st.push(dropPrimaryKeyConvertor.convert({ pk: it.pk }) as string);
+	st.push(addPrimaryKeyConvertor.convert({ pk: it.pk }) as string);
+	return st;
 });
 
 const renameConstraintConvertor = convertor('rename_constraint', (st) => {

drizzle-kit/src/dialects/postgres/diff.ts

@@ -27,7 +27,7 @@ import type {
 } from './ddl';
 import { createDDL, tableFromDDL } from './ddl';
 import { defaults, defaultsCommutative } from './grammar';
-import type { JsonRecreateIndex, JsonStatement } from './statements';
+import type { JsonAlterPrimaryKey, JsonRecreateIndex, JsonStatement } from './statements';
 import { prepareStatement } from './statements';
 
 export const ddlDiffDry = async (ddlFrom: PostgresDDL, ddlTo: PostgresDDL, mode: 'default' | 'push') => {
@@ -728,7 +728,10 @@ export const ddlDiff = async (
 	const jsonAddColumnsStatemets = columnsToCreate.filter(tablesFilter('created')).map((it) =>
 		prepareStatement('add_column', {
 			column: it,
-			isPK: ddl2.pks.one({ schema: it.schema, table: it.table, columns: [it.name] }) !== null,
+			// if pk existed before and new column now has pk, this will trigger alter_pk, that will automatically add pk
+			// this flag is needed for column recreation (generated)
+			// see tests: "drizzle-kit/tests/postgres/pg-constraints.test.ts" => "remove/add pk" and below
+			isPK: false, // ddl2.pks.one({ schema: it.schema, table: it.table, columns: [it.name] }) !== null,
 			isCompositePK: ddl2.pks.one({ schema: it.schema, table: it.table, columns: { CONTAINS: it.name } }) !== null,
 		})
 	);
@@ -845,7 +848,11 @@ export const ddlDiff = async (
 	});
 
 	const alteredChecks = alters.filter((it) => it.entityType === 'checks');
-	const jsonAlteredPKs = alteredPKs.map((it) => prepareStatement('alter_pk', { diff: it, pk: it.$right }));
+	const jsonAlteredPKs: JsonAlterPrimaryKey[] = alteredPKs.map((it) => {
+		const deleted = columnsToDelete.some((x) => it.columns?.from.includes(x.name));
+
+		return prepareStatement('alter_pk', { diff: it, pk: it.$right, deleted });
+	});
 
 	const jsonRecreateFKs = alters.filter((it) => it.entityType === 'fks').filter((x) => {
 		if (
@@ -1220,13 +1227,17 @@ export const ddlDiff = async (
 	jsonStatements.push(...jsonDropIndexes);
 	jsonStatements.push(...jsonDropPrimaryKeys);
 
-	jsonStatements.push(...jsonAddPrimaryKeys);
-	jsonStatements.push(...jsonRenamePrimaryKey);
 	jsonStatements.push(...jsonRenameReferences);
 	jsonStatements.push(...jsonAddColumnsStatemets);
+	jsonStatements.push(...jsonAddPrimaryKeys);
+	jsonStatements.push(...jsonRenamePrimaryKey);
 	jsonStatements.push(...recreateEnums);
 	jsonStatements.push(...jsonRecreateColumns);
+
+	jsonStatements.push(...jsonDropColumnsStatemets);
+	jsonStatements.push(...jsonAlteredPKs);
 	jsonStatements.push(...jsonAlterColumns);
+
 	jsonStatements.push(...jsonRecreateIndex);
 
 	jsonStatements.push(...jsonRenamedUniqueConstraints);
@@ -1237,9 +1248,6 @@ export const ddlDiff = async (
 	jsonStatements.push(...jsonCreateFKs);
 	jsonStatements.push(...jsonRecreateFKs);
 
-	jsonStatements.push(...jsonDropColumnsStatemets);
-	jsonStatements.push(...jsonAlteredPKs);
-
 	jsonStatements.push(...jsonCreatedCheckConstraints);
 
 	jsonStatements.push(...jsonAlterCheckConstraints);

drizzle-kit/src/dialects/postgres/statements.ts

@@ -278,6 +278,7 @@ export interface JsonAlterPrimaryKey {
 	type: 'alter_pk';
 	pk: PrimaryKey;
 	diff: DiffEntities['pks'];
+	deleted?: boolean;
 }
 
 export interface JsonMoveTable {

drizzle-kit/tests/cockroach/identity.test.ts

@@ -228,7 +228,6 @@ test.concurrent('drop identity from a column - no params', async ({ db }) => {
 	const { sqlStatements: pst } = await push({
 		db,
 		to,
-		log: 'statements',
 	});
 
 	const st0 = [

drizzle-kit/tests/cockroach/views.test.ts

@@ -148,7 +148,6 @@ test.concurrent('create materialized view', async ({ db }) => {
 	const { sqlStatements: pst } = await push({
 		db,
 		to: schema2,
-		log: 'statements',
 	});
 
 	const st0: string[] = [

drizzle-kit/tests/postgres/pg-columns.test.ts

@@ -342,75 +342,6 @@ test('alter table add composite pk', async (t) => {
 	expect(pst).toStrictEqual(st0);
 });
 
-// https://github.com/drizzle-team/drizzle-orm/issues/3496
-test('remove/add pk', async (t) => {
-	const Step = pgTable('Step', {
-		id: bigint({ mode: 'number' }).primaryKey(),
-	});
-	const schema1 = {
-		Step1: Step,
-		Branch: pgTable('Branch', {
-			id: bigint({ mode: 'number' }).primaryKey(),
-			stepId: bigint({ mode: 'number' }).references(() => Step.id, { onDelete: 'cascade' }),
-		}),
-	};
-
-	const { next: n1 } = await diff({}, schema1, []);
-	await push({ db, to: schema1 });
-
-	const schema2 = {
-		Step,
-		Branch: pgTable('Branch', {
-			stepId: bigint({ mode: 'number' }).primaryKey().references(() => Step.id, { onDelete: 'cascade' }),
-		}),
-	};
-
-	const { sqlStatements: st2 } = await diff(n1, schema2, []);
-	const { sqlStatements: pst2 } = await push({ db, to: schema2, log: 'statements' });
-
-	const expectedSt2 = [
-		'ALTER TABLE "Branch" DROP COLUMN "id";',
-		'ALTER TABLE "Branch" ADD PRIMARY KEY ("stepId");',
-	];
-	expect(st2).toStrictEqual(expectedSt2);
-	expect(pst2).toStrictEqual(expectedSt2);
-});
-test('remove/add pk #2', async (t) => {
-	const Step = pgTable('Step', {
-		id: bigint({ mode: 'number' }).primaryKey(),
-	});
-	const schema1 = {
-		Step1: Step,
-		Branch: pgTable('Branch', {
-			id: bigint({ mode: 'number' }).primaryKey(),
-			stepId: bigint({ mode: 'number' }).references(() => Step.id, { onDelete: 'cascade' }),
-		}),
-	};
-
-	const { next: n1 } = await diff({}, schema1, []);
-	await push({ db, to: schema1 });
-
-	const schema2 = {
-		Step,
-		Branch: pgTable('Branch', {
-			stepId: bigint({ mode: 'number' }).references(() => Step.id, { onDelete: 'cascade' }),
-			stepId2: bigint({ mode: 'number' }).primaryKey(),
-		}),
-	};
-
-	const { sqlStatements: st2 } = await diff(n1, schema2, []);
-	const { sqlStatements: pst2 } = await push({ db, to: schema2, log: 'statements' });
-
-	const expectedSt2 = [
-		'ALTER TABLE "Branch" ADD COLUMN "stepId2" bigint PRIMARY KEY;',
-		'ALTER TABLE "Branch" DROP CONSTRAINT "Branch_pkey";',
-		'ALTER TABLE "Branch" ADD PRIMARY KEY ("stepId2");',
-		'ALTER TABLE "Branch" DROP COLUMN "id";',
-	];
-	expect(st2).toStrictEqual(expectedSt2);
-	expect(pst2).toStrictEqual(expectedSt2);
-});
-
 test('rename table rename column #1', async (t) => {
 	const schema1 = {
 		users: pgTable('users', {