agent-sql 0.1.2 → 0.2.1

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Chris Arderne
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/README.md

@@ -8,10 +8,14 @@ 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.
 
-Then we ensure that that the needed tenant table is somewhere in the query,
-and add a `WHERE` clause ensuring that only values from the supplied ID are returned.
+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).
+Finally, we throw in a `LIMIT` clause (configurable) to prevent accidental LLM denial-of-service.
+
+Apparently this is how [Trigger.dev does it](https://x.com/mattaitken/status/2033928542975639785).
+And [Cloudflare](https://x.com/thomas_ankcorn/status/2033931057133748330).
 
 ## Quickstart
 
@@ -20,34 +24,45 @@ npm install agent-sql
 ```
 
 ```ts
-import { sanitise } from "agent-sql";
+import { agentSql } from "agent-sql";
 
-const sql = sanitise(`SELECT id, name FROM users WHERE status = 'active' LIMIT 10`, {
-  tables: { users: {} },
-  where: { table: "users", col: "tenant_id", value: "acme" },
-});
+const sql = agentSql(`SELECT * FROM msg`, "msg.user_id", 123);
 
 console.log(sql);
-// SELECT id, name
-// FROM users
-// WHERE (users.tenant_id = 'acme' AND status = 'active')
-// LIMIT 10
+// SELECT *
+// FROM msg
+// WHERE msg.user_id = 123
+// LIMIT 10000
 ```
 
-Or, more usefully:
+`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
+
+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.
 
 ```ts
-import { sanitiserFactory } from "agent-sql";
+import { createAgentSql, defineSchema } from "agent-sql";
 import { tool } from "ai";
 import { sql } from "drizzle-orm";
 import { db } from "@/db";
 
-function makeSqlTool(orgId: string) {
+// 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" } },
+});
+
+function makeSqlTool(userId: string) {
   // Create a sanitiser function for this tenant
-  const sanitise = sanitiserFactory({
-    tables: {},
-    where: { table: "org", col: "id", value: orgId },
-  });
+  // Specify one or more column->value pairs that will be enforced
+  const agentSql = createAgentSql(schema, { "user.id": userId });
 
   return tool({
     description: "Run raw SQL against the DB",
@@ -55,7 +70,7 @@ function makeSqlTool(orgId: 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 = sanitise(query);
+      const sanitised = 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));
@@ -64,7 +79,29 @@ function makeSqlTool(orgId: string) {
 }
 ```
 
-`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.
+### It works with 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 { 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"],
+});
+```
 
 ## Development
 

package/dist/drizzle.d.mts

@@ -0,0 +1,22 @@
+import { t as Schema } from "./joins-Cu_0yAgN.mjs";
+import { Table, TableConfig } from "drizzle-orm";
+
+//#region src/drizzle.d.ts
+/** Extract all table name strings from a drizzle schema module. */
+type TableNames<T> = { [K in keyof T]: T[K] extends Table<infer C extends TableConfig> ? C["name"] : never }[keyof T];
+/**
+ * Build an agent-sql schema from a drizzle-orm schema module.
+ *
+ * Usage:
+ * ```ts
+ * import * as drizzleSchema from "./schema";
+ * const schema = defineSchemaFromDrizzle(drizzleSchema);
+ * ```
+ */
+declare function defineSchemaFromDrizzle<T extends Record<string, unknown>>(drizzleSchema: T, {
+  exclude
+}?: {
+  exclude?: TableNames<T>[];
+}): Schema;
+//#endregion
+export { defineSchemaFromDrizzle };

package/dist/drizzle.mjs

@@ -0,0 +1,66 @@
+import { getTableName, isTable } from "drizzle-orm";
+//#region src/drizzle.ts
+/**
+* Convert a camelCase string to snake_case.
+* This matches drizzle-orm's internal conversion for columns without explicit DB names.
+*/
+function camelToSnake(input) {
+	return (input.replace(/['\u2019]/g, "").match(/[\da-z]+|[A-Z]+(?![a-z])|[A-Z][\da-z]+/g) ?? []).map((w) => w.toLowerCase()).join("_");
+}
+/** Get the DB column name, applying camelCase → snake_case when the key was used as name. */
+function getDbColumnName(col) {
+	return col.keyAsName ? camelToSnake(col.name) : col.name;
+}
+const ColumnsSymbol = Symbol.for("drizzle:Columns");
+const InlineForeignKeysSymbol = Symbol.for("drizzle:PgInlineForeignKeys");
+/**
+* Build an agent-sql schema from a drizzle-orm schema module.
+*
+* Usage:
+* ```ts
+* import * as drizzleSchema from "./schema";
+* const schema = defineSchemaFromDrizzle(drizzleSchema);
+* ```
+*/
+function defineSchemaFromDrizzle(drizzleSchema, { exclude } = {}) {
+	const schema = {};
+	const excluded = new Set(exclude);
+	const tables = [];
+	for (const value of Object.values(drizzleSchema)) if (isTable(value) && !excluded.has(getTableName(value))) tables.push(value);
+	const fkMap = /* @__PURE__ */ new Map();
+	for (const table of tables) {
+		const tableName = getTableName(table);
+		const fks = table[InlineForeignKeysSymbol] ?? [];
+		const colFks = /* @__PURE__ */ new Map();
+		for (const fk of fks) {
+			const ref = fk.reference();
+			const fromCol = ref.columns[0];
+			const foreignTableName = getTableName(ref.foreignTable);
+			if (excluded.has(foreignTableName)) colFks.set(getDbColumnName(fromCol), "excluded");
+			else {
+				const toCol = ref.foreignColumns[0];
+				colFks.set(getDbColumnName(fromCol), {
+					ft: foreignTableName,
+					fc: getDbColumnName(toCol)
+				});
+			}
+		}
+		fkMap.set(tableName, colFks);
+	}
+	for (const table of tables) {
+		const tableName = getTableName(table);
+		const columns = Object.values(table[ColumnsSymbol]);
+		const colFks = fkMap.get(tableName);
+		const tableSchema = {};
+		for (const col of columns) {
+			const dbName = getDbColumnName(col);
+			const fk = colFks.get(dbName);
+			if (fk === "excluded") continue;
+			tableSchema[dbName] = fk ?? null;
+		}
+		schema[tableName] = tableSchema;
+	}
+	return schema;
+}
+//#endregion
+export { defineSchemaFromDrizzle };

package/dist/index.d.mts

@@ -1,267 +1,18 @@
-//#region src/ast.d.ts
-interface SelectStatement {
-  readonly type: "select";
-  distinct: Distinct | DistinctOn | null;
-  columns: Column[];
-  from: SelectFrom;
-  joins: JoinClause[];
-  where: WhereRoot | null;
-  groupBy: GroupByClause | null;
-  having: HavingClause | null;
-  orderBy: OrderByClause | null;
-  limit: LimitClause | null;
-  offset: OffsetClause | null;
-}
-interface Distinct {
-  readonly type: "distinct";
-}
-/** PostgreSQL: DISTINCT ON (col1, col2, ...) */
-interface DistinctOn {
-  readonly type: "distinct_on";
-  columns: WhereValue[];
-}
-interface GroupByClause {
-  readonly type: "group_by";
-  items: WhereValue[];
-}
-interface HavingClause {
-  readonly type: "having";
-  expr: WhereExpr;
-}
-interface OrderByClause {
-  readonly type: "order_by";
-  items: OrderByItem[];
-}
-type SortDirection = "asc" | "desc";
-type NullsOrder = "nulls_first" | "nulls_last";
-interface OrderByItem {
-  readonly type: "order_by_item";
-  expr: WhereValue;
-  direction?: SortDirection;
-  nulls?: NullsOrder;
-}
-interface OffsetClause {
-  readonly type: "offset";
-  value: number;
-}
-type JoinType = "inner" | "inner_outer" | "left" | "left_outer" | "right" | "right_outer" | "full" | "full_outer" | "cross" | "natural";
-type JoinCondition = {
-  readonly type: "join_on";
-  expr: WhereExpr;
-} | {
-  readonly type: "join_using";
-  columns: string[];
-};
-interface JoinClause {
-  readonly type: "join";
-  joinType: JoinType;
-  table: TableRef;
-  condition: JoinCondition | null;
-}
-type SelectFrom = {
-  readonly type: "select_from";
-  table: TableRef;
-};
-type LimitClause = {
-  readonly type: "limit";
-  value: number;
-};
-type WhereRoot = {
-  readonly type: "where_root";
-  inner: WhereExpr;
-};
-type WhereExpr = WhereAnd | WhereOr | WhereNot | WhereComparison | WhereIsNull | WhereIsBool | WhereBetween | WhereIn | WhereLike | WhereTsMatch;
-interface WhereAnd {
-  readonly type: "where_and";
-  left: WhereExpr;
-  right: WhereExpr;
-}
-interface WhereOr {
-  readonly type: "where_or";
-  left: WhereExpr;
-  right: WhereExpr;
-}
-type ComparisonOperator = "=" | "<>" | "!=" | "<" | ">" | "<=" | ">=";
-interface WhereNot {
-  readonly type: "where_not";
-  expr: WhereExpr;
-}
-interface WhereIsNull {
-  readonly type: "where_is_null";
-  not: boolean;
-  expr: WhereValue;
-}
-type IsBoolTarget = boolean | "unknown";
-interface WhereIsBool {
-  readonly type: "where_is_bool";
-  not: boolean;
-  expr: WhereValue;
-  target: IsBoolTarget;
-}
-interface WhereBetween {
-  readonly type: "where_between";
-  not: boolean;
-  expr: WhereValue;
-  low: WhereValue;
-  high: WhereValue;
-}
-interface WhereIn {
-  readonly type: "where_in";
-  not: boolean;
-  expr: WhereValue;
-  list: WhereValue[];
-}
-/** "ilike" is PostgreSQL-specific (case-insensitive LIKE) */
-type LikeOp = "like" | "ilike";
-interface WhereLike {
-  readonly type: "where_like";
-  not: boolean;
-  op: LikeOp;
-  expr: WhereValue;
-  pattern: WhereValue;
-}
-/** PostgreSQL: JSONB operators */
-type JsonbOp = "->" | "->>" | "#>" | "#>>" | "?" | "?|" | "?&" | "@>";
-/** PostgreSQL: JSONB binary operator expression */
-interface WhereJsonbOp {
-  readonly type: "where_jsonb_op";
-  op: JsonbOp;
-  left: WhereValue;
-  right: WhereValue;
-}
-/** PostgreSQL: text search match (@@) */
-interface WhereTsMatch {
-  readonly type: "where_ts_match";
-  left: WhereValue;
-  right: WhereValue;
-}
-/** pgvector: distance operators */
-type PgvectorOp = "<->" | "<#>" | "<=>" | "<+>" | "<~>" | "<%>";
-/** pgvector: distance operator expression */
-interface WherePgvectorOp {
-  readonly type: "where_pgvector_op";
-  op: PgvectorOp;
-  left: WhereValue;
-  right: WhereValue;
-}
-type ArithOp = "+" | "-" | "*" | "/" | "%" | "||";
-interface WhereArith {
-  readonly type: "where_arith";
-  op: ArithOp;
-  left: WhereValue;
-  right: WhereValue;
-}
-interface WhereUnaryMinus {
-  readonly type: "where_unary_minus";
-  expr: WhereValue;
-}
-interface CaseWhen {
-  condition: WhereValue;
-  result: WhereValue;
-}
-interface CaseExpr {
-  readonly type: "case_expr";
-  subject: WhereValue | null;
-  whens: CaseWhen[];
-  else: WhereValue | null;
-}
-interface CastExpr {
-  readonly type: "cast_expr";
-  expr: WhereValue;
-  typeName: string;
-}
-type WhereValue = {
-  readonly type: "where_value";
-  kind: "string";
-  value: string;
-} | {
-  readonly type: "where_value";
-  kind: "integer";
-  value: number;
-} | {
-  readonly type: "where_value";
-  kind: "float";
-  value: number;
-} | {
-  readonly type: "where_value";
-  kind: "bool";
-  value: boolean;
-} | {
-  readonly type: "where_value";
-  kind: "null";
-} | {
-  readonly type: "where_value";
-  kind: "column_ref";
-  ref: ColumnRef;
-} | {
-  readonly type: "where_value";
-  kind: "func_call";
-  func: FuncCall;
-} | WhereArith | WhereUnaryMinus | WhereJsonbOp | WherePgvectorOp | CaseExpr | CastExpr;
-interface WhereComparison {
-  readonly type: "where_comparison";
-  operator: ComparisonOperator;
-  left: WhereValue;
-  right: WhereValue;
-}
-interface ColumnRef {
-  readonly type: "column_ref";
-  schema?: string;
-  table?: string;
-  name: string;
-}
-type FuncCallArg = {
-  kind: "wildcard";
-} | {
-  kind: "args";
-  distinct: boolean;
-  args: WhereValue[];
-};
-interface FuncCall {
-  readonly type: "func_call";
-  name: string;
-  args: FuncCallArg;
-}
-interface ColumnExpr {
-  readonly type: "column_expr";
-  kind: "wildcard" | "qualified_wildcard" | "expr";
-  table?: string;
-  expr?: WhereValue;
-}
-interface Column {
-  readonly type: "column";
-  expr: ColumnExpr;
-  alias?: Alias;
-}
-interface Alias {
-  readonly type: "alias";
-  name: string;
-}
-interface TableRef {
-  readonly type: "table_ref";
+import { i as SelectStatement, n as defineSchema, r as Result, t as Schema } from "./joins-Cu_0yAgN.mjs";
+
+//#region src/guard.d.ts
+type GuardVal = string | number;
+interface GuardCol {
   schema?: string;
-  name: string;
+  table: string;
+  column: string;
 }
-//#endregion
-//#region src/result.d.ts
-type Failure<E = unknown> = {
-  ok: false;
-  error: E;
-  unwrap(): never;
-};
-type Success<T = void> = {
-  ok: true;
-  data: T;
-  unwrap(): T;
+type WhereGuard = GuardCol & {
+  value: GuardVal;
 };
-type Result<T = void, E = Error> = Failure<E> | Success<T>;
-//#endregion
-//#region src/graph.d.ts
-type TableDefs = ReturnType<typeof defineTables>;
-declare function defineTables<T extends { [Table in keyof T]: Record<string, null | { [FK in keyof T & string]: {
-  ft: FK;
-  fc: keyof T[FK] & string;
-} }[keyof T & string]> }>(tables: T): T;
+type OneOrTwoDots<S extends string> = S extends `${infer A}.${infer B}.${infer C}` ? A extends `${string}.${string}` ? never : B extends `${string}.${string}` ? never : C extends `${string}.${string}` ? never : S : S extends `${infer A}.${infer B}` ? A extends `${string}.${string}` ? never : B extends `${string}.${string}` ? never : S : never;
+type SchemaGuardKeys<T> = { [Table in keyof T & string]: `${Table}.${keyof T[Table] & string}` }[keyof T & string];
+declare function applyGuards(ast: SelectStatement, guards: WhereGuard[], limit?: number): Result<SelectStatement>;
 //#endregion
 //#region src/output.d.ts
 declare function outputSql(ast: SelectStatement): string;
@@ -269,54 +20,21 @@ declare function outputSql(ast: SelectStatement): string;
 //#region src/parse.d.ts
 declare function parseSql(expr: string): Result<SelectStatement>;
 //#endregion
-//#region src/sanitise.d.ts
-interface GuardCol {
-  schema?: string;
-  table: string;
-  col: string;
-}
-interface WhereGuard {
-  schema?: string;
-  table: string;
-  col: string;
-  value: string | number;
-}
-declare function sanitiseSql(ast: SelectStatement, guard: WhereGuard): Result<SelectStatement>;
-//#endregion
 //#region src/index.d.ts
-declare function sanitise(sql: string, {
-  tables,
-  where
-}: {
-  tables: TableDefs;
-  where: WhereGuard;
+declare function agentSql<S extends string>(sql: string, column: S & OneOrTwoDots<S>, value: GuardVal, {
+  schema,
+  limit
+}?: {
+  schema?: Schema;
+  limit?: number;
 }): string;
-declare function safeSanitise(sql: string, {
-  tables,
-  where
-}: {
-  tables: TableDefs;
-  where: WhereGuard;
-}): Result<string>;
-declare function sanitiserFactory(_: {
-  tables: TableDefs;
-  where: WhereGuard;
+declare function createAgentSql<T extends Schema, S extends SchemaGuardKeys<T>>(schema: T, guards: Record<S, GuardVal>, opts: {
+  limit?: number;
   throws: false;
 }): (expr: string) => Result<string>;
-declare function sanitiserFactory(_: {
-  tables: TableDefs;
-  where: WhereGuard;
+declare function createAgentSql<T extends Schema, S extends SchemaGuardKeys<T>>(schema: T, guards: Record<S, GuardVal>, opts?: {
+  limit?: number;
   throws?: true;
 }): (expr: string) => string;
-declare function makeSanitiserFactory(_: {
-  tables: TableDefs;
-  guardCol: GuardCol;
-  throws: false;
-}): (guardVal: string | number) => (expr: string) => Result<string>;
-declare function makeSanitiserFactory(_: {
-  tables: TableDefs;
-  guardCol: GuardCol;
-  throws?: true;
-}): (guardVal: string | number) => (expr: string) => string;
 //#endregion
-export { defineTables, makeSanitiserFactory, outputSql, parseSql, safeSanitise, sanitise, sanitiseSql, sanitiserFactory };
+export { agentSql, createAgentSql, defineSchema, outputSql, parseSql, applyGuards as sanitiseSql };

package/dist/index.mjs

@@ -6,65 +6,9 @@ var ParseError = class extends Error {
 var SanitiseError = class extends Error {
 	type = "sanitise_error";
 };
-//#endregion
-//#region src/result.ts
-function Err(error) {
-	return {
-		ok: false,
-		error,
-		unwrap() {
-			throw new Error(String(error));
-		}
-	};
-}
-function Ok(data) {
-	return {
-		ok: true,
-		data,
-		unwrap() {
-			return data;
-		}
-	};
-}
-function returnOrThrow(result, throws) {
-	if (!throws) return result;
-	if (result.ok) return result.data;
-	throw result.error;
-}
-//#endregion
-//#region src/graph.ts
-function defineTables(tables) {
-	return tables;
-}
-function checkJoins(ast, tables) {
-	if (!(ast.from.table.name in tables)) return Err(new SanitiseError(`Table ${ast.from.table.name} is not allowed`));
-	for (const join of ast.joins) {
-		const joinSettings = tables[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 = tables[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
-	};
-}
+var AgentSqlError = class extends Error {
+	type = "agent_sql_error";
+};
 //#endregion
 //#region src/utils.ts
 function unreachable(x) {
@@ -276,6 +220,176 @@ function handleColumnExpr(node) {
 	}
 }
 //#endregion
+//#region src/result.ts
+function Err(error) {
+	return {
+		ok: false,
+		error,
+		unwrap() {
+			throw new Error(String(error));
+		}
+	};
+}
+function Ok(data) {
+	return {
+		ok: true,
+		data,
+		unwrap() {
+			return data;
+		}
+	};
+}
+function returnOrThrow(result, throws) {
+	if (!throws) return result;
+	if (result.ok) return result.data;
+	throw result.error;
+}
+//#endregion
+//#region src/guard.ts
+const DEFAULT_LIMIT = 1e4;
+function applyGuards(ast, guards, limit = DEFAULT_LIMIT) {
+	const ast2 = addWhereGuard(ast, guards);
+	if (!ast2.ok) return ast2;
+	return addLimitGuard(ast2.data, limit);
+}
+function resolveGuards(guards) {
+	if (guards.length === 0) return Err(new AgentSqlError("At least one guard must be provided."));
+	const result = [];
+	for (const [column, value] of Object.entries(guards)) {
+		const guardCol = resolveSingleGuardCol(column);
+		if (!guardCol.ok) return guardCol;
+		result.push({
+			...guardCol.data,
+			value
+		});
+	}
+	return Ok(result);
+}
+function resolveSingleGuardCol(column) {
+	const [a, b, c] = column.split(".");
+	if (a === void 0 || b === void 0) return Err(new AgentSqlError(`Malformed column string: '${column}'. Pass 'table.column'.`));
+	if (c === void 0) return Ok({
+		table: a,
+		column: b
+	});
+	return Err(new AgentSqlError("Specifying guard as schema.table.name not yet supported"));
+}
+function addWhereGuard(ast, guards) {
+	for (const guard of guards) {
+		const tableRef = {
+			type: "table_ref",
+			schema: guard.schema,
+			name: guard.table
+		};
+		if (!checkIfTableRefExists(ast, tableRef)) return Err(new SanitiseError(`The table '${handleTableRef(tableRef)}' must appear in the FROM or JOIN clauses.`));
+	}
+	const [first, ...rest] = guards.map((guard) => {
+		const { schema, table, column, value } = guard;
+		return {
+			type: "where_comparison",
+			operator: "=",
+			left: {
+				type: "where_value",
+				kind: "column_ref",
+				ref: {
+					type: "column_ref",
+					schema,
+					table,
+					name: column
+				}
+			},
+			right: typeof value === "string" ? {
+				type: "where_value",
+				kind: "string",
+				value
+			} : {
+				type: "where_value",
+				kind: "integer",
+				value
+			}
+		};
+	});
+	if (!first) return Err(new AgentSqlError("No guards were provided"));
+	const combined = rest.reduce((acc, clause) => ({
+		type: "where_and",
+		left: acc,
+		right: clause
+	}), first);
+	return Ok({
+		...ast,
+		where: ast.where ? {
+			type: "where_root",
+			inner: {
+				type: "where_and",
+				left: combined,
+				right: ast.where.inner
+			}
+		} : {
+			type: "where_root",
+			inner: combined
+		}
+	});
+}
+function addLimitGuard(ast, limit) {
+	const limitClause = {
+		type: "limit",
+		value: limit
+	};
+	if (ast.limit === null) return Ok({
+		...ast,
+		limit: limitClause
+	});
+	if (ast.limit.value > limit) return Ok({
+		...ast,
+		limit: limitClause
+	});
+	return Ok(ast);
+}
+function checkIfTableRefExists(ast, tableRef) {
+	return tableEquals(ast.from.table, tableRef) || ast.joins.some((join) => tableEquals(join.table, tableRef));
+}
+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",
@@ -6333,114 +6447,40 @@ function parseSql(expr) {
 	}
 }
 //#endregion
-//#region src/sanitise.ts
-function sanitiseSql(ast, guard) {
-	const { schema, table, col, value } = guard;
-	const tableRef = {
-		type: "table_ref",
+//#region src/index.ts
+function agentSql(sql, column, value, { schema, limit } = {}) {
+	return privateAgentSql(sql, {
+		guards: { [column]: value },
 		schema,
-		name: table
-	};
-	if (!checkIfTableRefExists(ast, tableRef)) return Err(new SanitiseError(`The table '${handleTableRef(tableRef)}' must appear in the FROM or JOIN clauses.`));
-	const newClause = {
-		type: "where_comparison",
-		operator: "=",
-		left: {
-			type: "where_value",
-			kind: "column_ref",
-			ref: {
-				type: "column_ref",
-				schema,
-				table,
-				name: col
-			}
-		},
-		right: typeof value === "string" ? {
-			type: "where_value",
-			kind: "string",
-			value
-		} : {
-			type: "where_value",
-			kind: "integer",
-			value
-		}
-	};
-	return Ok({
-		...ast,
-		where: ast.where ? {
-			type: "where_root",
-			inner: {
-				type: "where_and",
-				left: newClause,
-				right: ast.where.inner
-			}
-		} : {
-			type: "where_root",
-			inner: newClause
-		}
+		limit,
+		throws: true
 	});
 }
-function checkIfTableRefExists(ast, tableRef) {
-	return tableEquals(ast.from.table, tableRef) || ast.joins.some((join) => tableEquals(join.table, tableRef));
-}
-function tableEquals(a, b) {
-	return a.schema == b.schema && a.name == b.name;
+function createAgentSql(schema, guards, { limit, throws = true } = {}) {
+	return (expr) => throws ? privateAgentSql(expr, {
+		guards,
+		schema,
+		limit,
+		throws
+	}) : privateAgentSql(expr, {
+		guards,
+		schema,
+		limit,
+		throws
+	});
 }
-//#endregion
-//#region src/index.ts
-function privateSanitise(sql, { tables, where, throws }) {
+function privateAgentSql(sql, { guards: guardsRaw, schema, limit, throws }) {
+	const guards = resolveGuards(guardsRaw);
+	if (!guards.ok) throw guards.error;
 	const ast = parseSql(sql);
 	if (!ast.ok) return returnOrThrow(ast, throws);
-	const ast2 = checkJoins(ast.data, tables);
+	const ast2 = checkJoins(ast.data, schema);
 	if (!ast2.ok) return returnOrThrow(ast2, throws);
-	const san = sanitiseSql(ast2.data, where);
+	const san = applyGuards(ast2.data, guards.data, limit);
 	if (!san.ok) return returnOrThrow(san, throws);
 	const res = outputSql(san.data);
 	if (throws) return res;
 	return Ok(res);
 }
-function sanitise(sql, { tables, where }) {
-	return privateSanitise(sql, {
-		tables,
-		where,
-		throws: true
-	});
-}
-function safeSanitise(sql, { tables, where }) {
-	return privateSanitise(sql, {
-		tables,
-		where,
-		throws: false
-	});
-}
-function sanitiserFactory({ tables, where, throws = true }) {
-	return (expr) => throws ? privateSanitise(expr, {
-		tables,
-		where,
-		throws
-	}) : privateSanitise(expr, {
-		tables,
-		where,
-		throws
-	});
-}
-function makeSanitiserFactory({ tables, guardCol, throws = true }) {
-	function factory(guardVal) {
-		const where = {
-			...guardCol,
-			value: guardVal
-		};
-		return throws ? sanitiserFactory({
-			tables,
-			where,
-			throws
-		}) : sanitiserFactory({
-			tables,
-			where,
-			throws
-		});
-	}
-	return factory;
-}
 //#endregion
-export { defineTables, makeSanitiserFactory, outputSql, parseSql, safeSanitise, sanitise, sanitiseSql, sanitiserFactory };
+export { agentSql, createAgentSql, defineSchema, outputSql, parseSql, applyGuards as sanitiseSql };

package/dist/joins-Cu_0yAgN.d.mts

@@ -0,0 +1,266 @@
+//#region src/ast.d.ts
+interface SelectStatement {
+  readonly type: "select";
+  distinct: Distinct | DistinctOn | null;
+  columns: Column[];
+  from: SelectFrom;
+  joins: JoinClause[];
+  where: WhereRoot | null;
+  groupBy: GroupByClause | null;
+  having: HavingClause | null;
+  orderBy: OrderByClause | null;
+  limit: LimitClause | null;
+  offset: OffsetClause | null;
+}
+interface Distinct {
+  readonly type: "distinct";
+}
+/** PostgreSQL: DISTINCT ON (col1, col2, ...) */
+interface DistinctOn {
+  readonly type: "distinct_on";
+  columns: WhereValue[];
+}
+interface GroupByClause {
+  readonly type: "group_by";
+  items: WhereValue[];
+}
+interface HavingClause {
+  readonly type: "having";
+  expr: WhereExpr;
+}
+interface OrderByClause {
+  readonly type: "order_by";
+  items: OrderByItem[];
+}
+type SortDirection = "asc" | "desc";
+type NullsOrder = "nulls_first" | "nulls_last";
+interface OrderByItem {
+  readonly type: "order_by_item";
+  expr: WhereValue;
+  direction?: SortDirection;
+  nulls?: NullsOrder;
+}
+interface OffsetClause {
+  readonly type: "offset";
+  value: number;
+}
+type JoinType = "inner" | "inner_outer" | "left" | "left_outer" | "right" | "right_outer" | "full" | "full_outer" | "cross" | "natural";
+type JoinCondition = {
+  readonly type: "join_on";
+  expr: WhereExpr;
+} | {
+  readonly type: "join_using";
+  columns: string[];
+};
+interface JoinClause {
+  readonly type: "join";
+  joinType: JoinType;
+  table: TableRef;
+  condition: JoinCondition | null;
+}
+type SelectFrom = {
+  readonly type: "select_from";
+  table: TableRef;
+};
+type LimitClause = {
+  readonly type: "limit";
+  value: number;
+};
+type WhereRoot = {
+  readonly type: "where_root";
+  inner: WhereExpr;
+};
+type WhereExpr = WhereAnd | WhereOr | WhereNot | WhereComparison | WhereIsNull | WhereIsBool | WhereBetween | WhereIn | WhereLike | WhereTsMatch;
+interface WhereAnd {
+  readonly type: "where_and";
+  left: WhereExpr;
+  right: WhereExpr;
+}
+interface WhereOr {
+  readonly type: "where_or";
+  left: WhereExpr;
+  right: WhereExpr;
+}
+type ComparisonOperator = "=" | "<>" | "!=" | "<" | ">" | "<=" | ">=";
+interface WhereNot {
+  readonly type: "where_not";
+  expr: WhereExpr;
+}
+interface WhereIsNull {
+  readonly type: "where_is_null";
+  not: boolean;
+  expr: WhereValue;
+}
+type IsBoolTarget = boolean | "unknown";
+interface WhereIsBool {
+  readonly type: "where_is_bool";
+  not: boolean;
+  expr: WhereValue;
+  target: IsBoolTarget;
+}
+interface WhereBetween {
+  readonly type: "where_between";
+  not: boolean;
+  expr: WhereValue;
+  low: WhereValue;
+  high: WhereValue;
+}
+interface WhereIn {
+  readonly type: "where_in";
+  not: boolean;
+  expr: WhereValue;
+  list: WhereValue[];
+}
+/** "ilike" is PostgreSQL-specific (case-insensitive LIKE) */
+type LikeOp = "like" | "ilike";
+interface WhereLike {
+  readonly type: "where_like";
+  not: boolean;
+  op: LikeOp;
+  expr: WhereValue;
+  pattern: WhereValue;
+}
+/** PostgreSQL: JSONB operators */
+type JsonbOp = "->" | "->>" | "#>" | "#>>" | "?" | "?|" | "?&" | "@>";
+/** PostgreSQL: JSONB binary operator expression */
+interface WhereJsonbOp {
+  readonly type: "where_jsonb_op";
+  op: JsonbOp;
+  left: WhereValue;
+  right: WhereValue;
+}
+/** PostgreSQL: text search match (@@) */
+interface WhereTsMatch {
+  readonly type: "where_ts_match";
+  left: WhereValue;
+  right: WhereValue;
+}
+/** pgvector: distance operators */
+type PgvectorOp = "<->" | "<#>" | "<=>" | "<+>" | "<~>" | "<%>";
+/** pgvector: distance operator expression */
+interface WherePgvectorOp {
+  readonly type: "where_pgvector_op";
+  op: PgvectorOp;
+  left: WhereValue;
+  right: WhereValue;
+}
+type ArithOp = "+" | "-" | "*" | "/" | "%" | "||";
+interface WhereArith {
+  readonly type: "where_arith";
+  op: ArithOp;
+  left: WhereValue;
+  right: WhereValue;
+}
+interface WhereUnaryMinus {
+  readonly type: "where_unary_minus";
+  expr: WhereValue;
+}
+interface CaseWhen {
+  condition: WhereValue;
+  result: WhereValue;
+}
+interface CaseExpr {
+  readonly type: "case_expr";
+  subject: WhereValue | null;
+  whens: CaseWhen[];
+  else: WhereValue | null;
+}
+interface CastExpr {
+  readonly type: "cast_expr";
+  expr: WhereValue;
+  typeName: string;
+}
+type WhereValue = {
+  readonly type: "where_value";
+  kind: "string";
+  value: string;
+} | {
+  readonly type: "where_value";
+  kind: "integer";
+  value: number;
+} | {
+  readonly type: "where_value";
+  kind: "float";
+  value: number;
+} | {
+  readonly type: "where_value";
+  kind: "bool";
+  value: boolean;
+} | {
+  readonly type: "where_value";
+  kind: "null";
+} | {
+  readonly type: "where_value";
+  kind: "column_ref";
+  ref: ColumnRef;
+} | {
+  readonly type: "where_value";
+  kind: "func_call";
+  func: FuncCall;
+} | WhereArith | WhereUnaryMinus | WhereJsonbOp | WherePgvectorOp | CaseExpr | CastExpr;
+interface WhereComparison {
+  readonly type: "where_comparison";
+  operator: ComparisonOperator;
+  left: WhereValue;
+  right: WhereValue;
+}
+interface ColumnRef {
+  readonly type: "column_ref";
+  schema?: string;
+  table?: string;
+  name: string;
+}
+type FuncCallArg = {
+  kind: "wildcard";
+} | {
+  kind: "args";
+  distinct: boolean;
+  args: WhereValue[];
+};
+interface FuncCall {
+  readonly type: "func_call";
+  name: string;
+  args: FuncCallArg;
+}
+interface ColumnExpr {
+  readonly type: "column_expr";
+  kind: "wildcard" | "qualified_wildcard" | "expr";
+  table?: string;
+  expr?: WhereValue;
+}
+interface Column {
+  readonly type: "column";
+  expr: ColumnExpr;
+  alias?: Alias;
+}
+interface Alias {
+  readonly type: "alias";
+  name: string;
+}
+interface TableRef {
+  readonly type: "table_ref";
+  schema?: string;
+  name: string;
+}
+//#endregion
+//#region src/result.d.ts
+type Failure<E = unknown> = {
+  ok: false;
+  error: E;
+  unwrap(): never;
+};
+type Success<T = void> = {
+  ok: true;
+  data: T;
+  unwrap(): T;
+};
+type Result<T = void, E = Error> = Failure<E> | Success<T>;
+//#endregion
+//#region src/joins.d.ts
+type Schema = ReturnType<typeof defineSchema>;
+declare function defineSchema<T extends { [Table in keyof T]: Record<string, null | { [FK in keyof T & string]: {
+  ft: FK;
+  fc: keyof T[FK] & string;
+} }[keyof T & string]> }>(schema: T): T;
+//#endregion
+export { SelectStatement as i, defineSchema as n, Result as r, Schema as t };

package/package.json

@@ -1,7 +1,13 @@
 {
   "name": "agent-sql",
-  "version": "0.1.2",
+  "version": "0.2.1",
   "description": "A starter for creating a TypeScript package.",
+  "keywords": [
+    "agent",
+    "ai",
+    "postgres",
+    "sql"
+  ],
   "homepage": "https://github.com/carderne/agent-sql#readme",
   "bugs": {
     "url": "https://github.com/carderne/agent-sql/issues"
@@ -18,6 +24,7 @@
   "type": "module",
   "exports": {
     ".": "./dist/index.mjs",
+    "./drizzle": "./dist/drizzle.mjs",
     "./package.json": "./package.json"
   },
   "publishConfig": {
@@ -42,11 +49,20 @@
     "@types/pg": "^8.18.0",
     "@typescript/native-preview": "7.0.0-dev.20260316.1",
     "bumpp": "^11.0.1",
+    "drizzle-orm": "^0.45.1",
     "pg": "^8.20.0",
     "tsx": "^4.21.0",
     "typescript": "^5.9.3",
     "vite-plus": "^0.1.11"
   },
+  "peerDependencies": {
+    "drizzle-orm": ">=0.45"
+  },
+  "peerDependenciesMeta": {
+    "drizzle-orm": {
+      "optional": true
+    }
+  },
   "packageManager": "pnpm@10.32.1",
   "pnpm": {
     "overrides": {