Batch PostgreSQL constraint introspection

changelogs/drizzle-kit/0.32.4.md

@@ -3,3 +3,5 @@
 Added columns are now created before primary keys that use them. Unique constraints and indexes are now created before dependent foreign keys.
 
 PostgreSQL introspection now preserves JSON and JSONB expression defaults instead of parsing them as JSON literals.
+
+PostgreSQL introspection can preload primary and unique constraints once per schema to reduce query volume for large databases.

drizzle-kit/src/api.ts

@@ -17,6 +17,7 @@ import { originUUID } from './global';
 import { MySqlSchema as MySQLSchemaKit } from './serializer/mysqlSchema';
 import { PgSchema as PgSchemaKit, pgSchema, Role, squashPgScheme, View } from './serializer/pgSchema';
 import { fromDatabase } from './serializer/pgSerializer';
+import type { PgIntrospectionOptions } from './serializer/pgSerializer';
 import { SingleStoreSchema as SingleStoreSchemaKit } from './serializer/singlestoreSchema';
 import { SQLiteSchema as SQLiteSchemaKit } from './serializer/sqliteSchema';
 import { ProxyParams } from './serializer/studio';
@@ -40,6 +41,7 @@ export type DrizzlePgDBIntrospectSchema = Omit<
 	PgSchemaKit,
 	'internal'
 >;
