agent-sql 0.1.2 → 0.2.0

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,33 +24,47 @@ 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 },
+  msg: { userId: { user: "id" } },
+});
+
+function makeSqlTool(userId: string) {
   // Create a sanitiser function for this tenant
-  const sanitise = sanitiserFactory({
-    tables: {},
-    where: { table: "org", col: "id", value: orgId },
+  const agentSql = createAgentSql({
+    column: "user.id",
+    value: userId,
+    schema,
   });
 
   return tool({
@@ -55,7 +73,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 +82,33 @@ 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({
+  column: "user.id",
+  value: userId,
+  schema,
+});
+```
+
+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,19 @@
-//#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 WhereGuard = GuardCol & {
+  value: GuardVal;
 };
-type Success<T = void> = {
-  ok: true;
-  data: T;
-  unwrap(): T;
-};
-type Result<T = void, E = Error> = Failure<E> | Success<T>;
+declare function addGuards(ast: SelectStatement, guard: WhereGuard, limit?: number): Result<SelectStatement>;
 //#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;
+//#region src/namespec.d.ts
+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;
 //#endregion
 //#region src/output.d.ts
 declare function outputSql(ast: SelectStatement): string;
@@ -269,54 +21,41 @@ 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<S extends string>(_: {
+  column: S & OneOrTwoDots<S>;
+  value: GuardVal;
+  schema?: Schema;
+  limit?: number;
   throws: false;
 }): (expr: string) => Result<string>;
-declare function sanitiserFactory(_: {
-  tables: TableDefs;
-  where: WhereGuard;
+declare function createAgentSql<S extends string>(_: {
+  column: S & OneOrTwoDots<S>;
+  value: GuardVal;
+  schema?: Schema;
+  limit?: number;
   throws?: true;
 }): (expr: string) => string;
-declare function makeSanitiserFactory(_: {
-  tables: TableDefs;
-  guardCol: GuardCol;
+declare function createAgentSql<S extends string>(_: {
+  column: S & OneOrTwoDots<S>;
+  value?: undefined;
+  schema?: Schema;
+  limit?: number;
   throws: false;
-}): (guardVal: string | number) => (expr: string) => Result<string>;
-declare function makeSanitiserFactory(_: {
-  tables: TableDefs;
-  guardCol: GuardCol;
+}): (guardVal: GuardVal) => (expr: string) => Result<string>;
+declare function createAgentSql<S extends string>(_: {
+  column: S & OneOrTwoDots<S>;
+  value?: undefined;
+  schema?: Schema;
+  limit?: number;
   throws?: true;
-}): (guardVal: string | number) => (expr: string) => string;
+}): (guardVal: GuardVal) => (expr: string) => string;
 //#endregion
-export { defineTables, makeSanitiserFactory, outputSql, parseSql, safeSanitise, sanitise, sanitiseSql, sanitiserFactory };
+export { agentSql, createAgentSql, defineSchema, outputSql, parseSql, addGuards 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,159 @@ 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 addGuards(ast, guard, limit = DEFAULT_LIMIT) {
+	const ast2 = addWhereGuard(ast, guard);
+	if (!ast2.ok) return ast2;
+	return addLimitGuard(ast2.data, limit);
+}
+function addWhereGuard(ast, guard) {
+	const { schema, table, column, value } = guard;
+	const tableRef = {
+		type: "table_ref",
+		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: column
+			}
+		},
+		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
+		}
+	});
+}
+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/namespec.ts
+function getQualifiedColumnFromString(column) {
+	const [a, b, c] = column.split(".");
+	if (a === void 0 || b === void 0) throw new AgentSqlError(`Malformed column string: '${column}'. Pass 'table.column'.`);
+	if (c === void 0) return Ok({
+		table: a,
+		column: b
+	});
+	return Ok({
+		schema: a,
+		table: b,
+		column: c
+	});
+}
+//#endregion
 //#region src/sql.ohm-bundle.js
 const result = makeRecipe([
 	"grammar",
@@ -6333,114 +6430,63 @@ function parseSql(expr) {
 	}
 }
 //#endregion
-//#region src/sanitise.ts
-function sanitiseSql(ast, guard) {
-	const { schema, table, col, value } = guard;
-	const tableRef = {
-		type: "table_ref",
-		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
-		}
-	});
-}
-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/index.ts
-function privateSanitise(sql, { tables, where, throws }) {
-	const ast = parseSql(sql);
-	if (!ast.ok) return returnOrThrow(ast, throws);
-	const ast2 = checkJoins(ast.data, tables);
-	if (!ast2.ok) return returnOrThrow(ast2, throws);
-	const san = sanitiseSql(ast2.data, where);
-	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,
+function agentSql(sql, column, value, { schema, limit } = {}) {
+	return privateAgentSql(sql, {
+		column,
+		value,
+		schema,
+		limit,
 		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,
+function createAgentSql({ column, schema, value, limit, throws = true }) {
+	if (value !== void 0) return (expr) => throws ? privateAgentSql(expr, {
+		column,
+		value,
+		schema,
+		limit,
 		throws
-	}) : privateSanitise(expr, {
-		tables,
-		where,
+	}) : privateAgentSql(expr, {
+		column,
+		value,
+		schema,
+		limit,
 		throws
 	});
-}
-function makeSanitiserFactory({ tables, guardCol, throws = true }) {
 	function factory(guardVal) {
-		const where = {
-			...guardCol,
-			value: guardVal
-		};
-		return throws ? sanitiserFactory({
-			tables,
-			where,
+		return throws ? createAgentSql({
+			column,
+			schema,
+			value: guardVal,
+			limit,
 			throws
-		}) : sanitiserFactory({
-			tables,
-			where,
+		}) : createAgentSql({
+			column,
+			schema,
+			value: guardVal,
+			limit,
 			throws
 		});
 	}
 	return factory;
 }
+function privateAgentSql(sql, { column, value, schema, limit, throws }) {
+	const guardCol = getQualifiedColumnFromString(column);
+	if (!guardCol.ok) throw guardCol.error;
+	const ast = parseSql(sql);
+	if (!ast.ok) return returnOrThrow(ast, throws);
+	const ast2 = checkJoins(ast.data, schema);
+	if (!ast2.ok) return returnOrThrow(ast2, throws);
+	const where = {
+		...guardCol.data,
+		value
+	};
+	const san = addGuards(ast2.data, where, limit);
+	if (!san.ok) return returnOrThrow(san, throws);
+	const res = outputSql(san.data);
+	if (throws) return res;
+	return Ok(res);
+}
 //#endregion
-export { defineTables, makeSanitiserFactory, outputSql, parseSql, safeSanitise, sanitise, sanitiseSql, sanitiserFactory };
+export { agentSql, createAgentSql, defineSchema, outputSql, parseSql, addGuards 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,6 +1,6 @@
 {
   "name": "agent-sql",
-  "version": "0.1.2",
+  "version": "0.2.0",
   "description": "A starter for creating a TypeScript package.",
   "homepage": "https://github.com/carderne/agent-sql#readme",
   "bugs": {
@@ -18,6 +18,7 @@
   "type": "module",
   "exports": {
     ".": "./dist/index.mjs",
+    "./drizzle": "./dist/drizzle.mjs",
     "./package.json": "./package.json"
   },
   "publishConfig": {
@@ -42,11 +43,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": {