agent-sql 0.1.2

package/README.md

@@ -0,0 +1,99 @@
+# agent-sql
+
+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.
+
+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.
+
+Apparently this is how [Trigger.dev does it](https://x.com/mattaitken/status/2033928542975639785). And [Cloudflare](https://x.com/thomas_ankcorn/status/2033931057133748330).
+
+## Quickstart
+
+```bash
+npm install agent-sql
+```
+
+```ts
+import { sanitise } 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" },
+});
+
+console.log(sql);
+// SELECT id, name
+// FROM users
+// WHERE (users.tenant_id = 'acme' AND status = 'active')
+// LIMIT 10
+```
+
+Or, more usefully:
+
+```ts
+import { sanitiserFactory } from "agent-sql";
+import { tool } from "ai";
+import { sql } from "drizzle-orm";
+import { db } from "@/db";
+
+function makeSqlTool(orgId: string) {
+  // Create a sanitiser function for this tenant
+  const sanitise = sanitiserFactory({
+    tables: {},
+    where: { table: "org", col: "id", value: orgId },
+  });
+
+  return tool({
+    description: "Run raw SQL against the DB",
+    inputSchema: z.object({ query: z.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);
+      // 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));
+    },
+  });
+}
+```
+
+`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.
+
+## Development
+
+First install [Vite+](https://viteplus.dev/guide/):
+
+```bash
+curl -fsSL https://vite.plus | bash
+```
+
+Install dependencies:
+
+```bash
+vp install
+```
+
+Format, lint, typecheck:
+
+```bash
+vp check --fix
+```
+
+Run the unit tests:
+
+```bash
+vp test
+```
+
+Build the library:
+
+```bash
+vp pack
+```

package/dist/index.d.mts

@@ -0,0 +1,322 @@
+//#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/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;
+//#endregion
+//#region src/output.d.ts
+declare function outputSql(ast: SelectStatement): string;
+//#endregion
+//#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;
+}): string;
+declare function safeSanitise(sql: string, {
+  tables,
+  where
+}: {
+  tables: TableDefs;
+  where: WhereGuard;
+}): Result<string>;
+declare function sanitiserFactory(_: {
+  tables: TableDefs;
+  where: WhereGuard;
+  throws: false;
+}): (expr: string) => Result<string>;
+declare function sanitiserFactory(_: {
+  tables: TableDefs;
+  where: WhereGuard;
+  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 };