From fdbb7919641fbc52d0a4209a8679b4e97ad6f07b Mon Sep 17 00:00:00 2001 From: Dan Lynch Date: Thu, 30 Jul 2026 01:54:01 +0000 Subject: [PATCH] feat(transform): extension + role routing (schema-portable installs, symbol qualification, role renaming) --- .../__tests__/extension-router.test.ts | 97 ++++++ .../__tests__/extension-transform.test.ts | 165 +++++++++ packages/transform/__tests__/facts.test.ts | 44 +++ .../__tests__/role-transform.test.ts | 133 +++++++ packages/transform/src/extension-router.ts | 329 ++++++++++++++++++ packages/transform/src/extension-transform.ts | 232 ++++++++++++ packages/transform/src/facts.ts | 93 +++++ packages/transform/src/index.ts | 26 ++ packages/transform/src/role-router.ts | 69 ++++ packages/transform/src/role-transform.ts | 138 ++++++++ 10 files changed, 1326 insertions(+) create mode 100644 packages/transform/__tests__/extension-router.test.ts create mode 100644 packages/transform/__tests__/extension-transform.test.ts create mode 100644 packages/transform/__tests__/role-transform.test.ts create mode 100644 packages/transform/src/extension-router.ts create mode 100644 packages/transform/src/extension-transform.ts create mode 100644 packages/transform/src/role-router.ts create mode 100644 packages/transform/src/role-transform.ts diff --git a/packages/transform/__tests__/extension-router.test.ts b/packages/transform/__tests__/extension-router.test.ts new file mode 100644 index 00000000..73a1009c --- /dev/null +++ b/packages/transform/__tests__/extension-router.test.ts @@ -0,0 +1,97 @@ +import { + COMMON_EXTENSIONS, + ExtensionDefinition, + ExtensionRouter +} from '../src/extension-router'; + +describe('ExtensionRouter.resolveInstall', () => { + it('returns the target schema for a routed extension', () => { + const router = new ExtensionRouter({ pgcrypto: { to: 'extensions' } }); + expect(router.resolveInstall('pgcrypto')).toBe('extensions'); + }); + + it('returns null to strip the SCHEMA clause (repollute)', () => { + const router = new ExtensionRouter({ pgcrypto: { to: null } }); + expect(router.resolveInstall('pgcrypto')).toBeNull(); + }); + + it('returns undefined (leave unchanged) for an unrouted extension', () => { + const router = new ExtensionRouter({ pgcrypto: { to: 'extensions' } }); + expect(router.resolveInstall('pg_trgm')).toBeUndefined(); + expect(router.resolveInstall(undefined)).toBeUndefined(); + }); + + it('refuses to move an extension pinned to a fixed schema', () => { + const inventory: ExtensionDefinition[] = [ + { name: 'postgis_tiger_geocoder', fixedSchema: 'tiger', symbols: [] } + ]; + const router = new ExtensionRouter( + { postgis_tiger_geocoder: { to: 'extensions' } }, + { inventory } + ); + expect(router.resolveInstall('postgis_tiger_geocoder')).toBeUndefined(); + }); +}); + +describe('ExtensionRouter.resolveSymbol', () => { + it('routes a bare extension symbol to the target schema', () => { + const router = ExtensionRouter.toSchema('extensions'); + expect(router.resolveSymbol(null, 'crypt', 'function')).toEqual({ to: 'extensions' }); + expect(router.resolveSymbol(null, 'gen_salt', 'function')).toEqual({ to: 'extensions' }); + }); + + it('requalifies a public-qualified symbol', () => { + const router = ExtensionRouter.toSchema('extensions'); + expect(router.resolveSymbol('public', 'digest', 'function')).toEqual({ to: 'extensions' }); + }); + + it('strips qualification when routing to bare (null)', () => { + const router = ExtensionRouter.toSchema(null, { from: ['extensions'] }); + expect(router.resolveSymbol('extensions', 'crypt', 'function')).toEqual({ to: null }); + }); + + it('leaves symbols already in the target schema unchanged', () => { + const router = ExtensionRouter.toSchema('extensions'); + expect(router.resolveSymbol('extensions', 'crypt', 'function')).toBeUndefined(); + }); + + it('never routes a symbol that graduated into core at the target version', () => { + const pg13 = ExtensionRouter.toSchema('extensions', { serverVersion: 13 }); + expect(pg13.resolveSymbol(null, 'gen_random_uuid', 'function')).toBeUndefined(); + + const pg12 = ExtensionRouter.toSchema('extensions', { serverVersion: 12 }); + expect(pg12.resolveSymbol(null, 'gen_random_uuid', 'function')).toEqual({ to: 'extensions' }); + }); + + it('routes extension-provided types (citext)', () => { + const router = ExtensionRouter.toSchema('extensions'); + expect(router.resolveSymbol(null, 'citext', 'type')).toEqual({ to: 'extensions' }); + // citext is a type, not a function — namespace must match + expect(router.resolveSymbol(null, 'citext', 'function')).toBeUndefined(); + }); + + it('ignores symbols not in the inventory (user-defined lookalikes)', () => { + const router = ExtensionRouter.toSchema('extensions'); + expect(router.resolveSymbol(null, 'my_helper', 'function')).toBeUndefined(); + expect(router.resolveSymbol('app', 'crypt', 'function')).toBeUndefined(); + }); + + it('honors an explicit `from` allowlist', () => { + const router = new ExtensionRouter({ pgcrypto: { to: 'extensions', from: ['public'] } }); + // only public-qualified refs are rewritten; bare refs are left alone + expect(router.resolveSymbol('public', 'crypt', 'function')).toEqual({ to: 'extensions' }); + expect(router.resolveSymbol(null, 'crypt', 'function')).toBeUndefined(); + }); +}); + +describe('ExtensionRouter inventory', () => { + it('exposes a curated set of common extensions', () => { + const names = COMMON_EXTENSIONS.map(d => d.name); + expect(names).toEqual(expect.arrayContaining(['pgcrypto', 'uuid-ossp', 'citext'])); + }); + + it('reports whether configured extensions have symbol routes', () => { + expect(ExtensionRouter.toSchema('extensions', { extensions: ['pgcrypto'] }).hasSymbolRoutes()).toBe(true); + expect(new ExtensionRouter({ nonexistent_ext: { to: 'x' } }).hasSymbolRoutes()).toBe(false); + }); +}); diff --git a/packages/transform/__tests__/extension-transform.test.ts b/packages/transform/__tests__/extension-transform.test.ts new file mode 100644 index 00000000..7c70ac32 --- /dev/null +++ b/packages/transform/__tests__/extension-transform.test.ts @@ -0,0 +1,165 @@ +import { Deparser, loadModule, parseSql } from 'plpgsql-parser'; + +import { ExtensionRouter } from '../src/extension-router'; +import { transformExtensions } from '../src/extension-transform'; + +beforeAll(async () => { + await loadModule(); +}); + +/** parse -> deparse -> parse must be stable (structural round-trip). */ +function assertRoundTrip(sql: string): void { + const once = parseSql(sql); + const deparsed = once.stmts.map((s: any) => Deparser.deparse(s.stmt)).join(';\n') + ';'; + const twice = parseSql(deparsed); + expect(twice.stmts.length).toBe(once.stmts.length); +} + +describe('transformExtensions — install schema (node construction)', () => { + it('adds a SCHEMA clause to CREATE EXTENSION that has none', () => { + const { sql, result } = transformExtensions( + 'CREATE EXTENSION pgcrypto;', + { pgcrypto: { to: 'extensions' } } + ); + expect(sql).toMatch(/CREATE EXTENSION pgcrypto\s+(WITH\s+)?SCHEMA extensions/i); + expect(result.installsMoved.get('pgcrypto')).toBe('extensions'); + assertRoundTrip(sql); + }); + + it('changes an existing SCHEMA clause', () => { + const { sql } = transformExtensions( + 'CREATE EXTENSION pgcrypto WITH SCHEMA public;', + { pgcrypto: { to: 'extensions' } } + ); + expect(sql).toMatch(/SCHEMA extensions/i); + expect(sql).not.toMatch(/SCHEMA public/i); + assertRoundTrip(sql); + }); + + it('removes the SCHEMA clause when routing to bare (null)', () => { + const { sql } = transformExtensions( + 'CREATE EXTENSION pgcrypto WITH SCHEMA extensions;', + { pgcrypto: { to: null } } + ); + expect(sql).not.toMatch(/SCHEMA/i); + expect(sql).toMatch(/CREATE EXTENSION pgcrypto/i); + assertRoundTrip(sql); + }); + + it('preserves IF NOT EXISTS when adding a schema', () => { + const { sql } = transformExtensions( + 'CREATE EXTENSION IF NOT EXISTS pgcrypto;', + { pgcrypto: { to: 'extensions' } } + ); + expect(sql).toMatch(/IF NOT EXISTS/i); + expect(sql).toMatch(/SCHEMA extensions/i); + assertRoundTrip(sql); + }); + + it('rewrites ALTER EXTENSION ... SET SCHEMA', () => { + const { sql, result } = transformExtensions( + 'ALTER EXTENSION pgcrypto SET SCHEMA public;', + { pgcrypto: { to: 'extensions' } } + ); + expect(sql).toMatch(/ALTER EXTENSION pgcrypto SET SCHEMA extensions/i); + expect(result.installsMoved.get('pgcrypto')).toBe('extensions'); + assertRoundTrip(sql); + }); + + it('leaves unrouted extensions untouched', () => { + const { sql, result } = transformExtensions( + 'CREATE EXTENSION pg_trgm;', + { pgcrypto: { to: 'extensions' } } + ); + expect(sql).toMatch(/CREATE EXTENSION pg_trgm/i); + expect(sql).not.toMatch(/SCHEMA/i); + expect(result.installsMoved.size).toBe(0); + }); +}); + +describe('transformExtensions — symbol references (node construction)', () => { + it('qualifies a bare extension function call', () => { + const { sql, result } = transformExtensions( + "SELECT crypt('pw', gen_salt('bf'));", + ExtensionRouter.toSchema('extensions') + ); + expect(sql).toMatch(/extensions\.crypt/); + expect(sql).toMatch(/extensions\.gen_salt/); + expect(result.symbolsRewritten.get('crypt')).toBe(1); + assertRoundTrip(sql); + }); + + it('requalifies a public-qualified extension call', () => { + const { sql } = transformExtensions( + "SELECT public.digest('x', 'sha256');", + ExtensionRouter.toSchema('extensions') + ); + expect(sql).toMatch(/extensions\.digest/); + expect(sql).not.toMatch(/public\.digest/); + assertRoundTrip(sql); + }); + + it('strips qualification when routing to bare', () => { + const { sql } = transformExtensions( + "SELECT extensions.crypt('pw', extensions.gen_salt('bf'));", + ExtensionRouter.toSchema(null, { from: ['extensions'] }) + ); + expect(sql).toMatch(/\bcrypt\(/); + expect(sql).not.toMatch(/extensions\.crypt/); + assertRoundTrip(sql); + }); + + it('does not touch gen_random_uuid on modern PostgreSQL (core symbol)', () => { + const { sql, result } = transformExtensions( + 'SELECT gen_random_uuid();', + ExtensionRouter.toSchema('extensions', { serverVersion: 16 }) + ); + expect(sql).not.toMatch(/extensions\.gen_random_uuid/); + expect(result.symbolsRewritten.has('gen_random_uuid')).toBe(false); + }); + + it('routes an extension-provided type (citext) in a column definition', () => { + const { sql } = transformExtensions( + 'CREATE TABLE t (email citext NOT NULL);', + ExtensionRouter.toSchema('extensions') + ); + expect(sql).toMatch(/extensions\.citext/); + assertRoundTrip(sql); + }); + + it('rewrites extension calls inside a LANGUAGE sql body', () => { + const { sql } = transformExtensions( + `CREATE FUNCTION hash_pw(pw text) RETURNS text AS $$ + SELECT crypt(pw, gen_salt('bf')) + $$ LANGUAGE sql;`, + ExtensionRouter.toSchema('extensions') + ); + expect(sql).toMatch(/extensions\.crypt/); + expect(sql).toMatch(/extensions\.gen_salt/); + assertRoundTrip(sql); + }); + + it('rewrites extension calls inside a PL/pgSQL body', () => { + const { sql } = transformExtensions( + `CREATE FUNCTION hash_pw(pw text) RETURNS text AS $$ + BEGIN + RETURN crypt(pw, gen_salt('bf')); + END; + $$ LANGUAGE plpgsql;`, + ExtensionRouter.toSchema('extensions') + ); + expect(sql).toMatch(/extensions\.crypt/); + expect(sql).toMatch(/extensions\.gen_salt/); + assertRoundTrip(sql); + }); + + it('leaves user-defined lookalikes alone', () => { + const { sql, result } = transformExtensions( + 'SELECT app.crypt(x), my_helper();', + ExtensionRouter.toSchema('extensions') + ); + expect(sql).toMatch(/app\.crypt/); + expect(sql).toMatch(/my_helper\(\)/); + expect(result.symbolsRewritten.size).toBe(0); + }); +}); diff --git a/packages/transform/__tests__/facts.test.ts b/packages/transform/__tests__/facts.test.ts index 6901da3b..c7a97835 100644 --- a/packages/transform/__tests__/facts.test.ts +++ b/packages/transform/__tests__/facts.test.ts @@ -167,4 +167,48 @@ describe('classifyStatements', () => { expect(facts[0].securityRelevant).toBe(true); expect(facts[0].roles).toEqual(['administrator']); }); + + it('classifies CREATE EXTENSION with and without a schema clause', () => { + const facts = classifyStatements(` + CREATE EXTENSION pgcrypto; + CREATE EXTENSION IF NOT EXISTS pg_partman WITH SCHEMA partman; + `); + expect(facts.map(f => f.kind)).toEqual(['extension', 'extension']); + expect(facts[0].extension).toEqual({ + name: 'pgcrypto', + schema: null, + action: 'create', + ifNotExists: false + }); + expect(facts[1].extension).toEqual({ + name: 'pg_partman', + schema: 'partman', + action: 'create', + ifNotExists: true + }); + }); + + it('classifies ALTER EXTENSION ... SET SCHEMA and DROP EXTENSION', () => { + const facts = classifyStatements(` + ALTER EXTENSION pg_partman SET SCHEMA public; + DROP EXTENSION IF EXISTS pgcrypto CASCADE; + `); + expect(facts.map(f => f.kind)).toEqual(['extension', 'extension']); + expect(facts[0].extension).toEqual({ + name: 'pg_partman', + schema: 'public', + action: 'set_schema' + }); + expect(facts[1].extension).toEqual({ + name: 'pgcrypto', + schema: null, + action: 'drop' + }); + }); + + it('does not classify non-extension ALTER ... SET SCHEMA as extension', () => { + const facts = classifyStatements(`ALTER TABLE app.t SET SCHEMA app2;`); + expect(facts[0].kind).not.toBe('extension'); + expect(facts[0].extension).toBeUndefined(); + }); }); diff --git a/packages/transform/__tests__/role-transform.test.ts b/packages/transform/__tests__/role-transform.test.ts new file mode 100644 index 00000000..6f17496d --- /dev/null +++ b/packages/transform/__tests__/role-transform.test.ts @@ -0,0 +1,133 @@ +import { Deparser, loadModule, parseSql } from 'plpgsql-parser'; + +import { RoleRouter } from '../src/role-router'; +import { transformRoles } from '../src/role-transform'; + +beforeAll(async () => { + await loadModule(); +}); + +function assertRoundTrip(sql: string): void { + const once = parseSql(sql); + const deparsed = once.stmts.map((s: any) => Deparser.deparse(s.stmt)).join(';\n') + ';'; + const twice = parseSql(deparsed); + expect(twice.stmts.length).toBe(once.stmts.length); +} + +// A generic, platform-neutral mapping between two role naming conventions. +const MAP = { anonymous: 'anon', administrator: 'service_role' }; + +describe('RoleRouter', () => { + it('resolves routed names and leaves others', () => { + const router = new RoleRouter(MAP); + expect(router.resolve('anonymous')).toBe('anon'); + expect(router.resolve('authenticated')).toBeUndefined(); + expect(router.resolve(undefined)).toBeUndefined(); + }); + + it('inverts a one-to-one mapping', () => { + const inv = new RoleRouter(MAP).invert(); + expect(inv.resolve('anon')).toBe('anonymous'); + expect(inv.resolve('service_role')).toBe('administrator'); + }); + + it('refuses to invert a non-injective mapping', () => { + expect(() => new RoleRouter({ a: 'x', b: 'x' }).invert()).toThrow(/one-to-one/); + }); +}); + +describe('transformRoles', () => { + it('renames grantees in GRANT and REVOKE', () => { + const { sql, result } = transformRoles( + 'GRANT SELECT ON t TO anonymous, authenticated;', + MAP + ); + expect(sql).toMatch(/TO anon, authenticated/); + expect(result.rolesRenamed.get('anon')).toBe(1); + assertRoundTrip(sql); + }); + + it('renames ownership targets', () => { + const { sql } = transformRoles('ALTER TABLE t OWNER TO administrator;', MAP); + expect(sql).toMatch(/OWNER TO service_role/); + assertRoundTrip(sql); + }); + + it('renames policy roles', () => { + const { sql } = transformRoles( + 'CREATE POLICY p ON t FOR SELECT TO anonymous USING (true);', + MAP + ); + expect(sql).toMatch(/TO anon/); + assertRoundTrip(sql); + }); + + it('renames ALTER DEFAULT PRIVILEGES FOR ROLE and grantees', () => { + const { sql } = transformRoles( + 'ALTER DEFAULT PRIVILEGES FOR ROLE administrator IN SCHEMA s GRANT SELECT ON TABLES TO anonymous;', + MAP + ); + expect(sql).toMatch(/FOR ROLE service_role/); + expect(sql).toMatch(/TO anon/); + assertRoundTrip(sql); + }); + + it('renames CREATE ROLE and DROP ROLE', () => { + const create = transformRoles('CREATE ROLE anonymous NOLOGIN;', MAP); + expect(create.sql).toMatch(/CREATE ROLE anon/); + assertRoundTrip(create.sql); + + const drop = transformRoles('DROP ROLE anonymous;', MAP); + expect(drop.sql).toMatch(/DROP ROLE anon/); + assertRoundTrip(drop.sql); + }); + + it('renames role membership (granted and grantee roles)', () => { + const { sql, result } = transformRoles('GRANT administrator TO anonymous;', MAP); + // The deparser folds the granted-role token to upper case (a pre-existing + // quirk of AccessPriv rendering, independent of the rename); match + // case-insensitively. The identifier is still `service_role`. + expect(sql).toMatch(/GRANT service_role TO anon/i); + expect(result.rolesRenamed.get('service_role')).toBe(1); + expect(result.rolesRenamed.get('anon')).toBe(1); + assertRoundTrip(sql); + }); + + it('renames SET ROLE and SET SESSION AUTHORIZATION', () => { + const setRole = transformRoles('SET ROLE anonymous;', MAP); + expect(setRole.sql).toMatch(/SET ROLE (TO )?anon|SET ROLE anon/i); + assertRoundTrip(setRole.sql); + + const setAuth = transformRoles('SET SESSION AUTHORIZATION anonymous;', MAP); + expect(setAuth.sql).toMatch(/anon/); + assertRoundTrip(setAuth.sql); + }); + + it('renames ALTER ROLE ... RENAME TO', () => { + const { sql } = transformRoles('ALTER ROLE anonymous RENAME TO administrator;', MAP); + expect(sql).toMatch(/ALTER ROLE anon RENAME TO service_role/); + assertRoundTrip(sql); + }); + + it('never rewrites PUBLIC or CURRENT_USER', () => { + const { sql, result } = transformRoles( + 'GRANT SELECT ON t TO PUBLIC; ALTER TABLE t OWNER TO CURRENT_USER;', + { PUBLIC: 'anon', CURRENT_USER: 'anon' } + ); + expect(sql).toMatch(/TO PUBLIC/); + expect(sql).toMatch(/CURRENT_USER/); + expect(result.rolesRenamed.size).toBe(0); + }); + + it('leaves unrouted roles untouched', () => { + const { sql, result } = transformRoles('GRANT SELECT ON t TO authenticated;', MAP); + expect(sql).toMatch(/TO authenticated/); + expect(result.rolesRenamed.size).toBe(0); + }); + + it('is fully reversible via invert()', () => { + const forward = transformRoles('GRANT SELECT ON t TO anonymous;', MAP); + const back = transformRoles(forward.sql, new RoleRouter(MAP).invert()); + expect(back.sql).toMatch(/TO anonymous/); + }); +}); diff --git a/packages/transform/src/extension-router.ts b/packages/transform/src/extension-router.ts new file mode 100644 index 00000000..e81dbeb8 --- /dev/null +++ b/packages/transform/src/extension-router.ts @@ -0,0 +1,329 @@ +/** + * Extension routing for the core transform. + * + * A PostgreSQL extension is installed into exactly one schema, and the objects + * it provides (functions, types, operators) live in that schema. Two databases + * can install the *same* extension into *different* schemas — one leaves it in + * `public`, another isolates it in a dedicated schema — so portable SQL must be + * able to move an extension's install site **and** rewrite every reference to + * the symbols it provides. + * + * This is a distinct problem from {@link SchemaRouter}, which routes objects the + * SQL itself creates. Here the objects are owned by an extension and are not + * declared in the script; the only way to know that a bare `crypt(...)` call + * belongs to `pgcrypto` (and must move when `pgcrypto` moves) is an + * **inventory** of which symbols each extension provides. That inventory is + * also **version-aware**: a symbol can graduate into the server core and cease + * to be extension-owned (the canonical case is `gen_random_uuid()`, provided by + * `pgcrypto` before PostgreSQL 13 and part of core `pg_catalog` from 13 on — it + * must never be routed on 13+). + * + * An {@link ExtensionRouter} answers two questions: + * + * - **install**: where should `CREATE EXTENSION ` (or `ALTER EXTENSION + * SET SCHEMA`) place the extension? + * - **reference**: given a reference to a symbol `(schema, name)` in namespace + * `ns`, what schema should it live in now — including stripping the qualifier + * entirely (relying on `search_path`) when the target is "unqualified". + * + * A schema of `null` denotes the **unqualified** site: a bare reference + * (`crypt(...)`) or installing an extension with no explicit `SCHEMA` clause. + * This makes the router fully bidirectional: `null -> 'extensions'` qualifies + * bare references, and `'extensions' -> null` strips them back to bare. + */ + +/** PostgreSQL namespaces an extension symbol can occupy. */ +export type ExtensionSymbolNamespace = 'function' | 'type' | 'operator'; + +/** A single symbol provided by an extension. */ +export interface ExtensionSymbol { + /** Bare symbol name (e.g. `crypt`, `citext`, `gen_random_uuid`). */ + name: string; + /** Which namespace the symbol occupies. Defaults to `function`. */ + namespace?: ExtensionSymbolNamespace; + /** + * Major PostgreSQL version at/after which this symbol is part of the server + * core (`pg_catalog`) rather than the extension. At or above this version + * the symbol is **never routed** — it resolves from `pg_catalog` regardless + * of where the extension lives. The classic case is `gen_random_uuid` + * (`coreSince: 13`). + */ + coreSince?: number; +} + +/** The set of symbols an extension provides. */ +export interface ExtensionDefinition { + /** Extension name as used in `CREATE EXTENSION `. */ + name: string; + /** + * Whether the extension is relocatable (`ALTER EXTENSION ... SET SCHEMA` + * succeeds). Informational for callers deciding *how* to move an existing + * install; the router never emits `SET SCHEMA` itself. + */ + relocatable?: boolean; + /** + * A schema the extension is pinned to and cannot be moved out of (fixed by + * its control file, e.g. a geocoder that always creates its own schema). + * When set, the router refuses to route the extension's install or symbols. + */ + fixedSchema?: string; + /** Symbols the extension provides. */ + symbols: ExtensionSymbol[]; +} + +/** + * A routing rule for one extension. `to` is the destination install/reference + * schema (or `null` to strip qualification / install without a `SCHEMA` + * clause). `from` lists the schemas currently holding the extension that + * should be rewritten; a `null` entry additionally routes bare references. + * When `from` is omitted, references in **any** schema (and bare ones) are + * routed — useful when the source layout is unknown. + */ +export interface ExtensionRoute { + to: string | null; + from?: (string | null)[]; +} + +/** Routing specification keyed by extension name. */ +export type ExtensionRouteSpec = Record; + +/** A concrete rewrite instruction: from a schema (or bare) to a schema (or bare). */ +export interface SymbolRewrite { + /** Target schema, or `null` to make the reference unqualified. */ + to: string | null; +} + +export interface ExtensionRouterOptions { + /** + * The symbol inventory. Defaults to {@link COMMON_EXTENSIONS}. Callers can + * pass an augmented or replacement inventory (e.g. extended from live-catalog + * introspection). + */ + inventory?: ExtensionDefinition[]; + /** + * Target PostgreSQL major version, used to apply {@link ExtensionSymbol.coreSince}. + * Defaults to a high value so version-graduated symbols (e.g. `gen_random_uuid`) + * are treated as core and never routed unless a caller opts into an older + * version. + */ + serverVersion?: number; +} + +const DEFAULT_SERVER_VERSION = 9999; + +/** + * A curated, platform-agnostic inventory of common contrib/extension symbols. + * Intentionally minimal and conservative — only well-known, stable symbols are + * listed so routing never touches something it cannot prove belongs to the + * extension. Callers extend this via {@link ExtensionRouterOptions.inventory}. + */ +export const COMMON_EXTENSIONS: ExtensionDefinition[] = [ + { + name: 'pgcrypto', + relocatable: true, + symbols: [ + { name: 'crypt' }, + { name: 'gen_salt' }, + { name: 'digest' }, + { name: 'hmac' }, + { name: 'encrypt' }, + { name: 'decrypt' }, + { name: 'encrypt_iv' }, + { name: 'decrypt_iv' }, + { name: 'gen_random_bytes' }, + // Provided by pgcrypto before PG13; part of core pg_catalog from 13 on. + { name: 'gen_random_uuid', coreSince: 13 }, + { name: 'pgp_sym_encrypt' }, + { name: 'pgp_sym_decrypt' }, + { name: 'pgp_pub_encrypt' }, + { name: 'pgp_pub_decrypt' } + ] + }, + { + name: 'uuid-ossp', + relocatable: true, + symbols: [ + { name: 'uuid_generate_v1' }, + { name: 'uuid_generate_v1mc' }, + { name: 'uuid_generate_v3' }, + { name: 'uuid_generate_v4' }, + { name: 'uuid_generate_v5' }, + { name: 'uuid_nil' } + ] + }, + { + name: 'citext', + relocatable: true, + symbols: [ + { name: 'citext', namespace: 'type' }, + { name: 'citextin', namespace: 'function' }, + { name: 'citext_hash', namespace: 'function' } + ] + }, + { + name: 'pg_trgm', + relocatable: true, + symbols: [ + { name: 'similarity' }, + { name: 'show_trgm' }, + { name: 'word_similarity' }, + { name: 'strict_word_similarity' } + ] + }, + { + name: 'ltree', + relocatable: true, + symbols: [ + { name: 'ltree', namespace: 'type' }, + { name: 'lquery', namespace: 'type' }, + { name: 'ltxtquery', namespace: 'type' }, + { name: 'subltree' }, + { name: 'subpath' }, + { name: 'nlevel' }, + { name: 'lca' } + ] + }, + { + name: 'hstore', + relocatable: true, + symbols: [{ name: 'hstore', namespace: 'type' }] + }, + { + name: 'unaccent', + relocatable: true, + symbols: [{ name: 'unaccent' }] + } +]; + +interface IndexedSymbol { + extension: string; + coreSince?: number; +} + +/** + * Resolves install-schema and symbol-reference routing for extensions, backed + * by a version-aware symbol inventory. + */ +export class ExtensionRouter { + private readonly routes: Map; + private readonly serverVersion: number; + private readonly definitions: Map; + /** `${namespace}:${name}` -> providing extension (+ version predicate). */ + private readonly symbolIndex: Map; + + constructor(routes: ExtensionRouteSpec | Map = {}, options: ExtensionRouterOptions = {}) { + this.routes = routes instanceof Map ? new Map(routes) : new Map(Object.entries(routes)); + this.serverVersion = options.serverVersion ?? DEFAULT_SERVER_VERSION; + const inventory = options.inventory ?? COMMON_EXTENSIONS; + this.definitions = new Map(inventory.map(def => [def.name, def])); + this.symbolIndex = new Map(); + for (const def of inventory) { + for (const sym of def.symbols) { + const key = symbolKey(sym.namespace ?? 'function', sym.name); + // First definition wins; inventories should not double-claim a symbol. + if (!this.symbolIndex.has(key)) { + this.symbolIndex.set(key, { extension: def.name, coreSince: sym.coreSince }); + } + } + } + } + + /** + * Build a router that moves the given extensions (default: every inventoried + * one) to a single `targetSchema`, rewriting their references from `public` + * and from bare (unqualified) sites. Passing `null` as `targetSchema` strips + * qualification instead (the "repollute public / rely on search_path" + * direction). + */ + static toSchema( + targetSchema: string | null, + options: ExtensionRouterOptions & { extensions?: string[]; from?: (string | null)[] } = {} + ): ExtensionRouter { + const inventory = options.inventory ?? COMMON_EXTENSIONS; + const names = options.extensions ?? inventory.map(d => d.name); + const from = options.from ?? ['public', null]; + const spec: ExtensionRouteSpec = {}; + for (const name of names) spec[name] = { to: targetSchema, from }; + return new ExtensionRouter(spec, options); + } + + /** Coerce a spec, map, or existing router into a router. */ + static from( + source: ExtensionRouter | ExtensionRouteSpec | Map, + options?: ExtensionRouterOptions + ): ExtensionRouter { + if (source instanceof ExtensionRouter) return source; + return new ExtensionRouter(source, options); + } + + /** True when no extension routes are configured. */ + get size(): number { + return this.routes.size; + } + + /** Extension names this router may rewrite. */ + routedExtensions(): string[] { + return [...this.routes.keys()]; + } + + /** True when any configured extension has inventoried symbols to rewrite. */ + hasSymbolRoutes(): boolean { + for (const name of this.routes.keys()) { + const def = this.definitions.get(name); + if (def && !def.fixedSchema && def.symbols.length > 0) return true; + } + return false; + } + + /** + * Resolve the install schema for `CREATE EXTENSION ` / + * `ALTER EXTENSION SET SCHEMA`. Returns: + * - a string to place/keep the extension in that schema, + * - `null` to install with no explicit `SCHEMA` clause (server default), + * - `undefined` to leave the statement unchanged (no route, or the extension + * is pinned to a fixed schema). + */ + resolveInstall(extname: string | undefined | null): string | null | undefined { + if (!extname) return undefined; + const route = this.routes.get(extname); + if (!route) return undefined; + if (this.definitions.get(extname)?.fixedSchema !== undefined) return undefined; + return route.to; + } + + /** + * Resolve a rewrite for a reference to `(schema, name)` in namespace `ns`, + * where `schema` is `null` for a bare reference. Returns the target + * ({@link SymbolRewrite}) or `undefined` to leave the reference unchanged. + * + * A reference is rewritten only when the inventory proves the symbol belongs + * to a routed extension, the symbol has not graduated to core at the target + * server version, the current schema is one the route rewrites `from`, and + * the destination actually differs from where it already is. + */ + resolveSymbol( + schema: string | null | undefined, + name: string | undefined, + ns: ExtensionSymbolNamespace + ): SymbolRewrite | undefined { + if (!name) return undefined; + const indexed = this.symbolIndex.get(symbolKey(ns, name)); + if (!indexed) return undefined; + // Graduated into core at/after the target server version — resolves from + // pg_catalog and must not be tied to the extension's schema. + if (indexed.coreSince !== undefined && this.serverVersion >= indexed.coreSince) return undefined; + + const route = this.routes.get(indexed.extension); + if (!route) return undefined; + if (this.definitions.get(indexed.extension)?.fixedSchema !== undefined) return undefined; + + const current = schema ?? null; + if (route.from !== undefined && !route.from.includes(current)) return undefined; + if (current === route.to) return undefined; + return { to: route.to }; + } +} + +function symbolKey(ns: ExtensionSymbolNamespace, name: string): string { + return `${ns}:${name}`; +} diff --git a/packages/transform/src/extension-transform.ts b/packages/transform/src/extension-transform.ts new file mode 100644 index 00000000..a5b7c4db --- /dev/null +++ b/packages/transform/src/extension-transform.ts @@ -0,0 +1,232 @@ +/** + * AST transform that applies an {@link ExtensionRouter}. + * + * Unlike the schema transform — which only rewrites string *fields* already + * present on a node — moving an extension changes the **shape** of the AST: + * + * - `CREATE EXTENSION ` with no `SCHEMA` clause must gain one, which means + * constructing a new `DefElem` node and attaching it to a (possibly absent) + * `options` array; the reverse strips the `DefElem` (and empties `options`). + * - a bare `crypt(...)` reference gains a schema qualifier (a `String` node is + * prepended to `funcname`), and the reverse removes it. + * + * All rewrites are AST-precise and driven by the router's version-aware symbol + * inventory, so built-ins that merely share a name with an extension symbol + * (and symbols that have graduated into core) are left untouched. + */ + +import { walk as walkSql } from '@pgsql/traverse'; +import { Deparser, parseSql, transformSync, walk as walkPlpgsql } from 'plpgsql-parser'; + +import type { ExtensionRouteSpec, ExtensionRouterOptions, ExtensionSymbolNamespace } from './extension-router'; +import { ExtensionRouter } from './extension-router'; + +/** What an extension transform changed. */ +export interface ExtensionTransformResult { + /** Extensions whose install schema was rewritten: extname -> target (null = unqualified). */ + installsMoved: Map; + /** Symbol references rewritten, counted by bare symbol name. */ + symbolsRewritten: Map; +} + +/** Create a fresh result tracker. */ +export function createExtensionResult(): ExtensionTransformResult { + return { installsMoved: new Map(), symbolsRewritten: new Map() }; +} + +function schemaStringNode(schema: string): any { + return { String: { sval: schema } }; +} + +/** + * Rewrite a (possibly schema-qualified) name list in place against the router. + * `names` is the node's own array (`funcname`, `TypeName.names`, ...); it is + * mutated in place so the enclosing node keeps referring to it. + */ +function rewriteNameList( + names: any[] | undefined, + router: ExtensionRouter, + ns: ExtensionSymbolNamespace, + result: ExtensionTransformResult +): void { + if (!Array.isArray(names) || names.length === 0) return; + const parts = names.map((n: any) => n?.String?.sval); + if (parts.some(p => typeof p !== 'string')) return; // non-identifier element (e.g. %type) + + const name = parts[parts.length - 1]; + const schema = parts.length >= 2 ? parts[parts.length - 2] : null; + const rewrite = router.resolveSymbol(schema, name, ns); + if (!rewrite) return; + + if (rewrite.to === null) { + // Strip qualification: keep only the bare name node. + names.splice(0, names.length - 1); + } else if (schema === null) { + // Qualify a bare reference: prepend the target schema. + names.unshift(schemaStringNode(rewrite.to)); + } else { + // Requalify: rewrite the existing schema element. + names[names.length - 2].String.sval = rewrite.to; + } + result.symbolsRewritten.set(name, (result.symbolsRewritten.get(name) ?? 0) + 1); +} + +/** Read/replace the `schema` DefElem inside a `CreateExtensionStmt.options`. */ +function setExtensionSchemaOption(node: any, target: string | null): boolean { + const options: any[] = Array.isArray(node.options) ? node.options : []; + const idx = options.findIndex((opt: any) => opt?.DefElem?.defname === 'schema'); + + if (target === null) { + // Install with no explicit SCHEMA clause: remove any existing option. + if (idx === -1) return false; + options.splice(idx, 1); + if (options.length === 0) delete node.options; + else node.options = options; + return true; + } + + const defElem = { + DefElem: { + defname: 'schema', + arg: { String: { sval: target } }, + defaction: 'DEFELEM_UNSPEC' + } + }; + if (idx === -1) { + options.push(defElem); + } else { + const existing = options[idx].DefElem; + if (existing.arg?.String?.sval === target) return false; + existing.arg = { String: { sval: target } }; + } + node.options = options; + return true; +} + +/** + * Create a SQL AST visitor that applies extension routing. Composable with the + * walkers used by the core transform and by PL/pgSQL body traversal. + */ +export function createExtensionVisitor( + router: ExtensionRouter, + result: ExtensionTransformResult +) { + const rewriteBody = router.hasSymbolRoutes(); + + return { + CreateExtensionStmt: (path: any) => { + const node = path.node; + const target = router.resolveInstall(node.extname); + if (target === undefined) return; + if (setExtensionSchemaOption(node, target)) { + result.installsMoved.set(node.extname, target); + } + }, + + // ALTER EXTENSION SET SCHEMA . The walker does not recurse + // into AlterObjectSchemaStmt.object; a null target has no SET SCHEMA form + // (a schema must be named), so it is left unchanged. + AlterObjectSchemaStmt: (path: any) => { + const node = path.node; + if (node.objectType !== 'OBJECT_EXTENSION') return; + const extname = node.object?.String?.sval; + const target = router.resolveInstall(extname); + if (typeof target !== 'string' || target === node.newschema) return; + node.newschema = target; + result.installsMoved.set(extname, target); + }, + + FuncCall: (path: any) => { + rewriteNameList(path.node.funcname, router, 'function', result); + }, + + // The walker does not auto-recurse into CallStmt.funccall. + CallStmt: (path: any) => { + if (path.node.funccall) { + rewriteNameList(path.node.funccall.funcname, router, 'function', result); + } + }, + + TypeName: (path: any) => { + rewriteNameList(path.node.names, router, 'type', result); + }, + + CreateFunctionStmt: (path: any) => { + const node = path.node; + if (!rewriteBody) return; + // LANGUAGE sql bodies are opaque strings to the walker; parse and rewrite + // them separately (PL/pgSQL bodies ride the hydrated walk in transformExtensions). + const isPlpgsql = (node.options ?? []).some( + (opt: any) => + opt?.DefElem?.defname === 'language' && + opt.DefElem.arg?.String?.sval === 'plpgsql' + ); + if (isPlpgsql || !Array.isArray(node.options)) return; + for (const opt of node.options) { + if (opt?.DefElem?.defname === 'as' && opt.DefElem.arg?.List?.items) { + for (const item of opt.DefElem.arg.List.items) { + if (typeof item?.String?.sval === 'string') { + try { + item.String.sval = rewriteSqlBodyString(item.String.sval, router, result); + } catch { + // Not parseable standalone (e.g. C symbol names) — leave as-is. + } + } + } + } + } + } + }; +} + +/** Rewrite extension references inside a `LANGUAGE sql` body string. */ +function rewriteSqlBodyString( + body: string, + router: ExtensionRouter, + result: ExtensionTransformResult +): string { + const stmts: any[] = parseSql(body)?.stmts ?? []; + if (stmts.length === 0) return body; + const visitor = createExtensionVisitor(router, result); + const pieces: string[] = []; + for (const stmt of stmts) { + if (!stmt?.stmt) continue; + walkSql(stmt.stmt, visitor); + pieces.push(Deparser.deparse(stmt.stmt)); + } + return pieces.join(';\n'); +} + +/** + * Apply extension routing to a SQL string: parse -> walk -> deparse, including + * PL/pgSQL and `LANGUAGE sql` bodies. Returns the rewritten SQL and a summary + * of what changed. + */ +export function transformExtensions( + sql: string, + router: ExtensionRouter | ExtensionRouteSpec, + options: ExtensionRouterOptions = {} +): { sql: string; result: ExtensionTransformResult } { + const resolved = ExtensionRouter.from(router, options); + const result = createExtensionResult(); + + const out = transformSync(sql, (ctx: any) => { + const stmts: any[] = ctx.sql?.stmts ?? []; + const visitor = createExtensionVisitor(resolved, result); + for (const stmt of stmts) { + if (stmt?.stmt) walkSql(stmt.stmt, visitor); + } + if (resolved.hasSymbolRoutes()) { + for (const fn of ctx.functions ?? []) { + if (fn.plpgsql?.hydrated) { + walkPlpgsql(fn.plpgsql.hydrated, {}, { + walkSqlExpressions: true, + sqlVisitor: visitor + }); + } + } + } + }, { hydrate: true, pretty: true }); + + return { sql: out, result }; +} diff --git a/packages/transform/src/facts.ts b/packages/transform/src/facts.ts index 9256974e..351a1f68 100644 --- a/packages/transform/src/facts.ts +++ b/packages/transform/src/facts.ts @@ -14,6 +14,7 @@ export interface QualifiedName { */ export type StatementKind = | 'schema' + | 'extension' | 'table' | 'view' | 'index' @@ -29,6 +30,39 @@ export type StatementKind = | 'seed_dml' | 'other'; +/** + * The action a statement performs on a PostgreSQL extension. + * + * - `create` — `CREATE EXTENSION` (`CreateExtensionStmt`). + * - `set_schema` — `ALTER EXTENSION ... SET SCHEMA` (`AlterObjectSchemaStmt` + * with `objectType: OBJECT_EXTENSION`); only succeeds for relocatable + * extensions, so it is surfaced as its own action. + * - `drop` — `DROP EXTENSION` (`DropStmt` with `removeType: OBJECT_EXTENSION`). + */ +export type ExtensionAction = 'create' | 'set_schema' | 'drop'; + +/** + * Facts about an extension-level statement. Unlike ordinary objects, an + * extension is installed into exactly one schema and its member objects are + * renamed with it, so the relevant fact is the extension name plus the schema + * it is (being) placed in. + */ +export interface ExtensionFact { + /** The extension name (`CREATE EXTENSION `). */ + name: string; + /** + * The schema the statement places the extension in, or `null` when none is + * specified (`CREATE EXTENSION ` with no `SCHEMA` clause installs into + * the current default — typically `public` or the extension's fixed schema). + * `DROP EXTENSION` carries no schema. + */ + schema: string | null; + /** Which extension operation this statement performs. */ + action: ExtensionAction; + /** `CREATE EXTENSION IF NOT EXISTS`. */ + ifNotExists?: boolean; +} + /** * AST-derived facts about a single top-level SQL statement. * @@ -62,6 +96,11 @@ export interface StatementFacts { referencedSchemas: string[]; /** Role names granted to, owning, or bound by this statement. */ roles: string[]; + /** + * For extension-level statements (`kind: 'extension'`): the extension being + * created, relocated, or dropped. Absent for every other statement. + */ + extension?: ExtensionFact; /** Foreign-key target tables (from column/table FK constraints). */ fkTargets: QualifiedName[]; /** @@ -94,6 +133,7 @@ const SECURITY_TAGS = new Set([ const KIND_BY_TAG: Record = { CreateSchemaStmt: 'schema', + CreateExtensionStmt: 'extension', CreateStmt: 'table', ViewStmt: 'view', IndexStmt: 'index', @@ -120,6 +160,23 @@ function qn(schema: string | null | undefined, name: string): QualifiedName { return { schema: schema ?? null, name }; } +/** + * Read the `SCHEMA ` clause of a `CreateExtensionStmt` from its options + * (`{ DefElem: { defname: 'schema', arg: { String: { sval } } } }`), or `null` + * when the statement specifies no schema. + */ +function extensionSchemaOption(options: any[] | undefined): string | null { + if (!Array.isArray(options)) return null; + for (const opt of options) { + const def = opt?.DefElem; + if (def?.defname === 'schema') { + const sval = def.arg?.String?.sval; + return typeof sval === 'string' ? sval : null; + } + } + return null; +} + function nameListToQualified(names: any[] | undefined): QualifiedName | null { if (!Array.isArray(names) || names.length === 0) return null; const parts = names @@ -222,6 +279,14 @@ function classifyOne(nodeTag: string, node: any): StatementFacts { case 'CreateSchemaStmt': facts.creates.push(qn(null, node.schemaname)); break; + case 'CreateExtensionStmt': + facts.extension = { + name: node.extname, + schema: extensionSchemaOption(node.options), + action: 'create', + ifNotExists: node.if_not_exists === true + }; + break; case 'CreateStmt': case 'ViewStmt': { const rel = nodeTag === 'ViewStmt' ? node.view : node.relation; @@ -284,6 +349,34 @@ function classifyOne(nodeTag: string, node: any): StatementFacts { if (name) facts.creates.push(name); break; } + case 'AlterObjectSchemaStmt': + // ALTER EXTENSION SET SCHEMA . Other object types + // (TABLE/TYPE/FUNCTION ... SET SCHEMA) keep their default classification. + if (node.objectType === 'OBJECT_EXTENSION') { + const name = node.object?.String?.sval; + if (typeof name === 'string') { + facts.kind = 'extension'; + facts.extension = { + name, + schema: typeof node.newschema === 'string' ? node.newschema : null, + action: 'set_schema' + }; + } + } + break; + case 'DropStmt': + // DROP EXTENSION [, ...]. Extension names are bare String nodes. + // Other DROP object types keep their default classification. When + // several extensions are dropped in one statement the first names the + // fact; the full list stays available via the raw node. + if (node.removeType === 'OBJECT_EXTENSION' && Array.isArray(node.objects)) { + const name = node.objects[0]?.String?.sval; + if (typeof name === 'string') { + facts.kind = 'extension'; + facts.extension = { name, schema: null, action: 'drop' }; + } + } + break; case 'AlterTableStmt': { if (node.relation) { facts.creates.push(qn(node.relation.schemaname ?? null, node.relation.relname)); diff --git a/packages/transform/src/index.ts b/packages/transform/src/index.ts index ced506c1..27550252 100644 --- a/packages/transform/src/index.ts +++ b/packages/transform/src/index.ts @@ -1,9 +1,35 @@ export type { + ExtensionAction, + ExtensionFact, QualifiedName, StatementFacts, StatementKind, } from './facts'; export { classifyStatements } from './facts'; +export type { + ExtensionDefinition, + ExtensionRoute, + ExtensionRouteSpec, + ExtensionRouterOptions, + ExtensionSymbol, + ExtensionSymbolNamespace, + SymbolRewrite, +} from './extension-router'; +export { COMMON_EXTENSIONS, ExtensionRouter } from './extension-router'; +export type { ExtensionTransformResult } from './extension-transform'; +export { + createExtensionResult, + createExtensionVisitor, + transformExtensions, +} from './extension-transform'; +export type { RoleRouteSpec } from './role-router'; +export { RoleRouter } from './role-router'; +export type { RoleTransformResult } from './role-transform'; +export { + createRoleResult, + createRoleVisitor, + transformRoles, +} from './role-transform'; export type { ObjectInventory, QualifyResult, diff --git a/packages/transform/src/role-router.ts b/packages/transform/src/role-router.ts new file mode 100644 index 00000000..aee1bfd5 --- /dev/null +++ b/packages/transform/src/role-router.ts @@ -0,0 +1,69 @@ +/** + * Role-name routing for the core transform. + * + * Two databases can express the *same* access model with *different* role + * names — one calls the unauthenticated role `anonymous`, another `anon`; one + * calls the privileged role `administrator`, another something else. Portable + * SQL must be able to translate role identifiers between these conventions. + * + * Scope, deliberately narrow: a {@link RoleRouter} only renames role + * *identifiers* in the AST positions where a role name legitimately appears + * (grants, ownership, policies, default privileges, role membership, role + * settings). It does **not** alter role *attributes* (`BYPASSRLS`, `LOGIN`, + * ...): two roles that share a name across conventions do not necessarily + * share privilege scope, and inventing attribute changes would silently + * misrepresent the security model. Privilege-semantics reconciliation is a + * downstream, policy-level concern; the router restricts itself to the + * deterministic, reversible part — the name. + * + * Renaming is a pure identifier substitution, so a router is trivially + * invertible ({@link RoleRouter.invert}) for bidirectional translation. + */ + +/** Role-rename specification: source role name -> target role name. */ +export type RoleRouteSpec = Record; + +/** Resolves a target role name for a source role name. */ +export class RoleRouter { + private readonly map: Map; + + constructor(routes: RoleRouteSpec | Map = {}) { + this.map = routes instanceof Map ? new Map(routes) : new Map(Object.entries(routes)); + } + + /** Coerce a spec, map, or existing router into a router. */ + static from(source: RoleRouter | RoleRouteSpec | Map): RoleRouter { + return source instanceof RoleRouter ? source : new RoleRouter(source); + } + + /** Number of configured renames. */ + get size(): number { + return this.map.size; + } + + /** + * Resolve the target name for `role`, or `undefined` when it is not routed + * or already at its destination. + */ + resolve(role: string | undefined | null): string | undefined { + if (!role) return undefined; + const target = this.map.get(role); + if (target === undefined || target === role) return undefined; + return target; + } + + /** + * The inverse router (target -> source), for translating in the opposite + * direction. Throws if the mapping is not one-to-one. + */ + invert(): RoleRouter { + const inverted = new Map(); + for (const [from, to] of this.map) { + if (inverted.has(to)) { + throw new Error(`RoleRouter.invert: mapping is not one-to-one (two sources map to "${to}")`); + } + inverted.set(to, from); + } + return new RoleRouter(inverted); + } +} diff --git a/packages/transform/src/role-transform.ts b/packages/transform/src/role-transform.ts new file mode 100644 index 00000000..ae67bf3f --- /dev/null +++ b/packages/transform/src/role-transform.ts @@ -0,0 +1,138 @@ +/** + * AST transform that applies a {@link RoleRouter}. + * + * Role identifiers surface in many statement forms. Most flow through a single + * `RoleSpec` node (grantees, ownership `... OWNER TO`, `CREATE POLICY ... TO`, + * `ALTER DEFAULT PRIVILEGES [FOR ROLE ...] ... TO`, `ALTER ROLE`, `DROP ROLE`, + * role-membership grantees, `GRANTED BY`, `REASSIGN OWNED BY`), which the + * traversal reaches automatically. Four positions carry a role name outside a + * `RoleSpec` and are handled explicitly: + * + * - `CREATE ROLE ` — a bare string on `CreateRoleStmt.role`. + * - `GRANT TO ...` — the *granted* roles are `AccessPriv.priv_name` + * strings on `GrantRoleStmt.granted_roles` (not `RoleSpec`s). + * - `SET ROLE` / `SET SESSION AUTHORIZATION` — the role is a string constant + * in `VariableSetStmt.args`. + * - `ALTER ROLE RENAME TO ` — `RenameStmt` with string + * `subname`/`newname`. + * + * Reserved role specifications (`PUBLIC`, `CURRENT_USER`, `SESSION_USER`, ...) + * carry no `rolename` and a non-`CSTRING` `roletype`, so they are never + * rewritten. + */ + +import { walk as walkSql } from '@pgsql/traverse'; +import { transformSync } from 'plpgsql-parser'; + +import type { RoleRouteSpec } from './role-router'; +import { RoleRouter } from './role-router'; + +/** What a role transform changed: role name -> number of positions rewritten. */ +export interface RoleTransformResult { + rolesRenamed: Map; +} + +export function createRoleResult(): RoleTransformResult { + return { rolesRenamed: new Map() }; +} + +const ROLE_SET_PARAMS = new Set(['role', 'session_authorization']); + +/** + * Create a SQL AST visitor that renames role identifiers against a router. + * Composable with the walkers used by the core transform. + */ +export function createRoleVisitor(router: RoleRouter, result: RoleTransformResult) { + const record = (name: string): void => { + result.rolesRenamed.set(name, (result.rolesRenamed.get(name) ?? 0) + 1); + }; + + return { + RoleSpec: (path: any) => { + const node = path.node; + // Only named roles (ROLESPEC_CSTRING) carry a rolename; PUBLIC / + // CURRENT_USER / SESSION_USER are encoded by roletype and left alone. + if (node.roletype !== 'ROLESPEC_CSTRING') return; + const target = router.resolve(node.rolename); + if (target === undefined) return; + node.rolename = target; + record(node.rolename); + }, + + CreateRoleStmt: (path: any) => { + const node = path.node; + const target = router.resolve(node.role); + if (target === undefined) return; + node.role = target; + record(target); + }, + + GrantRoleStmt: (path: any) => { + // granted_roles are role names carried as AccessPriv.priv_name (the + // grantee_roles RoleSpecs are handled by the RoleSpec visitor). + const granted = path.node.granted_roles; + if (!Array.isArray(granted)) return; + for (const g of granted) { + const priv = g?.AccessPriv; + if (!priv || typeof priv.priv_name !== 'string') continue; + const target = router.resolve(priv.priv_name); + if (target === undefined) continue; + priv.priv_name = target; + record(target); + } + }, + + VariableSetStmt: (path: any) => { + const node = path.node; + if (!ROLE_SET_PARAMS.has(node.name) || !Array.isArray(node.args)) return; + for (const arg of node.args) { + const sval = arg?.A_Const?.sval?.sval ?? arg?.String?.sval; + if (typeof sval !== 'string') continue; + const target = router.resolve(sval); + if (target === undefined) continue; + if (arg.A_Const?.sval) arg.A_Const.sval.sval = target; + else if (arg.String) arg.String.sval = target; + record(target); + } + }, + + RenameStmt: (path: any) => { + const node = path.node; + if (node.renameType !== 'OBJECT_ROLE') return; + const subTarget = router.resolve(node.subname); + if (subTarget !== undefined) { + node.subname = subTarget; + record(subTarget); + } + const newTarget = router.resolve(node.newname); + if (newTarget !== undefined) { + node.newname = newTarget; + record(newTarget); + } + } + }; +} + +/** + * Apply role renaming to a SQL string: parse -> walk -> deparse. Role names in + * PL/pgSQL / `LANGUAGE sql` bodies are string literals (e.g. inside dynamic + * `EXECUTE`) and are intentionally not rewritten — see the module docs on why + * arbitrary string contents are out of scope. + */ +export function transformRoles( + sql: string, + router: RoleRouter | RoleRouteSpec | Map +): { sql: string; result: RoleTransformResult } { + const resolved = RoleRouter.from(router); + const result = createRoleResult(); + + const out = transformSync(sql, (ctx: any) => { + const stmts: any[] = ctx.sql?.stmts ?? []; + const visitor = createRoleVisitor(resolved, result); + for (const stmt of stmts) { + if (stmt?.stmt) walkSql(stmt.stmt, visitor); + } + }, { hydrate: true, pretty: true }); + + return { sql: out, result }; +}