agent-sql 0.3.2 → 0.3.4

package/README.md

@@ -19,6 +19,8 @@ agent-sql works by fully parsing the supplied SQL query into an AST and transfor
 - **`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.
 
+I plan to support inserts, updates, CTEs, subqueries once I can convincingly make them safe.
+
 ## Quickstart
 
 ```bash
@@ -28,7 +30,7 @@ npm install agent-sql
 ```ts
 import { agentSql } from "agent-sql";
 
-const sql = agentSql(`SELECT * FROM msg`, "msg.tenant_id", 123);
+const sql = agentSql("SELECT * FROM msg", "msg.tenant_id", 123);
 
 console.log(sql);
 // SELECT *
@@ -39,30 +41,81 @@ console.log(sql);
 
 ## Usage
 
-### Define once, use many times
+### Define a schema
 
-The simple approach above is enough to get started.
-But since no schema is provided, `JOIN`s will be blocked.
-A schema can be passed to `agentSql`, but typically you'll want to set it up once and re-use.
+In the simple example above, all `JOIN`s will be blocked.
+For agent-sql to know what joins and tables permit, you need to define a schema.
+Heads up: if you use Drizzle, you can just [use your Drizzle schema](#integration-with-ai-sdk-and-drizzle).
 
 ```ts
 import { createAgentSql, defineSchema } from "agent-sql";
-import { tool } from "ai";
-import { sql } from "drizzle-orm";
-import { db } from "@/db";
 
 // Define your schema.
 // Only the tables listed will be permitted
 // Joins can only use the FKs defined here
 const schema = defineSchema({
-  user: { id: null },
-  msg: { userId: { ft: "user", fc: "id" } },
+  tenant: { id: null },
+  msg: { tenant_id: { ft: "tenant", fc: "id" } },
 });
 
-function makeSqlTool(userId: string) {
+// Use your schema from above
+// Specify 1+ column->value pairs that will be enforced
+const agentSql = createAgentSql(schema, { "tenant.id": 123 });
+
+// Now use it
+const sql = agentSql("SELECT * FROM msg");
+```
+
+Outputs:
+
+```sql
+SELECT
+  msg.*                        -- qualify the *
+FROM msg
+INNER JOIN tenant              -- add the needed join for the guard
+  ON tenant.id = msg.tenant_id -- use the schema to join correctly
+WHERE tenant.id = 123          -- apply the guard
+LIMIT 10000                    -- limit the rows
+```
+
+### Bad stuff is blocked
+
+The following query will be blocked (many times over).
+
+```sql
+SELECT
+    sneaky_func('./bad_file')      -- won't pass whitelist
+FROM secret
+JOIN random                        -- not an approved table
+  ON random.id = secret.id         -- not an approved FK pair
+JOIN danger                        -- disconnected from join graph
+  ON true                          -- not allowed
+WHERE true                         -- won't trick anyone
+```
+
+### Integration with AI SDK and Drizzle
+
+If you're using Drizzle, you can skip the schema step and use the one you already have!
+
+Just pass it through, and `agentSql` will respect your schema.
+
+```ts
+import { tool } from "ai";
+import { sql } from "drizzle-orm";
+
+import { createAgentSql } from "agent-sql";
+import { defineSchemaFromDrizzle } from "agent-sql/drizzle";
+
+import { db } from "@/db";
+import * as drizzleSchema from "@/db/schema";
+
+// No need to re-enter your schema, we'll pull it in from Drizzle
+const schema = defineSchemaFromDrizzle(drizzleSchema);
+
+function makeSqlTool(tenantId: string) {
   // Create a sanitiser function for this tenant
   // Specify one or more column->value pairs that will be enforced
-  const agentSql = createAgentSql(schema, { "user.id": userId });
+  const agentSql = createAgentSql(schema, { "tenant.id": tenantId });
 
   return tool({
     description: "Run raw SQL against the DB",
@@ -70,34 +123,22 @@ function makeSqlTool(userId: string) {
     execute: async ({ query }) => {
       // The LLM can pass any query it likes, we'll sanitise it if possible
       // and return helpful error messages if not
-      const sanitised = agentSql(query);
+      const sql = agentSql(query);
       // Now we can throw that straight at the db and be confident it'll only
       // return data from the specified tenant
-      return db.execute(sql.raw(sanitised));
+      return db.execute(sql.raw(sql));
     },
   });
 }
 ```
 
-### It works with Drizzle
-
-If you're using Drizzle, you can skip the schema step and use the one you already have!
+### If you don't want your whole Drizzle schema available
 
-Just pass it through, and `agentSql` will respect your schema.
+You can also exclude tables if you don't want agents to see them:
 
 ```ts
 import { defineSchemaFromDrizzle } from "agent-sql/drizzle";
-import * as drizzleSchema from "@/db/schema";
-
-const schema = defineSchemaFromDrizzle(drizzleSchema);
 
-// The rest as before...
-const agentSql = createAgentSql(schema, { "user.id": userId });
-```
-
-You can also exclude tables if you don't want agents to see them:
-
-```ts
 const schema = defineSchemaFromDrizzle(drizzleSchema, {
   exclude: ["api_keys"],
 });

package/dist/index.mjs

@@ -793,7 +793,10 @@ function checkJoinColumns(ast, schema) {
 		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 leftRef = join.condition.expr.left.ref;
+		const rightRef = join.condition.expr.right.ref;
+		if (leftRef.table !== join.table.name && rightRef.table !== join.table.name) return Err(new SanitiseError(`JOIN ${join.table.name} ON clause does not reference ${join.table.name}`));
+		const { joining, foreign } = getJoinTableRef(join.table.name, leftRef, rightRef);
 		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) {
@@ -7448,7 +7451,10 @@ semantics.addOperation("toAST()", {
 });
 function parseSql(expr) {
 	const matchResult = result.match(expr);
-	if (matchResult.failed()) return Err(new ParseError(matchResult.message));
+	if (matchResult.failed()) {
+		const [message] = matchResult.message.split("\nExpected");
+		return Err(new ParseError(`Got invalid SQL. Some language features are deliberately disabled. Re-write your query without them.\n${message}`));
+	}
 	try {
 		return Ok(semantics(matchResult).toAST());
 	} catch (e) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "agent-sql",
-  "version": "0.3.2",
+  "version": "0.3.4",
   "description": "A starter for creating a TypeScript package.",
   "keywords": [
     "agent",
@@ -34,6 +34,7 @@
     "build": "vp pack",
     "dev": "vp pack --watch",
     "test": "vp test",
+    "coverage": "vp test --coverage",
     "check": "vp check",
     "prepublishOnly": "vp run build",
     "ohm": "ohm generateBundles --withTypes --esm 'src/sql.ohm'",
@@ -48,6 +49,7 @@
     "@types/node": "^25.5.0",
     "@types/pg": "^8.18.0",
     "@typescript/native-preview": "7.0.0-dev.20260316.1",
+    "@vitest/coverage-v8": "^4.1.0",
     "bumpp": "^11.0.1",
     "drizzle-orm": "^0.45.1",
     "pg": "^8.20.0",