diff --git a/changelogs/drizzle-kit/0.32.4.md b/changelogs/drizzle-kit/0.32.4.md index 491987e8a..685b514ce 100644 --- a/changelogs/drizzle-kit/0.32.4.md +++ b/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. diff --git a/drizzle-kit/src/api.ts b/drizzle-kit/src/api.ts index 70960756a..204a5571f 100644 --- a/drizzle-kit/src/api.ts +++ b/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 => { 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; diff --git a/drizzle-kit/src/serializer/pgSerializer.ts b/drizzle-kit/src/serializer/pgSerializer.ts index 60883596e..350109c87 100644 --- a/drizzle-kit/src/serializer/pgSerializer.ts +++ b/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 => { const result: Record = {}; const views: Record = {}; const policies: Record = {}; 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(sql: string) { + const preloads = new Map | undefined>>(); + + return async (schema: string, table: string): Promise => { + if (!options.batchConstraintQueries) return undefined; + + let pending = preloads.get(schema); + if (!pending) { + pending = db.query(sql, [schema]).then((rows) => { + const byTable = new Map(); + 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' diff --git a/drizzle-kit/tests/pg-constraint-batching.test.ts b/drizzle-kit/tests/pg-constraint-batching.test.ts new file mode 100644 index 000000000..40923ff97 --- /dev/null +++ b/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 (sql: string, params?: unknown[]): Promise => { + 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); +});