chore(pg-meta): migrate tests for roles (#33679)

* chore(pg-meta): migrate tests for roles

* chore: remove unused export

packages/pg-meta/src/pg-meta-roles.ts

@@ -20,6 +20,8 @@ const pgRoleZod = z.object({
 const pgRoleArrayZod = z.array(pgRoleZod)
 const pgRoleOptionalZod = z.optional(pgRoleZod)
 
+export type PGRole = z.infer<typeof pgRoleZod>
+
 function list({
   includeDefaultRoles: includeDefaultRoles = false,
   limit,
@@ -65,24 +67,25 @@ where
   }
 }
 
-function retrieve({ id }: { id: number }): { sql: string; zod: typeof pgRoleOptionalZod }
-function retrieve({ name }: { name: string }): { sql: string; zod: typeof pgRoleOptionalZod }
-function retrieve({ id, name }: { id?: number; name?: string }): {
+type RoleIdentifier = Pick<PGRole, 'id'> | Pick<PGRole, 'name'>
+
+function getIdentifierWhereClause(identifier: RoleIdentifier) {
+  if ('id' in identifier && identifier.id) {
+    return `${ident('id')} = ${literal(identifier.id)}`
+  } else if ('name' in identifier && identifier.name) {
+    return `${ident('name')} = ${literal(identifier.name)}`
+  }
+  throw new Error('Must provide either id or name')
+}
+
+function retrieve(identifier: RoleIdentifier): {
   sql: string
   zod: typeof pgRoleOptionalZod
 } {
-  if (id) {
-    const sql = `${ROLES_SQL} where r.oid = ${literal(id)};`
-    return {
-      sql,
-      zod: pgRoleOptionalZod,
-    }
-  } else {
-    const sql = `${ROLES_SQL} where rolname = ${literal(name)};`
-    return {
-      sql,
-      zod: pgRoleOptionalZod,
-    }
+  const sql = `with roles as (${ROLES_SQL}) select * from roles where ${getIdentifierWhereClause(identifier)};`
+  return {
+    sql,
+    zod: pgRoleOptionalZod,
   }
 }
 
@@ -156,17 +159,8 @@ type RoleUpdateParams = {
   password?: string
   validUntil?: string
 }
-function update({ id }: { id: number }, params: RoleUpdateParams): { sql: string }
-function update({ name }: { name: string }, params: RoleUpdateParams): { sql: string }
-function update(
-  {
-    id,
-    name,
-  }: {
-    id?: number
-    name?: string
-  },
-  {
+function update(identifier: RoleIdentifier, params: RoleUpdateParams): { sql: string } {
+  const {
     name: newName,
     isSuperuser,
     canCreateDb,
@@ -178,15 +172,14 @@ function update(
     connectionLimit,
     password,
     validUntil,
-  }: RoleUpdateParams
-): { sql: string } {
+  } = params
   const sql = `
 do $$
 declare
-  id oid := ${id === undefined ? `${literal(name)}::regrole` : literal(id)};
   old record;
 begin
-  select * into old from pg_roles where oid = id;
+  with roles as (${ROLES_SQL})
+  select * into old from roles where ${getIdentifierWhereClause(identifier)};
   if old is null then
     raise exception 'Cannot find role with id %', id;
   end if;
@@ -199,18 +192,18 @@ begin
     ${canLogin === undefined ? '' : canLogin ? 'login' : 'nologin'}
     ${isReplicationRole === undefined ? '' : isReplicationRole ? 'replication' : 'noreplication'}
     ${canBypassRls === undefined ? '' : canBypassRls ? 'bypassrls' : 'nobypassrls'}
-    ${connectionLimit === undefined ? '' : `connection limit ${literal(connectionLimit)}`}
+    ${connectionLimit === undefined ? '' : `connection limit ${connectionLimit}`}
     ${password === undefined ? '' : `password ${literal(password)}`}
-    ${validUntil === undefined ? '' : `valid until ${literal(validUntil)}`}
-  ', old.rolname));
+    ${validUntil === undefined ? '' : `valid until %L`}
+  ', old.name${validUntil === undefined ? '' : `, ${literal(validUntil)}`}));
 
   ${
     newName === undefined
       ? ''
       : `
   -- Using the same name in the rename clause gives an error, so only do it if the new name is different.
-  if new_name != old.nspname then
-    execute(format('alter role %I rename to ${ident(newName)};', old.nspname));
+  if ${literal(newName)} != old.name then
+    execute(format('alter role %I rename to %I;', old.name, ${literal(newName)}));
   end if;
   `
   }
@@ -223,30 +216,22 @@ $$;
 type RoleRemoveParams = {
   ifExists?: boolean
 }
-function remove({ id }: { id: number }, params?: RoleRemoveParams): { sql: string }
-function remove({ name }: { name: string }, params?: RoleRemoveParams): { sql: string }
 function remove(
-  {
-    id,
-    name,
-  }: {
-    id?: number
-    name?: string
-  },
+  identifier: RoleIdentifier,
   { ifExists = false }: RoleRemoveParams = {}
 ): { sql: string } {
   const sql = `
 do $$
 declare
-  id oid := ${id === undefined ? `${literal(name)}::regrole` : literal(id)};
   old record;
 begin
-  select * into old from pg_roles where oid = id;
+  with roles as (${ROLES_SQL})
+  select * into old from roles where ${getIdentifierWhereClause(identifier)};
   if old is null then
     raise exception 'Cannot find role with id %', id;
   end if;
 
-  execute(format('drop role ${ifExists ? 'if exists' : ''} %I;', old.rolname));
+  execute(format('drop role ${ifExists ? 'if exists' : ''} %I;', old.name));
 end
 $$;
 `

packages/pg-meta/test/roles.test.ts

@@ -0,0 +1,238 @@
+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 roles', async ({ executeQuery }) => {
+  const { sql, zod } = await pgMeta.roles.list()
+  const res = zod.parse(await executeQuery(sql))
+
+  let role = res.find(({ name }) => name === 'postgres')
+
+  expect(role).toMatchInlineSnapshot(
+    { activeConnections: expect.any(Number), id: expect.any(Number) },
+    `
+    {
+      "activeConnections": Any<Number>,
+      "canBypassRls": true,
+      "canCreateDb": true,
+      "canCreateRole": true,
+      "canLogin": true,
+      "config": {},
+      "connectionLimit": 100,
+      "id": Any<Number>,
+      "inheritRole": true,
+      "isReplicationRole": true,
+      "isSuperuser": true,
+      "name": "postgres",
+      "validUntil": null,
+    }
+  `
+  )
+
+  // pg_monitor is a predefined role. `includeDefaultRoles` defaults to false,
+  // so it shouldn't be included in the result.
+  role = res.find(({ name }) => name === 'pg_monitor')
+
+  expect(role).toMatchInlineSnapshot(`undefined`)
+})
+
+withTestDatabase('list roles w/ default roles', async ({ executeQuery }) => {
+  const { sql, zod } = await pgMeta.roles.list({ includeDefaultRoles: true })
+  const res = zod.parse(await executeQuery(sql))
+
+  const role = res.find(({ name }) => name === 'pg_monitor')
+
+  expect(role).toMatchInlineSnapshot(
+    {
+      activeConnections: expect.any(Number),
+      id: expect.any(Number),
+    },
+    `
+    {
+      "activeConnections": Any<Number>,
+      "canBypassRls": false,
+      "canCreateDb": false,
+      "canCreateRole": false,
+      "canLogin": false,
+      "config": {},
+      "connectionLimit": 100,
+      "id": Any<Number>,
+      "inheritRole": true,
+      "isReplicationRole": false,
+      "isSuperuser": false,
+      "name": "pg_monitor",
+      "validUntil": null,
+    }
+  `
+  )
+})
+
+withTestDatabase('retrieve, create, update, delete roles', async ({ executeQuery }) => {
+  // Create role
+  const { sql: createSql } = pgMeta.roles.create({
+    name: 'r1',
+    isSuperuser: true,
+    canCreateDb: true,
+    canCreateRole: true,
+    inheritRole: false,
+    canLogin: true,
+    isReplicationRole: true,
+    canBypassRls: true,
+    connectionLimit: 100,
+    validUntil: '2020-01-01T00:00:00.000Z',
+    config: { search_path: 'extension, public' },
+  })
+  await executeQuery(createSql)
+
+  // Retrieve the created role using list
+  const { sql: listSql, zod: listZod } = await pgMeta.roles.list()
+  const roles = listZod.parse(await executeQuery(listSql))
+  const createdRole = roles.find((role) => role.name === 'r1')
+  expect(createdRole).toMatchInlineSnapshot(
+    { id: expect.any(Number), activeConnections: expect.any(Number) },
+    `
+    {
+      "activeConnections": Any<Number>,
+      "canBypassRls": true,
+      "canCreateDb": true,
+      "canCreateRole": true,
+      "canLogin": true,
+      "config": {
+        "search_path": ""extension, public"",
+      },
+      "connectionLimit": 100,
+      "id": Any<Number>,
+      "inheritRole": false,
+      "isReplicationRole": true,
+      "isSuperuser": true,
+      "name": "r1",
+      "validUntil": "2020-01-01 00:00:00+00",
+    }
+  `
+  )
+
+  // Remove role
+  const { sql: removeSql } = pgMeta.roles.remove({ id: createdRole!.id })
+  await executeQuery(removeSql)
+
+  // Create a new role for update test
+  const { sql: createNewSql } = pgMeta.roles.create({
+    name: 'r1',
+  })
+  await executeQuery(createNewSql)
+
+  // Get the role ID for update
+  const { sql: getIdSql, zod: getIdZod } = await pgMeta.roles.list()
+  const roleForUpdate = getIdZod
+    .parse(await executeQuery(getIdSql))
+    .find((role) => role.name === 'r1')
+
+  // Update role with ISO string date
+  const { sql: updateSql } = pgMeta.roles.update(
+    { id: roleForUpdate!.id },
+    {
+      name: 'rr',
+      isSuperuser: true,
+      canCreateDb: true,
+      canCreateRole: true,
+      inheritRole: false,
+      canLogin: true,
+      isReplicationRole: true,
+      canBypassRls: true,
+      connectionLimit: 100,
+      validUntil: '2020-01-01T00:00:00.000Z',
+    }
+  )
+  await executeQuery(updateSql)
+
+  // Verify update using retrieve
+  const { sql: retrieveUpdatedSql, zod: retrieveZod } = pgMeta.roles.retrieve({
+    id: roleForUpdate!.id,
+  })
+  const updatedRole = retrieveZod.parse((await executeQuery(retrieveUpdatedSql))[0])
+  expect(updatedRole).toMatchInlineSnapshot(
+    { id: expect.any(Number), activeConnections: expect.any(Number) },
+    `
+    {
+      "activeConnections": Any<Number>,
+      "canBypassRls": true,
+      "canCreateDb": true,
+      "canCreateRole": true,
+      "canLogin": true,
+      "config": {},
+      "connectionLimit": 100,
+      "id": Any<Number>,
+      "inheritRole": false,
+      "isReplicationRole": true,
+      "isSuperuser": true,
+      "name": "rr",
+      "validUntil": "2020-01-01 00:00:00+00",
+    }
+  `
+  )
+
+  // Create role with config
+  const { sql: createConfigSql } = pgMeta.roles.create({
+    name: 'config_role',
+    config: { search_path: 'public', log_statement: 'all' },
+  })
+  await executeQuery(createConfigSql)
+
+  // Verify config role using list
+  const { sql: listConfigSql, zod: listConfigZod } = await pgMeta.roles.list()
+  const configRole = listConfigZod
+    .parse(await executeQuery(listConfigSql))
+    .find((role) => role.name === 'config_role')
+  expect(configRole).toMatchInlineSnapshot(
+    { id: expect.any(Number), activeConnections: expect.any(Number) },
+    `
+    {
+      "activeConnections": Any<Number>,
+      "canBypassRls": false,
+      "canCreateDb": false,
+      "canCreateRole": false,
+      "canLogin": false,
+      "config": {
+        "log_statement": "all",
+        "search_path": "public",
+      },
+      "connectionLimit": 100,
+      "id": Any<Number>,
+      "inheritRole": true,
+      "isReplicationRole": false,
+      "isSuperuser": false,
+      "name": "config_role",
+      "validUntil": null,
+    }
+  `
+  )
+
+  // Remove role and verify it's gone
+  const { sql: finalRemoveSql } = pgMeta.roles.remove({ id: configRole!.id })
+  await executeQuery(finalRemoveSql)
+
+  const { sql: finalListSql, zod: finalListZod } = await pgMeta.roles.list()
+  const finalRoles = finalListZod.parse(await executeQuery(finalListSql))
+  expect(finalRoles.find((role) => role.name === 'config_role')).toBeUndefined()
+})