agent-sql 0.2.3 → 0.3.1

package/README.md

@@ -4,20 +4,20 @@ Sanitise agent-written SQL for multi-tenant DBs.
 
 You provide a tenant ID, and the agent supplies the query.
 
-agent-sql works by fully parsing the supplied SQL query into an AST.
-The grammar ONLY accepts `SELECT` statements. Anything else is an error.
-CTEs and other complex things that we aren't confident of securing: error.
-
-It ensures that that the needed tenant table is somewhere in the query,
-and adds a `WHERE` clause ensuring that only values from the supplied ID are returned.
-Then it checks that the tables and `JOIN`s follow the schema, preventing sneaky joins.
+Apparently this is how [Trigger.dev does it](https://x.com/mattaitken/status/2033928542975639785).
+And [Cloudflare](https://x.com/thomas_ankcorn/status/2033931057133748330).
 
-Function calls also go through a whitelist (configurable).
+## How it works
 
-Finally, we throw in a `LIMIT` clause (configurable) to prevent accidental LLM denial-of-service.
+agent-sql works by fully parsing the supplied SQL query into an AST and transforming it:
 
-Apparently this is how [Trigger.dev does it](https://x.com/mattaitken/status/2033928542975639785).
-And [Cloudflare](https://x.com/thomas_ankcorn/status/2033931057133748330).
+- **Only `SELECT`:** it's impossible to insert, drop or anything else.
+- **Reduced subset:** CTEs, subqueries and other tricky things are rejected.
+- **Limited functions:** passed through a (configurable) whitelist.
+- **No DoS:** a default `LIMIT` is applied, but can be adjusted.
+- **`WHERE` guards:** insert multiple tenant/ownership conditions to be inserted.
+- **`JOIN`s added:** if needed to reach the guard tenant tables (save on tokens).
+- **No sneaky joins:** no `join secrets on true`. We have your back.
 
 ## Quickstart
 
@@ -28,17 +28,15 @@ npm install agent-sql
 ```ts
 import { agentSql } from "agent-sql";
 
-const sql = agentSql(`SELECT * FROM msg`, "msg.user_id", 123);
+const sql = agentSql(`SELECT * FROM msg`, "msg.tenant_id", 123);
 
 console.log(sql);
 // SELECT *
 // FROM msg
-// WHERE msg.user_id = 123
+// WHERE msg.tenant_id = 123
 // LIMIT 10000
 ```
 
-`agent-sql` parses the SQL, enforces a mandatory equality filter on the given column as the outermost `AND` condition (so it cannot be short-circuited by agent-supplied `OR` clauses), and returns the sanitised SQL string.
-
 ## Usage
 
 ### Define once, use many times

package/dist/index.d.mts

@@ -26,18 +26,21 @@ declare function parseSql(expr: string): Result<SelectStatement>;
 //#region src/index.d.ts
 declare function agentSql<S extends string>(sql: string, column: S & OneOrTwoDots<S>, value: GuardVal, {
   schema,
+  autoJoin,
   limit,
   pretty,
   db,
   allowExtraFunctions
 }?: {
   schema?: Schema;
+  autoJoin?: boolean;
   limit?: number;
   pretty?: boolean;
   db?: DbType;
   allowExtraFunctions?: string[];
 }): string;
 declare function createAgentSql<T extends Schema, S extends SchemaGuardKeys<T>>(schema: T, guards: Record<S, GuardVal>, opts: {
+  autoJoin?: boolean;
   limit?: number;
   pretty?: boolean;
   throws: false;
@@ -45,6 +48,7 @@ declare function createAgentSql<T extends Schema, S extends SchemaGuardKeys<T>>(
   allowExtraFunctions?: string[];
 }): (expr: string) => Result<string>;
 declare function createAgentSql<T extends Schema, S extends SchemaGuardKeys<T>>(schema: T, guards: Record<S, GuardVal>, opts?: {
+  autoJoin?: boolean;
   limit?: number;
   pretty?: boolean;
   throws?: true;

package/dist/index.mjs

@@ -774,6 +774,177 @@ function checkWhereExpr(expr, allowed) {
 	}
 }
 //#endregion
+//#region src/joins.ts
+function defineSchema(schema) {
+	return schema;
+}
+function checkJoins(ast, schema) {
+	if (schema === void 0) {
+		if (ast.joins.length > 0) return Err(new SanitiseError("No joins allowed when using simple API without schema."));
+		return Ok(ast);
+	}
+	if (!(ast.from.table.name in schema)) return Err(new SanitiseError(`Table ${ast.from.table.name} is not allowed`));
+	for (const join of ast.joins) {
+		const joinSettings = schema[join.table.name];
+		if (joinSettings === void 0) return Err(new SanitiseError(`Table ${join.table.name} is not allowed`));
+		if (join.condition === null || join.condition.type === "join_using" || join.condition.expr.type !== "where_comparison" || join.condition.expr.operator !== "=" || join.condition.expr.left.type !== "where_value" || join.condition.expr.left.kind !== "column_ref" || join.condition.expr.right.type !== "where_value" || join.condition.expr.right.kind !== "column_ref") return Err(new SanitiseError("Only JOIN ON column_ref = column_ref supported"));
+		const { joining, foreign } = getJoinTableRef(join.table.name, join.condition.expr.left.ref, join.condition.expr.right.ref);
+		const joinTableCol = joinSettings[joining.name];
+		if (joinTableCol === void 0) return Err(new SanitiseError(`Tried to join using ${join.table.name}.${joining.name}`));
+		if (joinTableCol === null) {
+			const foreignTableSettings = schema[foreign.table];
+			if (foreignTableSettings === void 0) return Err(new SanitiseError(`Table ${foreign.name} is not allowed`));
+			const foreignCol = foreignTableSettings[foreign.name];
+			if (foreignCol === void 0 || foreignCol === null) return Err(new SanitiseError(`Tried to join using ${foreign.table}.${foreign.name}`));
+			if (joining.table !== foreignCol.ft || joining.name !== foreignCol.fc) return Err(new SanitiseError(`Tried to join using ${joining.table}.${joining.name}`));
+		} else if (foreign.table !== joinTableCol.ft || foreign.name !== joinTableCol.fc) return Err(new SanitiseError(`Tried to join using ${foreign.table}.${foreign.name}`));
+	}
+	return Ok(ast);
+}
+function getJoinTableRef(joinTableName, left, right) {
+	if (left.table === joinTableName) return {
+		joining: left,
+		foreign: right
+	};
+	return {
+		joining: right,
+		foreign: left
+	};
+}
+//#endregion
+//#region src/graph.ts
+function insertNeededGuardJoins(ast, schema, guards, autoJoin) {
+	if (schema === void 0) return Ok(ast);
+	if (!autoJoin) return Ok(ast);
+	return resolveGraphForJoins(ast, schema, guards);
+}
+function resolveGraphForJoins(ast, schema, guards) {
+	const haveTables = /* @__PURE__ */ new Set();
+	haveTables.add(ast.from.table.name);
+	for (const join of ast.joins) haveTables.add(join.table.name);
+	const originalTables = new Set(haveTables);
+	const adj = buildAdjacency(schema);
+	const newJoins = [];
+	for (const guard of guards) {
+		if (haveTables.has(guard.table)) continue;
+		const path = bfsPath(adj, haveTables, guard.table);
+		if (path === null) return Err(new SanitiseError(`No join path from query tables to guard table '${guard.table}'`));
+		let current = [...haveTables].find((t) => path[0] && (path[0].tableA === t || path[0].tableB === t));
+		for (const edge of path) {
+			const neighbor = edge.tableA === current ? edge.tableB : edge.tableA;
+			if (!haveTables.has(neighbor)) {
+				newJoins.push(edgeToJoin(edge, neighbor));
+				haveTables.add(neighbor);
+			}
+			current = neighbor;
+		}
+	}
+	if (newJoins.length === 0) return Ok(ast);
+	const columns = qualifyWildcards(ast.columns, originalTables);
+	return Ok({
+		...ast,
+		columns,
+		joins: [...ast.joins, ...newJoins]
+	});
+}
+function buildAdjacency(schema) {
+	const adj = /* @__PURE__ */ new Map();
+	const tables = schema;
+	for (const [tableName, cols] of Object.entries(tables)) {
+		if (!adj.has(tableName)) adj.set(tableName, []);
+		for (const [colName, def] of Object.entries(cols)) if (def && typeof def === "object" && "ft" in def && "fc" in def) {
+			const edge = {
+				tableA: tableName,
+				colA: colName,
+				tableB: def.ft,
+				colB: def.fc
+			};
+			adj.get(tableName).push(edge);
+			if (!adj.has(edge.tableB)) adj.set(edge.tableB, []);
+			adj.get(edge.tableB).push(edge);
+		}
+	}
+	return adj;
+}
+function bfsPath(adj, startTables, target) {
+	if (startTables.has(target)) return [];
+	const visited = new Set(startTables);
+	const queue = [];
+	for (const t of startTables) queue.push([t, []]);
+	while (queue.length > 0) {
+		const [current, path] = queue.shift();
+		for (const edge of adj.get(current) ?? []) {
+			const neighbor = edge.tableA === current ? edge.tableB : edge.tableA;
+			if (visited.has(neighbor)) continue;
+			visited.add(neighbor);
+			const newPath = [...path, edge];
+			if (neighbor === target) return newPath;
+			queue.push([neighbor, newPath]);
+		}
+	}
+	return null;
+}
+function edgeToJoin(edge, fromTable) {
+	const [localTable, localCol, foreignTable, foreignCol] = edge.tableA === fromTable ? [
+		edge.tableA,
+		edge.colA,
+		edge.tableB,
+		edge.colB
+	] : [
+		edge.tableB,
+		edge.colB,
+		edge.tableA,
+		edge.colA
+	];
+	return {
+		type: "join",
+		joinType: "inner",
+		table: {
+			type: "table_ref",
+			name: localTable
+		},
+		condition: {
+			type: "join_on",
+			expr: {
+				type: "where_comparison",
+				operator: "=",
+				left: {
+					type: "where_value",
+					kind: "column_ref",
+					ref: {
+						type: "column_ref",
+						table: localTable,
+						name: localCol
+					}
+				},
+				right: {
+					type: "where_value",
+					kind: "column_ref",
+					ref: {
+						type: "column_ref",
+						table: foreignTable,
+						name: foreignCol
+					}
+				}
+			}
+		}
+	};
+}
+function qualifyWildcards(columns, tables) {
+	if (!columns.some((c) => c.expr.kind === "wildcard")) return columns;
+	const qualified = [];
+	for (const col of columns) if (col.expr.kind === "wildcard") for (const table of tables) qualified.push({
+		type: "column",
+		expr: {
+			type: "column_expr",
+			kind: "qualified_wildcard",
+			table
+		}
+	});
+	else qualified.push(col);
+	return qualified;
+}
+//#endregion
 //#region src/utils.ts
 function unreachable(x) {
 	throw new Error(`Unhandled variant: ${JSON.stringify(x)}`);
@@ -1108,44 +1279,6 @@ function tableEquals(a, b) {
 	return a.schema == b.schema && a.name == b.name;
 }
 //#endregion
-//#region src/joins.ts
-function defineSchema(schema) {
-	return schema;
-}
-function checkJoins(ast, schema) {
-	if (schema === void 0) {
-		if (ast.joins.length > 0) return Err(new SanitiseError("No joins allowed when using simple API without schema."));
-		return Ok(ast);
-	}
-	if (!(ast.from.table.name in schema)) return Err(new SanitiseError(`Table ${ast.from.table.name} is not allowed`));
-	for (const join of ast.joins) {
-		const joinSettings = schema[join.table.name];
-		if (joinSettings === void 0) return Err(new SanitiseError(`Table ${join.table.name} is not allowed`));
-		if (join.condition === null || join.condition.type === "join_using" || join.condition.expr.type !== "where_comparison" || join.condition.expr.operator !== "=" || join.condition.expr.left.type !== "where_value" || join.condition.expr.left.kind !== "column_ref" || join.condition.expr.right.type !== "where_value" || join.condition.expr.right.kind !== "column_ref") return Err(new SanitiseError("Only JOIN ON column_ref = column_ref supported"));
-		const { joining, foreign } = getJoinTableRef(join.table.name, join.condition.expr.left.ref, join.condition.expr.right.ref);
-		const joinTableCol = joinSettings[joining.name];
-		if (joinTableCol === void 0) return Err(new SanitiseError(`Tried to join using ${join.table.name}.${joining.name}`));
-		if (joinTableCol === null) {
-			const foreignTableSettings = schema[foreign.table];
-			if (foreignTableSettings === void 0) return Err(new SanitiseError(`Table ${foreign.name} is not allowed`));
-			const foreignCol = foreignTableSettings[foreign.name];
-			if (foreignCol === void 0 || foreignCol === null) return Err(new SanitiseError(`Tried to join using ${foreign.table}.${foreign.name}`));
-			if (joining.table !== foreignCol.ft || joining.name !== foreignCol.fc) return Err(new SanitiseError(`Tried to join using ${joining.table}.${joining.name}`));
-		} else if (foreign.table !== joinTableCol.ft || foreign.name !== joinTableCol.fc) return Err(new SanitiseError(`Tried to join using ${foreign.table}.${foreign.name}`));
-	}
-	return Ok(ast);
-}
-function getJoinTableRef(joinTableName, left, right) {
-	if (left.table === joinTableName) return {
-		joining: left,
-		foreign: right
-	};
-	return {
-		joining: right,
-		foreign: left
-	};
-}
-//#endregion
 //#region src/sql.ohm-bundle.js
 const result = makeRecipe([
 	"grammar",
@@ -7204,10 +7337,11 @@ function parseSql(expr) {
 }
 //#endregion
 //#region src/index.ts
-function agentSql(sql, column, value, { schema, limit = DEFAULT_LIMIT, pretty = false, db = DEFAULT_DB, allowExtraFunctions = [] } = {}) {
+function agentSql(sql, column, value, { schema, autoJoin = true, limit = DEFAULT_LIMIT, pretty = false, db = DEFAULT_DB, allowExtraFunctions = [] } = {}) {
 	return privateAgentSql(sql, {
 		guards: { [column]: value },
 		schema,
+		autoJoin,
 		limit,
 		pretty,
 		db,
@@ -7215,10 +7349,11 @@ function agentSql(sql, column, value, { schema, limit = DEFAULT_LIMIT, pretty =
 		throws: true
 	});
 }
-function createAgentSql(schema, guards, { limit = DEFAULT_LIMIT, pretty = false, db = DEFAULT_DB, allowExtraFunctions = [], throws = true } = {}) {
+function createAgentSql(schema, guards, { autoJoin = true, limit = DEFAULT_LIMIT, pretty = false, db = DEFAULT_DB, allowExtraFunctions = [], throws = true } = {}) {
 	return (expr) => throws ? privateAgentSql(expr, {
 		guards,
 		schema,
+		autoJoin,
 		limit,
 		pretty,
 		db,
@@ -7227,6 +7362,7 @@ function createAgentSql(schema, guards, { limit = DEFAULT_LIMIT, pretty = false,
 	}) : privateAgentSql(expr, {
 		guards,
 		schema,
+		autoJoin,
 		limit,
 		pretty,
 		db,
@@ -7234,7 +7370,7 @@ function createAgentSql(schema, guards, { limit = DEFAULT_LIMIT, pretty = false,
 		throws
 	});
 }
-function privateAgentSql(sql, { guards: guardsRaw, schema, limit, pretty, db, allowExtraFunctions, throws }) {
+function privateAgentSql(sql, { guards: guardsRaw, schema, autoJoin, limit, pretty, db, allowExtraFunctions, throws }) {
 	const guards = resolveGuards(guardsRaw);
 	if (!guards.ok) throw guards.error;
 	const ast = parseSql(sql);
@@ -7243,7 +7379,9 @@ function privateAgentSql(sql, { guards: guardsRaw, schema, limit, pretty, db, al
 	if (!ast2.ok) return returnOrThrow(ast2, throws);
 	const ast3 = checkFunctions(ast2.data, db, allowExtraFunctions);
 	if (!ast3.ok) return returnOrThrow(ast3, throws);
-	const san = applyGuards(ast3.data, guards.data, limit);
+	const ast4 = insertNeededGuardJoins(ast3.data, schema, guards.data, autoJoin);
+	if (!ast4.ok) return returnOrThrow(ast4, throws);
+	const san = applyGuards(ast4.data, guards.data, limit);
 	if (!san.ok) return returnOrThrow(san, throws);
 	const res = outputSql(san.data, pretty);
 	if (throws) return res;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "agent-sql",
-  "version": "0.2.3",
+  "version": "0.3.1",
   "description": "A starter for creating a TypeScript package.",
   "keywords": [
     "agent",
@@ -36,7 +36,7 @@
     "test": "vp test",
     "check": "vp check",
     "prepublishOnly": "vp run build",
-    "gen:types": "ohm generateBundles --withTypes --esm 'src/sql.ohm'",
+    "ohm": "ohm generateBundles --withTypes --esm 'src/sql.ohm'",
     "script": "vp exec tsx"
   },
   "dependencies": {