perf: optimize ACL role lookups with Map indexes and caching

.changeset/optimize-acl-role-lookups.md

@@ -0,0 +1,5 @@
+---
+"pg-introspection": patch
+---
+
+Optimize ACL role lookups with Map-based indexes and WeakMap caching. Replace O(n) linear scans in `getRole()`, `getRoleByName()`, and `expandRoles()` with O(1) Map lookups, and cache `expandRoles` results per introspection object. Significantly improves performance for schemas with many roles.

utils/pg-introspection/__tests__/acl-optimized-test.ts

@@ -0,0 +1,202 @@
+import assert from "node:assert";
+import { describe, it } from "node:test";
+
+import {
+  type AclObject,
+  PUBLIC_ROLE,
+  aclContainsRole,
+  expandRoles,
+  resolvePermissions,
+  parseAcl,
+} from "../dist/acl.js";
+import type { Introspection, PgRoles } from "../dist/introspection.js";
+
+function makeRole(
+  id: string,
+  name: string,
+  opts: Partial<PgRoles> = {},
+): PgRoles {
+  return {
+    _id: id,
+    rolname: name,
+    rolsuper: false,
+    rolinherit: true,
+    rolcreaterole: false,
+    rolcreatedb: false,
+    rolcanlogin: false,
+    rolreplication: false,
+    rolconnlimit: null,
+    rolpassword: null,
+    rolbypassrls: false,
+    rolconfig: null,
+    rolvaliduntil: null,
+    ...opts,
+  };
+}
+
+function makeIntrospection(
+  roles: PgRoles[],
+  authMembers: { member: string; roleid: string }[] = [],
+): Introspection {
+  return {
+    roles,
+    auth_members: authMembers.map((am) => ({
+      ...am,
+      admin_option: false,
+      grantor: "0",
+    })),
+  } as unknown as Introspection;
+}
+
+describe("expandRoles", () => {
+  it("always includes PUBLIC_ROLE", () => {
+    const role = makeRole("1", "alice");
+    const introspection = makeIntrospection([role]);
+    const result = expandRoles(introspection, [role]);
+    assert.ok(result.includes(PUBLIC_ROLE));
+  });
+
+  it("includes the role itself", () => {
+    const role = makeRole("1", "alice");
+    const introspection = makeIntrospection([role]);
+    const result = expandRoles(introspection, [role]);
+    assert.ok(result.includes(role));
+  });
+
+  it("expands inherited roles via auth_members", () => {
+    const alice = makeRole("1", "alice");
+    const editors = makeRole("2", "editors");
+    const introspection = makeIntrospection(
+      [alice, editors],
+      [{ member: "1", roleid: "2" }], // alice is member of editors
+    );
+    const result = expandRoles(introspection, [alice]);
+    assert.ok(result.includes(editors));
+    assert.ok(result.includes(alice));
+    assert.ok(result.includes(PUBLIC_ROLE));
+  });
+
+  it("expands transitive role inheritance", () => {
+    const alice = makeRole("1", "alice");
+    const editors = makeRole("2", "editors");
+    const admins = makeRole("3", "admins");
+    const introspection = makeIntrospection(
+      [alice, editors, admins],
+      [
+        { member: "1", roleid: "2" }, // alice -> editors
+        { member: "2", roleid: "3" }, // editors -> admins
+      ],
+    );
+    const result = expandRoles(introspection, [alice]);
+    assert.ok(result.includes(admins));
+    assert.ok(result.includes(editors));
+    assert.ok(result.includes(alice));
+  });
+
+  it("does not expand roles with NOINHERIT by default", () => {
+    const alice = makeRole("1", "alice", { rolinherit: false });
+    const editors = makeRole("2", "editors");
+    const introspection = makeIntrospection(
+      [alice, editors],
+      [{ member: "1", roleid: "2" }],
+    );
+    const result = expandRoles(introspection, [alice]);
+    // alice has NOINHERIT, so editors should NOT be included
+    assert.ok(!result.includes(editors));
+    assert.ok(result.includes(alice));
+  });
+
+  it("expands NOINHERIT roles when includeNoInherit=true", () => {
+    const alice = makeRole("1", "alice", { rolinherit: false });
+    const editors = makeRole("2", "editors");
+    const introspection = makeIntrospection(
+      [alice, editors],
+      [{ member: "1", roleid: "2" }],
+    );
+    const result = expandRoles(introspection, [alice], true);
+    assert.ok(result.includes(editors));
+  });
+
+  it("returns cached result for repeated calls", () => {
+    const alice = makeRole("1", "alice");
+    const editors = makeRole("2", "editors");
+    const introspection = makeIntrospection(
+      [alice, editors],
+      [{ member: "1", roleid: "2" }],
+    );
+    const result1 = expandRoles(introspection, [alice]);
+    const result2 = expandRoles(introspection, [alice]);
+    assert.strictEqual(result1, result2); // same reference = cached
+  });
+
+  it("handles circular role membership without infinite loop", () => {
+    const a = makeRole("1", "a");
+    const b = makeRole("2", "b");
+    const introspection = makeIntrospection(
+      [a, b],
+      [
+        { member: "1", roleid: "2" },
+        { member: "2", roleid: "1" },
+      ],
+    );
+    const result = expandRoles(introspection, [a]);
+    assert.ok(result.includes(a));
+    assert.ok(result.includes(b));
+  });
+});
+
+describe("aclContainsRole", () => {
+  it("returns true when role matches ACL directly", () => {
+    const alice = makeRole("1", "alice");
+    const introspection = makeIntrospection([alice]);
+    const acl = parseAcl("alice=r/postgres");
+    assert.ok(aclContainsRole(introspection, acl, alice));
+  });
+
+  it("returns true for PUBLIC ACL and any role", () => {
+    const alice = makeRole("1", "alice");
+    const introspection = makeIntrospection([alice]);
+    const acl = parseAcl("=r/postgres");
+    assert.ok(aclContainsRole(introspection, acl, alice));
+  });
+
+  it("returns true when role inherits ACL role", () => {
+    const alice = makeRole("1", "alice");
+    const editors = makeRole("2", "editors");
+    const introspection = makeIntrospection(
+      [alice, editors],
+      [{ member: "1", roleid: "2" }],
+    );
+    const acl = parseAcl("editors=rw/postgres");
+    assert.ok(aclContainsRole(introspection, acl, alice));
+  });
+
+  it("returns false when role does not match", () => {
+    const alice = makeRole("1", "alice");
+    const bob = makeRole("2", "bob");
+    const introspection = makeIntrospection([alice, bob]);
+    const acl = parseAcl("bob=r/postgres");
+    assert.ok(!aclContainsRole(introspection, acl, alice));
+  });
+});
+
+describe("resolvePermissions", () => {
+  it("grants permissions from matching ACLs", () => {
+    const alice = makeRole("1", "alice");
+    const introspection = makeIntrospection([alice]);
+    const acls = [parseAcl("alice=rw/postgres")];
+    const perms = resolvePermissions(introspection, acls, alice);
+    assert.strictEqual(perms.select, true);
+    assert.strictEqual(perms.update, true);
+    assert.strictEqual(perms.insert, false);
+  });
+
+  it("grants all permissions to superuser", () => {
+    const superuser = makeRole("1", "super", { rolsuper: true });
+    const introspection = makeIntrospection([superuser]);
+    const perms = resolvePermissions(introspection, [], superuser);
+    assert.strictEqual(perms.select, true);
+    assert.strictEqual(perms.insert, true);
+    assert.strictEqual(perms.delete, true);
+  });
+});

