Skip to content

Bound PostgreSQL schema suggestion queries #32

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
joe-thom wants to merge 1 commit into
base: main
Choose a base branch
from
+85 −38
Draft

Bound PostgreSQL schema suggestion queries #32

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 56 additions & 38 deletions drizzle-kit/src/cli/commands/pgPushUtils.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,12 +73,29 @@ export type SelectResolverOutput = {
};
};

export type PgSuggestionsRowCheck = 'exact' | 'exists';

async function getRowCount(db: DB, tableName: string, rowCheck: PgSuggestionsRowCheck) {
if (rowCheck === 'exists') {
const rows = await db.query(`select 1 from ${tableName} limit 1`);
return rows.length;
}

const rows = await db.query<{ count: string | number }>(`select count(*) as count from ${tableName}`);
return Number(rows[0].count);
}

function describeRowCount(count: number, rowCheck: PgSuggestionsRowCheck) {
return rowCheck === 'exact' ? `${count} items` : 'existing items';
}

export const pgSuggestions = async (
db: DB,
statements: JsonStatement[],
selectResolver?: (
input: SelectResolverInput,
) => Promise<SelectResolverOutput>,
rowCheck: PgSuggestionsRowCheck = 'exact',
) => {
let shouldAskForApprove = false;
const statementsToExecute: string[] = [];
Expand All @@ -99,43 +116,46 @@ export const pgSuggestions = async (
} else if (statement.type === 'rename_table') {
renamedTables[concatSchemaAndTableName(statement.toSchema, statement.tableNameTo)] = statement.tableNameFrom;
} else if (statement.type === 'drop_table') {
const res = await db.query(
`select count(*) as count from ${
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables)
}`,
const count = await getRowCount(
db,
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables),
rowCheck,
);
const count = Number(res[0].count);
if (count > 0) {
infoToPrint.push(`· You're about to delete ${chalk.underline(statement.tableName)} table with ${count} items`);
infoToPrint.push(
`· You're about to delete ${chalk.underline(statement.tableName)} table with ${
describeRowCount(count, rowCheck)
}`,
);
// statementsToExecute.push(
// `truncate table ${tableNameWithSchemaFrom(statement)} cascade;`
// );
tablesToRemove.push(statement.tableName);
shouldAskForApprove = true;
}
} else if (statement.type === 'drop_view' && statement.materialized) {
const res = await db.query(`select count(*) as count from "${statement.schema ?? 'public'}"."${statement.name}"`);
const count = Number(res[0].count);
const count = await getRowCount(db, `"${statement.schema ?? 'public'}"."${statement.name}"`, rowCheck);
if (count > 0) {
infoToPrint.push(
`· You're about to delete "${chalk.underline(statement.name)}" materialized view with ${count} items`,
`· You're about to delete "${chalk.underline(statement.name)}" materialized view with ${
describeRowCount(count, rowCheck)
}`,
);

matViewsToRemove.push(statement.name);
shouldAskForApprove = true;
}
} else if (statement.type === 'alter_table_drop_column') {
const res = await db.query(
`select count(*) as count from ${
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables)
}`,
const count = await getRowCount(
db,
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables),
rowCheck,
);
const count = Number(res[0].count);
if (count > 0) {
infoToPrint.push(
`· You're about to delete ${
chalk.underline(statement.columnName)
} column in ${statement.tableName} table with ${count} items`,
} column in ${statement.tableName} table with ${describeRowCount(count, rowCheck)}`,
);
columnsToRemove.push(`${statement.tableName}_${statement.columnName}`);
shouldAskForApprove = true;
Expand All @@ -151,12 +171,11 @@ export const pgSuggestions = async (
shouldAskForApprove = true;
}
} else if (statement.type === 'alter_table_alter_column_set_type') {
const res = await db.query(
`select count(*) as count from ${
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables)
}`,
const count = await getRowCount(
db,
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables),
rowCheck,
);
const count = Number(res[0].count);
if (count > 0) {
infoToPrint.push(
`· You're about to change ${chalk.underline(statement.columnName)} column type from ${
Expand All @@ -165,7 +184,7 @@ export const pgSuggestions = async (
chalk.underline(
statement.newDataType,
)
} with ${count} items`,
} with ${describeRowCount(count, rowCheck)}`,
);
statementsToExecute.push(
`truncate table ${
Expand All @@ -176,12 +195,11 @@ export const pgSuggestions = async (
shouldAskForApprove = true;
}
} else if (statement.type === 'alter_table_alter_column_drop_pk') {
const res = await db.query(
`select count(*) as count from ${
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables)
}`,
const count = await getRowCount(
db,
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables),
rowCheck,
);
const count = Number(res[0].count);
if (count > 0) {
infoToPrint.push(
`· You're about to change ${
Expand Down Expand Up @@ -216,17 +234,16 @@ export const pgSuggestions = async (
continue;
} else if (statement.type === 'alter_table_add_column') {
if (statement.column.notNull && typeof statement.column.default === 'undefined') {
const res = await db.query(
`select count(*) as count from ${
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables)
}`,
const count = await getRowCount(
db,
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables),
rowCheck,
);
const count = Number(res[0].count);
if (count > 0) {
infoToPrint.push(
`· You're about to add not-null ${
chalk.underline(statement.column.name)
} column without default value, which contains ${count} items`,
} column without default value, which contains ${describeRowCount(count, rowCheck)}`,
);

tablesToTruncate.push(statement.tableName);
Expand All @@ -240,20 +257,21 @@ export const pgSuggestions = async (
}
}
} else if (statement.type === 'create_unique_constraint') {
const res = await db.query(
`select count(*) as count from ${
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables)
}`,
const count = await getRowCount(
db,
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables),
rowCheck,
);
const count = Number(res[0].count);
if (count > 0) {
const unsquashedUnique = PgSquasher.unsquashUnique(statement.data);
console.log(
`· You're about to add ${
chalk.underline(
unsquashedUnique.name,
)
} unique constraint to the table, which contains ${count} items. If this statement fails, you will receive an error from the database. Do you want to truncate ${
} unique constraint to the table, which contains ${
describeRowCount(count, rowCheck)
}. If this statement fails, you will receive an error from the database. Do you want to truncate ${
chalk.underline(
statement.tableName,
)
Expand Down
29 changes: 29 additions & 0 deletions drizzle-kit/tests/pgSuggestions.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
import { describe, expect, test, vi } from 'vitest';
import { pgSuggestions } from '../src/cli/commands/pgPushUtils';
import type { JsonStatement } from '../src/jsonStatements';

const dropTable: JsonStatement = {
type: 'drop_table',
tableName: 'users',
schema: 'public',
};

describe('pgSuggestions row checks', () => {
test('counts rows by default', async () => {
const query = vi.fn().mockResolvedValue([{ count: '42' }]);

const result = await pgSuggestions({ query }, [dropTable]);

expect(query).toHaveBeenCalledWith('select count(*) as count from "public"."users"');
expect(result.infoToPrint).toEqual(["· You're about to delete users table with 42 items"]);
});

test('can stop after the first row', async () => {
const query = vi.fn().mockResolvedValue([{}]);

const result = await pgSuggestions({ query }, [dropTable], undefined, 'exists');

expect(query).toHaveBeenCalledWith('select 1 from "public"."users" limit 1');
expect(result.infoToPrint).toEqual(["· You're about to delete users table with existing items"]);
});
});
Loading