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

@@ -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"

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,
 }

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,
+}

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
+`

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<ReturnType<typeof createTestDatabase>>) => Promise<void>
+) => {
+  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)
+})