Skip to content

Commit 09b16d8

Browse files
joe-thomjoe-thom
committed
Bound Postgres schema suggestion queries
1 parent 1c1372f commit 09b16d8

2 files changed

Lines changed: 85 additions & 38 deletions

File tree

‎drizzle-kit/src/cli/commands/pgPushUtils.ts‎

Lines changed: 56 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -73,12 +73,29 @@ export type SelectResolverOutput = {
7373
};
7474
};
7575

76+
export type PgSuggestionsRowCheck = 'exact' | 'exists';
77+
78+
async function getRowCount(db: DB, tableName: string, rowCheck: PgSuggestionsRowCheck) {
79+
if (rowCheck === 'exists') {
80+
const rows = await db.query(`select 1 from ${tableName} limit 1`);
81+
return rows.length;
82+
}
83+
84+
const rows = await db.query<{ count: string | number }>(`select count(*) as count from ${tableName}`);
85+
return Number(rows[0].count);
86+
}
87+
88+
function describeRowCount(count: number, rowCheck: PgSuggestionsRowCheck) {
89+
return rowCheck === 'exact' ? `${count} items` : 'existing items';
90+
}
91+
7692
export const pgSuggestions = async (
7793
db: DB,
7894
statements: JsonStatement[],
7995
selectResolver?: (
8096
input: SelectResolverInput,
8197
) => Promise<SelectResolverOutput>,
98+
rowCheck: PgSuggestionsRowCheck = 'exact',
8299
) => {
83100
let shouldAskForApprove = false;
84101
const statementsToExecute: string[] = [];
@@ -99,43 +116,46 @@ export const pgSuggestions = async (
99116
} else if (statement.type === 'rename_table') {
100117
renamedTables[concatSchemaAndTableName(statement.toSchema, statement.tableNameTo)] = statement.tableNameFrom;
101118
} else if (statement.type === 'drop_table') {
102-
const res = await db.query(
103-
`select count(*) as count from ${
104-
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables)
105-
}`,
119+
const count = await getRowCount(
120+
db,
121+
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables),
122+
rowCheck,
106123
);
107-
const count = Number(res[0].count);
108124
if (count > 0) {
109-
infoToPrint.push(`· You're about to delete ${chalk.underline(statement.tableName)} table with ${count} items`);
125+
infoToPrint.push(
126+
`· You're about to delete ${chalk.underline(statement.tableName)} table with ${
127+
describeRowCount(count, rowCheck)
128+
}`,
129+
);
110130
// statementsToExecute.push(
111131
// `truncate table ${tableNameWithSchemaFrom(statement)} cascade;`
112132
// );
113133
tablesToRemove.push(statement.tableName);
114134
shouldAskForApprove = true;
115135
}
116136
} else if (statement.type === 'drop_view' && statement.materialized) {
117-
const res = await db.query(`select count(*) as count from "${statement.schema ?? 'public'}"."${statement.name}"`);
118-
const count = Number(res[0].count);
137+
const count = await getRowCount(db, `"${statement.schema ?? 'public'}"."${statement.name}"`, rowCheck);
119138
if (count > 0) {
120139
infoToPrint.push(
121-
`· You're about to delete "${chalk.underline(statement.name)}" materialized view with ${count} items`,
140+
`· You're about to delete "${chalk.underline(statement.name)}" materialized view with ${
141+
describeRowCount(count, rowCheck)
142+
}`,
122143
);
123144

124145
matViewsToRemove.push(statement.name);
125146
shouldAskForApprove = true;
126147
}
127148
} else if (statement.type === 'alter_table_drop_column') {
128-
const res = await db.query(
129-
`select count(*) as count from ${
130-
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables)
131-
}`,
149+
const count = await getRowCount(
150+
db,
151+
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables),
152+
rowCheck,
132153
);
133-
const count = Number(res[0].count);
134154
if (count > 0) {
135155
infoToPrint.push(
136156
`· You're about to delete ${
137157
chalk.underline(statement.columnName)
138-
} column in ${statement.tableName} table with ${count} items`,
158+
} column in ${statement.tableName} table with ${describeRowCount(count, rowCheck)}`,
139159
);
140160
columnsToRemove.push(`${statement.tableName}_${statement.columnName}`);
141161
shouldAskForApprove = true;
@@ -151,12 +171,11 @@ export const pgSuggestions = async (
151171
shouldAskForApprove = true;
152172
}
153173
} else if (statement.type === 'alter_table_alter_column_set_type') {
154-
const res = await db.query(
155-
`select count(*) as count from ${
156-
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables)
157-
}`,
174+
const count = await getRowCount(
175+
db,
176+
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables),
177+
rowCheck,
158178
);
159-
const count = Number(res[0].count);
160179
if (count > 0) {
161180
infoToPrint.push(
162181
`· You're about to change ${chalk.underline(statement.columnName)} column type from ${
@@ -165,7 +184,7 @@ export const pgSuggestions = async (
165184
chalk.underline(
166185
statement.newDataType,
167186
)
168-
} with ${count} items`,
187+
} with ${describeRowCount(count, rowCheck)}`,
169188
);
170189
statementsToExecute.push(
171190
`truncate table ${
@@ -176,12 +195,11 @@ export const pgSuggestions = async (
176195
shouldAskForApprove = true;
177196
}
178197
} else if (statement.type === 'alter_table_alter_column_drop_pk') {
179-
const res = await db.query(
180-
`select count(*) as count from ${
181-
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables)
182-
}`,
198+
const count = await getRowCount(
199+
db,
200+
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables),
201+
rowCheck,
183202
);
184-
const count = Number(res[0].count);
185203
if (count > 0) {
186204
infoToPrint.push(
187205
`· You're about to change ${
@@ -216,17 +234,16 @@ export const pgSuggestions = async (
216234
continue;
217235
} else if (statement.type === 'alter_table_add_column') {
218236
if (statement.column.notNull && typeof statement.column.default === 'undefined') {
219-
const res = await db.query(
220-
`select count(*) as count from ${
221-
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables)
222-
}`,
237+
const count = await getRowCount(
238+
db,
239+
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables),
240+
rowCheck,
223241
);
224-
const count = Number(res[0].count);
225242
if (count > 0) {
226243
infoToPrint.push(
227244
`· You're about to add not-null ${
228245
chalk.underline(statement.column.name)
229-
} column without default value, which contains ${count} items`,
246+
} column without default value, which contains ${describeRowCount(count, rowCheck)}`,
230247
);
231248

232249
tablesToTruncate.push(statement.tableName);
@@ -240,20 +257,21 @@ export const pgSuggestions = async (
240257
}
241258
}
242259
} else if (statement.type === 'create_unique_constraint') {
243-
const res = await db.query(
244-
`select count(*) as count from ${
245-
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables)
246-
}`,
260+
const count = await getRowCount(
261+
db,
262+
tableNameWithSchemaFrom(statement.schema, statement.tableName, renamedSchemas, renamedTables),
263+
rowCheck,
247264
);
248-
const count = Number(res[0].count);
249265
if (count > 0) {
250266
const unsquashedUnique = PgSquasher.unsquashUnique(statement.data);
251267
console.log(
252268
`· You're about to add ${
253269
chalk.underline(
254270
unsquashedUnique.name,
255271
)
256-
} 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 ${
272+
} unique constraint to the table, which contains ${
273+
describeRowCount(count, rowCheck)
274+
}. If this statement fails, you will receive an error from the database. Do you want to truncate ${
257275
chalk.underline(
258276
statement.tableName,
259277
)

‎drizzle-kit/tests/pgSuggestions.test.ts‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
import { describe, expect, test, vi } from 'vitest';
2+
import { pgSuggestions } from '../src/cli/commands/pgPushUtils';
3+
import type { JsonStatement } from '../src/jsonStatements';
4+
5+
const dropTable: JsonStatement = {
6+
type: 'drop_table',
7+
tableName: 'users',
8+
schema: 'public',
9+
};
10+
11+
describe('pgSuggestions row checks', () => {
12+
test('counts rows by default', async () => {
13+
const query = vi.fn().mockResolvedValue([{ count: '42' }]);
14+
15+
const result = await pgSuggestions({ query }, [dropTable]);
16+
17+
expect(query).toHaveBeenCalledWith('select count(*) as count from "public"."users"');
18+
expect(result.infoToPrint).toEqual(["· You're about to delete users table with 42 items"]);
19+
});
20+
21+
test('can stop after the first row', async () => {
22+
const query = vi.fn().mockResolvedValue([{}]);
23+
24+
const result = await pgSuggestions({ query }, [dropTable], undefined, 'exists');
25+
26+
expect(query).toHaveBeenCalledWith('select 1 from "public"."users" limit 1');
27+
expect(result.infoToPrint).toEqual(["· You're about to delete users table with existing items"]);
28+
});
29+
});

0 commit comments

Comments
 (0)