+export type DrizzlePgIntrospectionOptions = PgIntrospectionOptions;
 
 function createConcurrencyLimiter(concurrency?: number) {
 	if (concurrency === undefined) {
@@ -84,6 +86,7 @@ export const introspectPgDB = async (
 	db: DrizzlePgDB,
 	filters: string[],
 	schemaFilters: string[],
+	options: DrizzlePgIntrospectionOptions = {},
 ): Promise<DrizzlePgDBIntrospectSchema> => {
 	const matchers = filters.map((it) => {
 		return new Minimatch(it);
@@ -119,6 +122,7 @@ export const introspectPgDB = async (
 		undefined,
 		undefined,
 		undefined,
+		options,
 	);
 
 	const schema = { id: originUUID, prevId: '', ...res } as PgSchemaKit;

drizzle-kit/src/serializer/pgSerializer.ts

@@ -966,6 +966,11 @@ function prepareRoles(entities?: {
 	return { useRoles, includeRoles, excludeRoles };
 }
 
+export type PgIntrospectionOptions = {
+	batchConstraintQueries?: boolean;
+	onConstraintPreloadError?: (error: unknown) => void;
+};
+
 export const fromDatabase = async (
 	db: DB,
 	tablesFilter: (table: string) => boolean = () => true,
@@ -983,11 +988,66 @@ export const fromDatabase = async (
 		status: IntrospectStatus,
 	) => void,
 	tsSchema?: PgSchemaInternal,
+	options: PgIntrospectionOptions = {},
 ): Promise<PgSchemaInternal> => {
 	const result: Record<string, Table> = {};
 	const views: Record<string, View> = {};
 	const policies: Record<string, Policy> = {};
 	const internals: PgKitInternals = { tables: {} };
+	const constraintsSql = `
+		SELECT rel.relname AS table_name, att.attname::text AS column_name,
+			CASE con.contype WHEN 'p' THEN 'PRIMARY KEY' WHEN 'u' THEN 'UNIQUE' END AS constraint_type,
+			con.conname AS constraint_name
+		FROM pg_catalog.pg_constraint con
+		JOIN pg_catalog.pg_class rel ON rel.oid = con.conrelid
+		JOIN pg_catalog.pg_namespace nsp ON nsp.oid = rel.relnamespace
+		CROSS JOIN LATERAL unnest(con.conkey) WITH ORDINALITY AS key_pair(attnum, ordinality)
+		JOIN pg_catalog.pg_attribute att ON att.attrelid = con.conrelid AND att.attnum = key_pair.attnum
+		WHERE nsp.nspname = $1 AND con.contype IN ('p', 'u')
+			AND pg_has_role(rel.relowner, 'USAGE')
+		ORDER BY rel.relname, con.conname, key_pair.ordinality`;
+	const primaryKeyNamesSql = `
+		SELECT rel.relname AS table_name, con.conname AS primary_key
+		FROM pg_catalog.pg_constraint con
+		JOIN pg_catalog.pg_class rel ON rel.oid = con.conrelid
+		WHERE con.contype = 'p' AND con.connamespace = $1::regnamespace`;
+
+	function createBatchLoader<Row extends { table_name: string }>(sql: string) {
+		const preloads = new Map<string, Promise<Map<string, Row[]> | undefined>>();
+
+		return async (schema: string, table: string): Promise<Row[] | undefined> => {
+			if (!options.batchConstraintQueries) return undefined;
+
+			let pending = preloads.get(schema);
+			if (!pending) {
+				pending = db.query<Row>(sql, [schema]).then((rows) => {
+					const byTable = new Map<string, Row[]>();
+					for (const row of rows) {
+						const existing = byTable.get(row.table_name);
+						if (existing) existing.push(row);
+						else byTable.set(row.table_name, [row]);
+					}
+
+					return byTable;
+				}).catch((error) => {
+					options.onConstraintPreloadError?.(error);
+					return undefined;
+				});
+				preloads.set(schema, pending);
+			}
+
+			const byTable = await pending;
+			return byTable === undefined ? undefined : byTable.get(table) ?? [];
+		};
+	}
+
+	const loadConstraints = createBatchLoader<{
+		table_name: string;
+		column_name: string;
+		constraint_type: 'PRIMARY KEY' | 'UNIQUE';
+		constraint_name: string;
+	}>(constraintsSql);
+	const loadPrimaryKeyNames = createBatchLoader<{ table_name: string; primary_key: string }>(primaryKeyNamesSql);
 
 	const where = schemaFilters.map((t) => `n.nspname = '${t}'`).join(' or ');
 
@@ -1228,7 +1288,7 @@ WHERE
 
 					const tableResponse = await getColumnsInfoQuery({ schema: tableSchema, table: tableName, db });
 
-					const tableConstraints = await db.query(
+					const tableConstraints = await loadConstraints(tableSchema, tableName) ?? await db.query(
 						`SELECT c.column_name, c.data_type, constraint_type, constraint_name, constraint_schema
       FROM information_schema.table_constraints tc
       JOIN information_schema.constraint_column_usage AS ccu USING (constraint_schema, constraint_name)
@@ -1405,7 +1465,7 @@ WHERE
 						const cprimaryKey = tableConstraints.filter((mapRow) => mapRow.constraint_type === 'PRIMARY KEY');
 
 						if (cprimaryKey.length > 1) {
-							const tableCompositePkName = await db.query(
+							const tableCompositePkName = await loadPrimaryKeyNames(tableSchema, tableName) ?? await db.query(
 								`SELECT conname AS primary_key
             FROM   pg_constraint join pg_class on (pg_class.oid = conrelid)
             WHERE  contype = 'p' 

drizzle-kit/tests/pg-constraint-batching.test.ts

@@ -0,0 +1,107 @@
+import { expect, test, vi } from 'vitest';
+import { introspectPgDB } from '../src/api';
+import type { DrizzlePgDB } from '../src/api';
+
+const columnRow = {
+	is_nullable: 'NO',
+	array_dimensions: 0,
+	data_type: 'integer',
+	seq_name: null,
+	column_default: null,
+	additional_dt: 'integer',
+	enum_name: 'int4',
+	is_generated: 'NEVER',
+	generation_expression: null,
+	is_identity: 'NO',
+	identity_generation: null,
+	identity_start: null,
+	identity_increment: null,
+	identity_maximum: null,
+	identity_minimum: null,
+	identity_cycle: 'NO',
+	type_schema: 'pg_catalog',
+};
+
+const tables = ['first', 'second'];
+
+function constraintsFor(table: string) {
+	return ['a', 'b'].map((columnName) => ({
+		table_name: table,
+		column_name: columnName,
+		constraint_type: 'PRIMARY KEY',
+		constraint_name: `${table}_pkey`,
+	}));
+}
+
+function database(failPreloads = false) {
+	const queries: string[] = [];
+	const query = async <T>(sql: string, params?: unknown[]): Promise<T[]> => {
+		queries.push(sql);
+		const constraintPreload = sql.includes("pg_has_role(rel.relowner, 'USAGE')");
+		const primaryKeyPreload = sql.includes('rel.relname AS table_name, con.conname AS primary_key');
+		if (failPreloads && (constraintPreload || primaryKeyPreload)) {
+			throw new Error('preload failed');
+		}
+		if (sql.includes("c.relkind IN ('r', 'v', 'm')")) {
+			return tables.map((table_name) => ({
+				table_schema: 'public',
+				table_name,
+				type: 'table',
+				rls_enabled: false,
+			})) as T[];
+		}
+		if (sql.includes('a.attndims AS array_dimensions')) {
+			const table = sql.match(/cls\.relname = '([^']+)'/)?.[1];
+			return ['a', 'b'].map((column_name) => ({ ...columnRow, table_name: table, column_name })) as T[];
+		}
+		if (constraintPreload) {
+			return tables.flatMap(constraintsFor) as T[];
+		}
+		if (primaryKeyPreload) {
+			return tables.map((table_name) => ({ table_name, primary_key: `${table_name}_pkey` })) as T[];
+		}
+		if (sql.includes('information_schema.constraint_column_usage')) {
+			const table = sql.match(/tc\.table_name = '([^']+)'/)?.[1];
+			return constraintsFor(table ?? '') as T[];
+		}
+		if (sql.includes('SELECT conname AS primary_key')) {
+			return [{ primary_key: `${params?.[1]}_pkey` }] as T[];
+		}
+
+		return [];
+	};
+
+	return { db: { query } as DrizzlePgDB, queries };
+}
+
+test('preloads constraints once per schema only when enabled', async () => {
+	const baselineDatabase = database();
+	const baseline = await introspectPgDB(baselineDatabase.db, [], ['public']);
+	const batchedDatabase = database();
+	const batched = await introspectPgDB(batchedDatabase.db, [], ['public'], {
+		batchConstraintQueries: true,
+	});
+
+	expect(batched).toEqual(baseline);
+	expect(batchedDatabase.queries.filter((sql) => sql.includes("pg_has_role(rel.relowner, 'USAGE')"))).toHaveLength(1);
+	expect(
+		batchedDatabase.queries.filter((sql) => sql.includes('rel.relname AS table_name, con.conname AS primary_key')),
+	).toHaveLength(1);
+	expect(batchedDatabase.queries.filter((sql) => sql.includes('information_schema.constraint_column_usage')))
+		.toHaveLength(0);
+});
+
+test('falls back to per-table queries when preloading fails', async () => {
+	const baseline = await introspectPgDB(database().db, [], ['public']);
+	const fallbackDatabase = database(true);
+	const onConstraintPreloadError = vi.fn();
+	const fallback = await introspectPgDB(fallbackDatabase.db, [], ['public'], {
+		batchConstraintQueries: true,
+		onConstraintPreloadError,
+	});
+
+	expect(fallback).toEqual(baseline);
+	expect(onConstraintPreloadError).toHaveBeenCalledTimes(2);
+	expect(fallbackDatabase.queries.filter((sql) => sql.includes('information_schema.constraint_column_usage')))
+		.toHaveLength(2);
+});