feat(transform): route object identities by cast type

packages/transform/__tests__/identity-casts.test.ts

@@ -0,0 +1,121 @@
+import { loadModule } from 'plpgsql-parser';
+
+import { SchemaRouter, transformSql, transformSqlStatement } from '../src';
+
+beforeAll(async () => {
+  await loadModule();
+});
+
+const MAPPING = new Map([
+  ['my-schema', 'my_schema'],
+  ['other-schema', 'other_schema']
+]);
+
+function run(sql: string, mapping: Map<string, string> | SchemaRouter = MAPPING): string {
+  return transformSql(sql, mapping).content;
+}
+
+describe('object-identity casts', () => {
+  it('routes a bare schema name through ::regnamespace', () => {
+    expect(run(`SELECT assert_schema('my-schema'::regnamespace);`))
+      .toContain(`CAST('my_schema' AS regnamespace)`);
+  });
+
+  it('routes the qualifier of ::regclass', () => {
+    expect(run(`SELECT assert_table('my-schema.users'::regclass);`))
+      .toContain(`CAST('my_schema.users' AS regclass)`);
+  });
+
+  it('routes ::regprocedure and keeps the argument list', () => {
+    expect(run(`SELECT assert_function('my-schema.fn(uuid, text)'::regprocedure);`))
+      .toContain(`CAST('my_schema.fn(uuid, text)' AS regprocedure)`);
+  });
+
+  it('routes a schema-qualified argument type', () => {
+    expect(
+      run(`SELECT assert_function('my-schema.fn("other-schema".row_t)'::regprocedure);`)
+    ).toContain(`CAST('my_schema.fn("other_schema".row_t)' AS regprocedure)`);
+  });
+
+  it('routes an argument type under an unqualified function name', () => {
+    expect(
+      run(`SELECT assert_function('fn("other-schema".row_t)'::regprocedure);`)
+    ).toContain(`CAST('fn("other_schema".row_t)' AS regprocedure)`);
+  });
+
+  it('routes ::regproc', () => {
+    expect(run(`SELECT assert_trigger('my-schema.users'::regclass, 'stamps', 'other-schema.tg'::regproc);`))
+      .toContain(`CAST('other_schema.tg' AS regproc)`);
+  });
+
+  it('routes ::regtype', () => {
+    expect(run(`SELECT CAST('my-schema.status'::regtype AS text);`))
+      .toContain(`CAST('my_schema.status' AS regtype)`);
+  });
+
+  it('leaves an unqualified identity alone (search_path resolves it)', () => {
+    expect(run(`SELECT assert_table('users'::regclass);`)).toContain(`CAST('users' AS regclass)`);
+  });
+
+  it('leaves an unmapped schema alone', () => {
+    expect(run(`SELECT assert_schema('stamps'::regnamespace);`))
+      .toContain(`CAST('stamps' AS regnamespace)`);
+  });
+
+  it('quotes a target name that is not a bare identifier', () => {
+    const mapping = new Map([['my_schema', 'my-schema']]);
+    expect(run(`SELECT assert_schema('my_schema'::regnamespace);`, mapping))
+      .toContain(`CAST('"my-schema"' AS regnamespace)`);
+    expect(run(`SELECT assert_table('my_schema.users'::regclass);`, mapping))
+      .toContain(`CAST('"my-schema".users' AS regclass)`);
+  });
+
+  it('reads an already-quoted operand', () => {
+    expect(run(`SELECT assert_table('"my-schema".users'::regclass);`))
+      .toContain(`CAST('my_schema.users' AS regclass)`);
+  });
+
+  it('routes a pg_catalog-qualified cast', () => {
+    expect(run(`SELECT assert_schema('my-schema'::pg_catalog.regnamespace);`))
+      .toContain(`'my_schema'`);
+  });
+
+  it('does not treat a non-identity cast as a reference', () => {
+    expect(run(`SELECT 'my-schema'::text;`)).toContain(`'my-schema'::text`);
+  });
+
+  it('routes each identity exactly once under a cyclic mapping', () => {
+    const swap = new Map([['a', 'b'], ['b', 'a']]);
+    const out = transformSqlStatement(
+      `SELECT assert_schema('a'::regnamespace), assert_table('a.users'::regclass);`,
+      swap
+    ).sql;
+    expect(out).toContain(`CAST('b' AS regnamespace)`);
+    expect(out).toContain(`CAST('b.users' AS regclass)`);
+  });
+
+  it('applies object-level routes to a qualified identity', () => {
+    const router = new SchemaRouter({
+      'my-schema': { schema: 'my_schema', relations: { users: 'people_schema' } }
+    });
+    const out = run(
+      `SELECT assert_table('my-schema.users'::regclass), assert_table('my-schema.orders'::regclass);`,
+      router
+    );
+    expect(out).toContain(`CAST('people_schema.users' AS regclass)`);
+    expect(out).toContain(`CAST('my_schema.orders' AS regclass)`);
+  });
+
+  it('rebinds an object route that renames', () => {
+    const router = new SchemaRouter({
+      'my-schema': { functions: { uid: { schema: null, name: 'current_user_id' } } }
+    });
+    expect(run(`SELECT assert_function('my-schema.uid(uuid)'::regprocedure);`, router))
+      .toContain(`CAST('current_user_id(uuid)' AS regprocedure)`);
+  });
+
+  it('routes identities inside a DO body', () => {
+    const sql = `DO $$ BEGIN PERFORM 'my-schema.users'::regclass; END $$;`;
+    expect(run(sql)).toContain('my_schema.users');
+  });
+});

