agent-sql 0.3.0 → 0.3.2

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.mjs

@@ -778,7 +778,12 @@ function checkWhereExpr(expr, allowed) {
 function defineSchema(schema) {
 	return schema;
 }
-function checkJoins(ast, schema) {
+function validateJoins(ast, schema) {
+	const ast2 = checkJoinColumns(ast, schema);
+	if (!ast2.ok) return ast2;
+	return checkJoinContinuity(ast);
+}
+function checkJoinColumns(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);
@@ -811,6 +816,122 @@ function getJoinTableRef(joinTableName, left, right) {
 		foreign: left
 	};
 }
+/**
+* Validates that all tables in the query's FROM/JOIN clauses form a single
+* connected component via the ON predicates. A disconnected table would produce
+* an implicit cross product which can leak data across tenants.
+*
+* Returns the AST unchanged if valid, or a SanitiseError listing the
+* disconnected tables.
+*/
+function checkJoinContinuity(ast) {
+	if (ast.joins.length === 0) return Ok(ast);
+	const tables = /* @__PURE__ */ new Set();
+	tables.add(ast.from.table.name);
+	for (const join of ast.joins) tables.add(join.table.name);
+	if (tables.size <= 1) return Ok(ast);
+	const adjacency = /* @__PURE__ */ new Map();
+	for (const table of tables) adjacency.set(table, /* @__PURE__ */ new Set());
+	for (const join of ast.joins) {
+		if (join.condition === null) continue;
+		if (join.condition.type === "join_using") continue;
+		const referencedTables = /* @__PURE__ */ new Set();
+		collectTableRefsFromExpr(join.condition.expr, referencedTables);
+		const relevantTables = [];
+		for (const t of referencedTables) if (tables.has(t)) relevantTables.push(t);
+		for (let i = 0; i < relevantTables.length; i++) for (let j = i + 1; j < relevantTables.length; j++) {
+			adjacency.get(relevantTables[i]).add(relevantTables[j]);
+			adjacency.get(relevantTables[j]).add(relevantTables[i]);
+		}
+	}
+	const start = ast.from.table.name;
+	const visited = /* @__PURE__ */ new Set();
+	const queue = [start];
+	visited.add(start);
+	while (queue.length > 0) {
+		const current = queue.shift();
+		for (const neighbor of adjacency.get(current) ?? []) if (!visited.has(neighbor)) {
+			visited.add(neighbor);
+			queue.push(neighbor);
+		}
+	}
+	if (visited.size === tables.size) return Ok(ast);
+	const disconnected = [];
+	for (const table of tables) if (!visited.has(table)) disconnected.push(table);
+	disconnected.sort();
+	return Err(new SanitiseError(`Disconnected table(s) in query: ${disconnected.join(", ")}. All tables must be connected via JOIN ON predicates.`));
+}
+/** Recursively collect table names from column references in a WhereExpr. */
+function collectTableRefsFromExpr(expr, out) {
+	switch (expr.type) {
+		case "where_and":
+		case "where_or":
+			collectTableRefsFromExpr(expr.left, out);
+			collectTableRefsFromExpr(expr.right, out);
+			break;
+		case "where_not":
+			collectTableRefsFromExpr(expr.expr, out);
+			break;
+		case "where_comparison":
+			collectTableRefsFromValue(expr.left, out);
+			collectTableRefsFromValue(expr.right, out);
+			break;
+		case "where_is_null":
+			collectTableRefsFromValue(expr.expr, out);
+			break;
+		case "where_is_bool":
+			collectTableRefsFromValue(expr.expr, out);
+			break;
+		case "where_between":
+			collectTableRefsFromValue(expr.expr, out);
+			collectTableRefsFromValue(expr.low, out);
+			collectTableRefsFromValue(expr.high, out);
+			break;
+		case "where_in":
+			collectTableRefsFromValue(expr.expr, out);
+			for (const v of expr.list) collectTableRefsFromValue(v, out);
+			break;
+		case "where_like":
+			collectTableRefsFromValue(expr.expr, out);
+			collectTableRefsFromValue(expr.pattern, out);
+			break;
+		case "where_ts_match":
+			collectTableRefsFromValue(expr.left, out);
+			collectTableRefsFromValue(expr.right, out);
+			break;
+	}
+}
+/** Recursively collect table names from column references in a WhereValue. */
+function collectTableRefsFromValue(val, out) {
+	switch (val.type) {
+		case "where_value":
+			if (val.kind === "column_ref" && val.ref.table) out.add(val.ref.table);
+			else if (val.kind === "func_call") {
+				if (val.func.args.kind === "args") for (const arg of val.func.args.args) collectTableRefsFromValue(arg, out);
+			}
+			break;
+		case "where_arith":
+		case "where_jsonb_op":
+		case "where_pgvector_op":
+			collectTableRefsFromValue(val.left, out);
+			collectTableRefsFromValue(val.right, out);
+			break;
+		case "where_unary_minus":
+			collectTableRefsFromValue(val.expr, out);
+			break;
+		case "case_expr":
+			if (val.subject) collectTableRefsFromValue(val.subject, out);
+			for (const w of val.whens) {
+				collectTableRefsFromValue(w.condition, out);
+				collectTableRefsFromValue(w.result, out);
+			}
+			if (val.else) collectTableRefsFromValue(val.else, out);
+			break;
+		case "cast_expr":
+			collectTableRefsFromValue(val.expr, out);
+			break;
+	}
+}
 //#endregion
 //#region src/graph.ts
 function insertNeededGuardJoins(ast, schema, guards, autoJoin) {
@@ -822,6 +943,7 @@ 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) {
@@ -839,8 +961,10 @@ function resolveGraphForJoins(ast, schema, guards) {
 		}
 	}
 	if (newJoins.length === 0) return Ok(ast);
+	const columns = qualifyWildcards(ast.columns, originalTables);
 	return Ok({
 		...ast,
+		columns,
 		joins: [...ast.joins, ...newJoins]
 	});
 }
@@ -927,6 +1051,20 @@ function edgeToJoin(edge, fromTable) {
 		}
 	};
 }
+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) {
@@ -7358,7 +7496,7 @@ function privateAgentSql(sql, { guards: guardsRaw, schema, autoJoin, limit, pret
 	if (!guards.ok) throw guards.error;
 	const ast = parseSql(sql);
 	if (!ast.ok) return returnOrThrow(ast, throws);
-	const ast2 = checkJoins(ast.data, schema);
+	const ast2 = validateJoins(ast.data, schema);
 	if (!ast2.ok) return returnOrThrow(ast2, throws);
 	const ast3 = checkFunctions(ast2.data, db, allowExtraFunctions);
 	if (!ast3.ok) return returnOrThrow(ast3, throws);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "agent-sql",
-  "version": "0.3.0",
+  "version": "0.3.2",
   "description": "A starter for creating a TypeScript package.",
   "keywords": [
     "agent",