utils/pg-introspection/src/acl.ts

@@ -89,14 +89,80 @@ export const PUBLIC_ROLE: PgRoles = Object.freeze({
   _id: "0",
 });
 
+// ── Optimized lookup caches (WeakMap keyed by introspection — auto-GC) ──
+
+const _roleByIdCache = new WeakMap<Introspection, Map<string, PgRoles>>();
+const _roleByNameCache = new WeakMap<Introspection, Map<string, PgRoles>>();
+const _membersByMemberIdCache = new WeakMap<
+  Introspection,
+  Map<string, string[]>
+>();
+const _expandRolesCache = new WeakMap<Introspection, Map<string, PgRoles[]>>();
+
+function _getRoleByIdMap(introspection: Introspection): Map<string, PgRoles> {
+  let map = _roleByIdCache.get(introspection);
+  if (!map) {
+    map = new Map();
+    map.set("0", PUBLIC_ROLE);
+    for (const r of introspection.roles) {
+      map.set(r._id, r);
+    }
+    _roleByIdCache.set(introspection, map);
+  }
+  return map;
+}
+
+function _getRoleByNameMap(introspection: Introspection): Map<string, PgRoles> {
+  let map = _roleByNameCache.get(introspection);
+  if (!map) {
+    map = new Map();
+    map.set("public", PUBLIC_ROLE);
+    for (const r of introspection.roles) {
+      map.set(r.rolname, r);
+    }
+    _roleByNameCache.set(introspection, map);
+  }
+  return map;
+}
+
+function _getMembersByMemberIdMap(
+  introspection: Introspection,
+): Map<string, string[]> {
+  let map = _membersByMemberIdCache.get(introspection);
+  if (!map) {
+    map = new Map();
+    for (const am of introspection.auth_members) {
+      let arr = map.get(am.member);
+      if (!arr) {
+        arr = [];
+        map.set(am.member, arr);
+      }
+      arr.push(am.roleid);
+    }
+    _membersByMemberIdCache.set(introspection, map);
+  }
+  return map;
+}
+
+function _getExpandRolesCache(
+  introspection: Introspection,
+): Map<string, PgRoles[]> {
+  let cache = _expandRolesCache.get(introspection);
+  if (!cache) {
+    cache = new Map();
+    _expandRolesCache.set(introspection, cache);
+  }
+  return cache;
+}
+
 /**
  * Gets a role given an OID; throws an error if the role is not found.
  */
 function getRole(introspection: Introspection, oid: string): PgRoles {
   if (oid === "0") {
     return PUBLIC_ROLE;
   }
-  const role = introspection.roles.find((r) => r._id === oid);
+  const role = _getRoleByIdMap(introspection).get(oid);
   if (!role) {
     throw new Error(`Could not find role with identifier '${oid}'`);
   }
@@ -110,7 +176,7 @@ function getRoleByName(introspection: Introspection, name: string): PgRoles {
   if (name === "public") {
     return PUBLIC_ROLE;
   }
-  const role = introspection.roles.find((r) => r.rolname === name);
+  const role = _getRoleByNameMap(introspection).get(name);
   if (!role) {
     throw new Error(`Could not find role with name '${name}'`);
   }
@@ -464,33 +530,62 @@ export const Permission = {
 
 /**
  * Returns all the roles role has been granted (including PUBLIC),
- * respecting `NOINHERIT`
+ * respecting `NOINHERIT`.
+ *
+ * Uses Map-based indexes and per-introspection caching to avoid O(n)
+ * linear scans on every call — significant for schemas with many roles.
  */
 export function expandRoles(
   introspection: Introspection,
   roles: PgRoles[],
   includeNoInherit = false,
 ): PgRoles[] {
+  // For single-role calls (the common path), use per-introspection cache
+  if (roles.length === 1) {
+    const cacheKey = `${roles[0]._id}:${includeNoInherit ? 1 : 0}`;
+    const cache = _getExpandRolesCache(introspection);
+    const cached = cache.get(cacheKey);
+    if (cached) {
+      return cached;
+    }
+    const result = _expandRolesUncached(introspection, roles, includeNoInherit);
+    cache.set(cacheKey, result);
+    return result;
+  }
+  return _expandRolesUncached(introspection, roles, includeNoInherit);
+}
+
+function _expandRolesUncached(
+  introspection: Introspection,
+  roles: PgRoles[],
+  includeNoInherit: boolean,
+): PgRoles[] {
+  const visitedIds = new Set<string>(["0"]); // PUBLIC_ROLE._id
   const allRoles = [PUBLIC_ROLE];
+  const memberIndex = _getMembersByMemberIdMap(introspection);
+  const roleByIdMap = _getRoleByIdMap(introspection);
 
   const addRole = (member: PgRoles) => {
-    if (!allRoles.includes(member)) {
+    if (!visitedIds.has(member._id)) {
+      visitedIds.add(member._id);
       allRoles.push(member);
       if (includeNoInherit || member.rolinherit !== false) {
-        introspection.auth_members.forEach((am) => {
-          // auth_members - role `am.member` gains the privileges of
-          // `am.roleid`
-
-          if (am.member === member._id) {
-            const rol = getRole(introspection, am.roleid);
-            addRole(rol);
+        const grants = memberIndex.get(member._id);
+        if (grants) {
+          for (const roleid of grants) {
+            const rol = roleByIdMap.get(roleid);
+            if (rol) {
+              addRole(rol);
+            }
           }
-        });
+        }
       }
     }
   };
 
-  roles.forEach(addRole);
+  for (const r of roles) {
+    addRole(r);
+  }
 
   return allRoles;
 }