packages/transform/src/index.ts

@@ -96,8 +96,10 @@ export {
   createSqlVisitor,
   escapeRegexp,
   extractPgpmHeader,
+  identityCastNamespace,
   shouldTransformSchema,
   transformComments,
+  transformIdentityCastLiteral,
   transformJsonStringValues,
   transformNameList,
   transformPlpgsqlTypeAst,

packages/transform/src/transform.ts

@@ -304,6 +304,105 @@ export function transformSchemaRefsInString(
   return out;
 }
 
+/**
+ * Cast types whose string operand is an object reference rather than opaque
+ * text. `'sch.tbl'::regclass` is a relation the parser resolves, and the cast
+ * names the namespace it resolves in — so the reference can be routed from the
+ * cast, semantically, instead of pattern-matched. This reaches what the string
+ * passes cannot: a bare `'sch'::regnamespace` has no dot to match on.
+ */
+const IDENTITY_CAST_NAMESPACES: Record<string, RouteNamespace> = {
+  regclass: 'relation',
+  regcollation: 'unknown',
+  regconfig: 'unknown',
+  regdictionary: 'unknown',
+  regnamespace: 'schema',
+  regoper: 'unknown',
+  regoperator: 'unknown',
+  regproc: 'function',
+  regprocedure: 'function',
+  regtype: 'type'
+};
+
+/**
+ * The namespace a cast resolves its operand in, or `undefined` when the target
+ * type is not an object-identity type.
+ */
+export function identityCastNamespace(typeName: any): RouteNamespace | undefined {
+  const names = typeName?.names;
+  if (!Array.isArray(names) || names.length === 0) return undefined;
+  // Only `pg_catalog`-qualified (or unqualified) reg* types are the built-ins.
+  if (names.length > 1 && names[0]?.String?.sval !== 'pg_catalog') return undefined;
+  const typeIdent = names[names.length - 1]?.String?.sval;
+  if (typeof typeIdent !== 'string') return undefined;
+  return IDENTITY_CAST_NAMESPACES[typeIdent.toLowerCase()];
+}
+
+const IDENTITY_TOKEN = String.raw`(?:"(?:[^"]|"")*"|[^\s".,()[\]]+)`;
+const BARE_IDENTITY_RE = new RegExp(`^\\s*(${IDENTITY_TOKEN})\\s*$`);
+const QUALIFIED_IDENTITY_RE = new RegExp(`^\\s*(${IDENTITY_TOKEN})\\.([\\s\\S]+)$`);
+
+function unquoteIdentifier(token: string): string {
+  if (token.length > 1 && token.startsWith('"') && token.endsWith('"')) {
+    return token.slice(1, -1).replace(/""/g, '"');
+  }
+  return token;
+}
+
+/**
+ * Route the object reference carried in the string operand of an identity
+ * cast. `ns` comes from the cast's target type.
+ *
+ * Unqualified operands (`'tbl'::regclass`) name no schema — they resolve
+ * through `search_path` — so they are returned unchanged.
+ */
+export function transformIdentityCastLiteral(
+  sval: string,
+  ns: RouteNamespace,
+  schemaMapping: SchemaMappingInput,
+  result: SchemaTransformResult
+): string {
+  const router = asRouter(schemaMapping);
+
+  if (ns === 'schema') {
+    const bare = BARE_IDENTITY_RE.exec(sval);
+    if (!bare) return sval;
+    const schemaName = unquoteIdentifier(bare[1]);
+    const newName = router.resolve(schemaName, undefined, 'schema');
+    if (!newName || newName === schemaName) return sval;
+    result.schemasFound.add(schemaName);
+    result.schemasTransformed.set(schemaName, newName);
+    return QuoteUtils.quoteIdentifier(newName);
+  }
+
+  // `regprocedure` and `regoperator` carry an argument list after the name, and
+  // an argument type can itself be schema-qualified
+  // (`fn("my-schema".row_type)`), so the list is routed as ordinary schema refs.
+  const argsAt = sval.indexOf('(');
+  const head = argsAt === -1 ? sval : sval.slice(0, argsAt);
+  const args =
+    argsAt === -1 ? '' : transformSchemaRefsInString(sval.slice(argsAt), router, result);
+
+  const qualified = QUALIFIED_IDENTITY_RE.exec(head);
+  if (!qualified) return `${head}${args}`;
+  const schemaName = unquoteIdentifier(qualified[1]);
+  const objName = unquoteIdentifier(qualified[2].trim());
+
+  const target = router.resolveObject(schemaName, objName, ns);
+  if (!target) return `${head}${args}`;
+  const rebound = target.name !== undefined && target.name !== objName;
+  const requalified = target.schema === null || (!!target.schema && target.schema !== schemaName);
+  if (!rebound && !requalified) return `${head}${args}`;
+
+  result.schemasFound.add(schemaName);
+  const nameToken = QuoteUtils.quoteIdentifier(rebound ? target.name! : objName);
+  if (target.schema === null) return `${nameToken}${args}`;
+  if (target.schema && target.schema !== schemaName) {
+    result.schemasTransformed.set(schemaName, target.schema);
+  }
+  return `${QuoteUtils.quoteIdentifier(target.schema ?? schemaName)}.${nameToken}${args}`;
+}
+
 /**
  * Create a SQL AST visitor that transforms schema names.
  * 
@@ -569,10 +668,26 @@ export function createSqlVisitor(
     A_Const: (path: any) => {
       const node = path.node;
       if (typeof node.sval?.sval === 'string' && node.sval.sval.includes('.')) {
+        if (!claimSite(result, node.sval, 'sval')) return;
         node.sval.sval = transformSchemaRefsInString(node.sval.sval, schemaMapping, result);
       }
     },
 
+    // Object identities carried as the operand of an identity cast
+    // ('sch.tbl'::regclass, 'sch.fn(uuid)'::regprocedure, 'sch'::regnamespace).
+    // The cast names the namespace, so these route through the same object
+    // routes as a RangeVar rather than through a string pattern — and the bare
+    // schema form, invisible to every string pass, routes here.
+    TypeCast: (path: any) => {
+      const node = path.node;
+      const sval = node?.arg?.A_Const?.sval;
+      if (typeof sval?.sval !== 'string') return;
+      const ns = identityCastNamespace(node.typeName);
+      if (!ns) return;
+      if (!claimSite(result, sval, 'sval')) return;
+      sval.sval = transformIdentityCastLiteral(sval.sval, ns, router, result);
+    },
+
     // Transform DefineStmt (CREATE TYPE schema.name, CREATE AGGREGATE schema.agg)
     DefineStmt: (path: any) => {
       const node = path.node;