From 06ef94a858b06b3481d3613ddfb8aa0d2fa8cedd Mon Sep 17 00:00:00 2001 From: Andrew Valleteau Date: Tue, 4 Mar 2025 10:42:18 +0100 Subject: [PATCH] chore(pg-meta): migrate indexes (#33645) * chore(pg-meta): migrate indexes * chore(pg-meta): fix vector to array conversion * chore: use common naming --- packages/pg-meta/package.json | 3 +- packages/pg-meta/src/index.ts | 2 + packages/pg-meta/src/pg-meta-indexes.ts | 101 ++++++++++++++ packages/pg-meta/src/sql/indexes.ts | 46 +++++++ packages/pg-meta/test/indexes.test.ts | 167 ++++++++++++++++++++++++ 5 files changed, 318 insertions(+), 1 deletion(-) create mode 100644 packages/pg-meta/src/pg-meta-indexes.ts create mode 100644 packages/pg-meta/src/sql/indexes.ts create mode 100644 packages/pg-meta/test/indexes.test.ts diff --git a/packages/pg-meta/package.json b/packages/pg-meta/package.json index e46a2568b72..d47f1e32210 100644 --- a/packages/pg-meta/package.json +++ b/packages/pg-meta/package.json @@ -11,7 +11,8 @@ "test": "run-s db:clean db:run test:run db:clean", "db:clean": "cd test/db && docker compose down", "db:run": "cd test/db && docker compose up --detach --wait", - "test:run": "vitest run" + "test:run": "vitest run", + "test:update": "vitest run --update" }, "dependencies": { "zod": "^3.22.4" diff --git a/packages/pg-meta/src/index.ts b/packages/pg-meta/src/index.ts index 491c62e096b..9ec042ced89 100644 --- a/packages/pg-meta/src/index.ts +++ b/packages/pg-meta/src/index.ts @@ -2,6 +2,7 @@ import roles from './pg-meta-roles' import schemas from './pg-meta-schemas' import * as functions from './pg-meta-functions' import tablePrivileges from './pg-meta-table-privileges' +import indexes from './pg-meta-indexes' import columnPrivileges from './pg-meta-column-privileges' export default { @@ -9,5 +10,6 @@ export default { schemas, functions, tablePrivileges, + indexes, columnPrivileges, } diff --git a/packages/pg-meta/src/pg-meta-indexes.ts b/packages/pg-meta/src/pg-meta-indexes.ts new file mode 100644 index 00000000000..1a2b854c7c5 --- /dev/null +++ b/packages/pg-meta/src/pg-meta-indexes.ts @@ -0,0 +1,101 @@ +import { literal } from './pg-format' +import { DEFAULT_SYSTEM_SCHEMAS } from './constants' +import { filterByList } from './helpers' +import { INDEXES_SQL } from './sql/indexes' +import { z } from 'zod' + +const pgIndexZod = z.object({ + id: z.number(), + table_id: z.number(), + schema: z.string(), + number_of_attributes: z.number(), + number_of_key_attributes: z.number(), + is_unique: z.boolean(), + is_primary: z.boolean(), + is_exclusion: z.boolean(), + is_immediate: z.boolean(), + is_clustered: z.boolean(), + is_valid: z.boolean(), + check_xmin: z.boolean(), + is_ready: z.boolean(), + is_live: z.boolean(), + is_replica_identity: z.boolean(), + key_attributes: z.array(z.number()), + collation: z.array(z.number()), + class: z.array(z.number()), + options: z.array(z.number()), + index_predicate: z.string().nullable(), + comment: z.string().nullable(), + index_definition: z.string(), + access_method: z.string(), + index_attributes: z.array( + z.object({ + attribute_number: z.number(), + attribute_name: z.string(), + data_type: z.string(), + }) + ), +}) + +const pgIndexArrayZod = z.array(pgIndexZod) +const pgIndexOptionalZod = z.optional(pgIndexZod) + +function list({ + includeSystemSchemas = false, + includedSchemas, + excludedSchemas, + limit, + offset, +}: { + includeSystemSchemas?: boolean + includedSchemas?: string[] + excludedSchemas?: string[] + limit?: number + offset?: number +} = {}): { + sql: string + zod: typeof pgIndexArrayZod +} { + let sql = ` + with indexes as (${INDEXES_SQL}) + select * + from indexes + ` + const filter = filterByList( + includedSchemas, + excludedSchemas, + !includeSystemSchemas ? DEFAULT_SYSTEM_SCHEMAS : undefined + ) + if (filter) { + sql += ` where schema ${filter}` + } + if (limit) { + sql += ` limit ${limit}` + } + if (offset) { + sql += ` offset ${offset}` + } + return { + sql, + zod: pgIndexArrayZod, + } +} + +function retrieve({ id }: { id: number }): { sql: string; zod: typeof pgIndexOptionalZod } { + const sql = ` + with indexes as (${INDEXES_SQL}) + select * + from indexes + where id = ${literal(id)}; + ` + return { + sql, + zod: pgIndexOptionalZod, + } +} + +export default { + list, + retrieve, + zod: pgIndexZod, +} diff --git a/packages/pg-meta/src/sql/indexes.ts b/packages/pg-meta/src/sql/indexes.ts new file mode 100644 index 00000000000..ef7cd40e5f1 --- /dev/null +++ b/packages/pg-meta/src/sql/indexes.ts @@ -0,0 +1,46 @@ +export const INDEXES_SQL = /* SQL */ ` + SELECT + idx.indexrelid::int8 AS id, + idx.indrelid::int8 AS table_id, + n.nspname AS schema, + idx.indnatts AS number_of_attributes, + idx.indnkeyatts AS number_of_key_attributes, + idx.indisunique AS is_unique, + idx.indisprimary AS is_primary, + idx.indisexclusion AS is_exclusion, + idx.indimmediate AS is_immediate, + idx.indisclustered AS is_clustered, + idx.indisvalid AS is_valid, + idx.indcheckxmin AS check_xmin, + idx.indisready AS is_ready, + idx.indislive AS is_live, + idx.indisreplident AS is_replica_identity, + idx.indkey::smallint[] AS key_attributes, + idx.indcollation::integer[] AS collation, + idx.indclass::integer[] AS class, + idx.indoption::smallint[] AS options, + idx.indpred AS index_predicate, + obj_description(idx.indexrelid, 'pg_class') AS comment, + ix.indexdef as index_definition, + am.amname AS access_method, + jsonb_agg( + jsonb_build_object( + 'attribute_number', a.attnum, + 'attribute_name', a.attname, + 'data_type', format_type(a.atttypid, a.atttypmod) + ) + ORDER BY a.attnum + ) AS index_attributes + FROM + pg_index idx + JOIN pg_class c ON c.oid = idx.indexrelid + JOIN pg_namespace n ON c.relnamespace = n.oid + JOIN pg_am am ON c.relam = am.oid + JOIN pg_attribute a ON a.attrelid = c.oid AND a.attnum = ANY(idx.indkey) + JOIN pg_indexes ix ON c.relname = ix.indexname + GROUP BY + idx.indexrelid, idx.indrelid, n.nspname, idx.indnatts, idx.indnkeyatts, idx.indisunique, + idx.indisprimary, idx.indisexclusion, idx.indimmediate, idx.indisclustered, idx.indisvalid, + idx.indcheckxmin, idx.indisready, idx.indislive, idx.indisreplident, idx.indkey, + idx.indcollation, idx.indclass, idx.indoption, idx.indexprs, idx.indpred, ix.indexdef, am.amname +` diff --git a/packages/pg-meta/test/indexes.test.ts b/packages/pg-meta/test/indexes.test.ts new file mode 100644 index 00000000000..306e7d3774e --- /dev/null +++ b/packages/pg-meta/test/indexes.test.ts @@ -0,0 +1,167 @@ +import { expect, test, beforeAll, afterAll } from 'vitest' +import pgMeta from '../src/index' +import { createTestDatabase, cleanupRoot } from './db/utils' + +beforeAll(async () => { + // Any global setup if needed +}) + +afterAll(async () => { + await cleanupRoot() +}) + +const withTestDatabase = ( + name: string, + fn: (db: Awaited>) => Promise +) => { + test(name, async () => { + const db = await createTestDatabase() + try { + await fn(db) + } finally { + await db.cleanup() + } + }) +} + +withTestDatabase('list indexes', async ({ executeQuery }) => { + // List indexes + const { sql: listSql, zod: listZod } = await pgMeta.indexes.list() + const indexes = listZod.parse(await executeQuery(listSql)) + const usersPkeyIndex = indexes.find( + ({ index_definition }) => + index_definition === 'CREATE UNIQUE INDEX users_pkey ON public.users USING btree (id)' + )! + + expect(usersPkeyIndex).toMatchInlineSnapshot( + ` + { + "access_method": "btree", + "check_xmin": false, + "class": [ + 3124, + ], + "collation": [ + 0, + ], + "comment": null, + "id": 16399, + "index_attributes": [ + { + "attribute_name": "id", + "attribute_number": 1, + "data_type": "bigint", + }, + ], + "index_definition": "CREATE UNIQUE INDEX users_pkey ON public.users USING btree (id)", + "index_predicate": null, + "is_clustered": false, + "is_exclusion": false, + "is_immediate": true, + "is_live": true, + "is_primary": true, + "is_ready": true, + "is_replica_identity": false, + "is_unique": true, + "is_valid": true, + "key_attributes": [ + 1, + ], + "number_of_attributes": 1, + "number_of_key_attributes": 1, + "options": [ + 0, + ], + "schema": "public", + "table_id": 16393, + } + ` + ) +}) + +withTestDatabase('retrieve index', async ({ executeQuery }) => { + // Retrieve specific index + const { sql: retrieveSql, zod: retrieveZod } = await pgMeta.indexes.retrieve({ + id: 16399, + }) + const index = retrieveZod.parse((await executeQuery(retrieveSql))[0]) + expect(index).toMatchInlineSnapshot( + ` + { + "access_method": "btree", + "check_xmin": false, + "class": [ + 3124, + ], + "collation": [ + 0, + ], + "comment": null, + "id": 16399, + "index_attributes": [ + { + "attribute_name": "id", + "attribute_number": 1, + "data_type": "bigint", + }, + ], + "index_definition": "CREATE UNIQUE INDEX users_pkey ON public.users USING btree (id)", + "index_predicate": null, + "is_clustered": false, + "is_exclusion": false, + "is_immediate": true, + "is_live": true, + "is_primary": true, + "is_ready": true, + "is_replica_identity": false, + "is_unique": true, + "is_valid": true, + "key_attributes": [ + 1, + ], + "number_of_attributes": 1, + "number_of_key_attributes": 1, + "options": [ + 0, + ], + "schema": "public", + "table_id": 16393, + } + ` + ) +}) + +withTestDatabase('list with filters', async ({ executeQuery }) => { + // Test includeSystemSchemas + const { sql: withSystemSql, zod } = await pgMeta.indexes.list({ includeSystemSchemas: true }) + const withSystem = zod.parse(await executeQuery(withSystemSql)) + expect(withSystem.some((idx) => idx.schema === 'pg_catalog')).toBe(true) + + // Test without system schemas (default) + const { sql: withoutSystemSql, zod: withoutSystemZod } = await pgMeta.indexes.list() + const withoutSystem = withoutSystemZod.parse(await executeQuery(withoutSystemSql)) + expect(withoutSystem.some((idx) => idx.schema === 'pg_catalog')).toBe(false) + + // Test includedSchemas + const { sql: includedSchemasSql, zod: includedSchemasZod } = await pgMeta.indexes.list({ + includedSchemas: ['public'], + }) + const includedSchemas = includedSchemasZod.parse(await executeQuery(includedSchemasSql)) + expect(includedSchemas.every((idx) => idx.schema === 'public')).toBe(true) + + // Test excludedSchemas + const { sql: excludedSchemasSql, zod: excludedSchemasZod } = await pgMeta.indexes.list({ + excludedSchemas: ['public'], + }) + const excludedSchemas = excludedSchemasZod.parse(await executeQuery(excludedSchemasSql)) + expect(excludedSchemas.some((idx) => idx.schema === 'public')).toBe(false) + + // Test limit and offset + const { sql: limitSql, zod: limitZod } = await pgMeta.indexes.list({ limit: 1 }) + const limited = limitZod.parse(await executeQuery(limitSql)) + expect(limited).toHaveLength(1) + + const { sql: offsetSql, zod: offsetZod } = await pgMeta.indexes.list({ offset: 1 }) + const offset = offsetZod.parse(await executeQuery(offsetSql)) + expect(offset).toHaveLength(withoutSystem.length - 1) +})