Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 106 additions & 0 deletions packages/transform/__tests__/routing-claims.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
import { loadModule } from 'plpgsql-parser';

import { classifyStatements } from '../src/facts';
import { SchemaRouter } from '../src/router';
import { transformSqlStatement } from '../src/transform';

beforeAll(async () => {
await loadModule();
});

const swap = new Map([
['a', 'b'],
['b', 'a']
]);

describe('single routing pass (claims)', () => {
it('applies a cyclic schema mapping exactly once per site', () => {
const cases: Array<[string, string]> = [
['CREATE TABLE a.t (id int);', 'CREATE TABLE b.t (\n id int\n);'],
['SELECT * FROM a.t;', 'SELECT *\nFROM b.t;'],
['INSERT INTO a.t VALUES (1);', 'INSERT INTO b.t VALUES\n (1);'],
['ALTER TABLE a.t ADD COLUMN c b.mytype;', 'ALTER TABLE b.t\n ADD COLUMN c a.mytype;'],
['CREATE VIEW a.v AS SELECT * FROM b.t;', 'CREATE VIEW b.v AS SELECT * FROM a.t;'],
['SELECT a.f(NULL::b.tp);', 'SELECT b.f(CAST(NULL AS a.tp));'],
['DROP TABLE a.t;', 'DROP TABLE b.t;'],
['CREATE INDEX i ON a.t (c);', 'CREATE INDEX i ON b.t (c);'],
[
'CREATE TABLE t (o uuid REFERENCES a.pk (id));',
'CREATE TABLE t (\n o uuid REFERENCES b.pk (id)\n);'
]
];
const norm = (s: string) => s.replace(/\s+/g, ' ').trim();
for (const [input, expected] of cases) {
expect(norm(transformSqlStatement(input, swap).sql)).toBe(norm(expected));
}
});

it('swaps a two-schema module without leftover-validation errors', () => {
const sql = [
'CREATE TABLE a.users (id uuid PRIMARY KEY);',
'CREATE TABLE b.posts (author uuid REFERENCES a.users (id));',
'CREATE FUNCTION a.author_of(p uuid) RETURNS uuid LANGUAGE sql AS $$ SELECT author FROM b.posts WHERE id = p $$;'
].join('\n');
const out = sql
.split('\n')
.map(stmt => transformSqlStatement(stmt, swap).sql)
.join('\n');
expect(out).toContain('b.users');
expect(out).toContain('a.posts');
expect(out).toContain('b.author_of');
expect(out).toContain('REFERENCES b.users');
});

it('statement-level namespace context wins over generic visitors', () => {
// Only a *function* route for a.f exists. The DropStmt handler routes with
// ns 'function'; the generic ObjectWithArgs visitor (ns 'unknown') must
// not route it a second time.
const router = new SchemaRouter({
a: { functions: { f: 'fns' } }
});
expect(transformSqlStatement('DROP FUNCTION a.f(int);', router).sql.trim()).toBe(
'DROP FUNCTION fns.f(int);'
);
expect(transformSqlStatement('ALTER FUNCTION a.f(int) OWNER TO u;', router).sql.trim()).toBe(
'ALTER FUNCTION fns.f(int) OWNER TO u;'
);
});

it('rebind with a cyclic name swap stays single-pass', () => {
const router = new SchemaRouter({
auth: { functions: { uid: { schema: null, name: 'current_user_id' } } }
});
expect(
transformSqlStatement('SELECT auth.uid();', router).sql.trim()
).toBe('SELECT current_user_id();');
});
});

describe('StatementFacts spans', () => {
it('reports each statement source span verbatim', () => {
const sql = `CREATE SCHEMA app;\nCREATE TABLE app.users (id uuid);\n\nSELECT 1;`;
const facts = classifyStatements(sql);
expect(facts).toHaveLength(3);
for (const f of facts) {
const text = sql.slice(f.span.start, f.span.start + f.span.len);
expect(text.trim().length).toBeGreaterThan(0);
}
const [schema, table, select] = facts;
expect(sql.slice(schema.span.start, schema.span.start + schema.span.len).trim()).toBe(
'CREATE SCHEMA app'
);
expect(sql.slice(table.span.start, table.span.start + table.span.len).trim()).toBe(
'CREATE TABLE app.users (id uuid)'
);
expect(sql.slice(select.span.start, select.span.start + select.span.len).trim()).toBe(
'SELECT 1'
);
});

it('covers the tail of the script for the final statement', () => {
const sql = 'SELECT 1'; // no trailing semicolon
const [f] = classifyStatements(sql);
expect(f.span.start).toBe(0);
expect(sql.slice(f.span.start, f.span.start + f.span.len)).toBe('SELECT 1');
});
});
21 changes: 20 additions & 1 deletion packages/transform/src/facts.ts
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,22 @@ export interface StatementFacts {
* are incomplete and slicing should treat it conservatively.
*/
dynamicSql: boolean;
/**
* The statement's source span in the classified script, as reported by the
* parser: `start` is the byte offset of the statement's first token, `len`
* runs to the end of the statement (the parser excludes the trailing `;`;
* for the final statement the span extends to the end of the script).
* `sql.slice(span.start, span.start + span.len)` is the statement's
* verbatim source, so consumers can carry original text alongside the
* facts without a second parse.
*/
span: StatementSpan;
}

/** A statement's location in the source script (byte offsets). */
export interface StatementSpan {
start: number;
len: number;
}

const SECURITY_TAGS = new Set([
Expand Down Expand Up @@ -272,7 +288,8 @@ function classifyOne(nodeTag: string, node: any): StatementFacts {
bodyReferences: [],
securityRelevant: SECURITY_TAGS.has(nodeTag),
securityDefiner: false,
dynamicSql: false
dynamicSql: false,
span: { start: 0, len: 0 }
};

switch (nodeTag) {
Expand Down Expand Up @@ -467,6 +484,8 @@ export function classifyStatements(sql: string): StatementFacts[] {
const nodeTag = stmtNode ? Object.keys(stmtNode)[0] : 'other';
const node = stmtNode?.[nodeTag] ?? {};
const facts = classifyOne(nodeTag, node);
const start = stmt?.stmt_location ?? 0;
facts.span = { start, len: stmt?.stmt_len ?? Math.max(0, sql.length - start) };

if (stmtNode) {
walkSql(stmtNode, createFactsVisitor(facts));
Expand Down
Loading
Loading