diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml
index a598b1d..c733fa1 100644
--- a/.github/workflows/unit-test.yml
+++ b/.github/workflows/unit-test.yml
@@ -12,3 +12,4 @@ jobs:
           node-version: 16
       - run: npm i
       - run: npm test
+      - run: npm test:types
diff --git a/package.json b/package.json
index bc001b6..cf30d6c 100644
--- a/package.json
+++ b/package.json
@@ -8,6 +8,7 @@
     "build": "tsc",
     "test": "rome ci src test",
     "lint:fix": "rome check --apply src test",
+    "test:types": "tsc --noEmit",
     "test:integration": "jest --runInBand test/integration",
     "test:compose:integration": "docker compose -f docker-compose.yml --profile with-sut up db sut --exit-code-from sut",
     "setup": "prisma generate && prisma migrate dev",
@@ -30,8 +31,9 @@
     "uuid": "^9.0.0"
   },
   "dependencies": {
+    "lodash": "^4.17.21",
     "node-sql-parser": "^4.12.0",
-    "lodash": "^4.17.21"
+    "type-fest": "^4.10.3"
   },
   "peerDependencies": {
     "@prisma/client": "^5.0.0",
diff --git a/src/expressions.ts b/src/expressions.ts
index e9e589b..367af3e 100644
--- a/src/expressions.ts
+++ b/src/expressions.ts
@@ -1,7 +1,8 @@
-import { PrismaClient } from "@prisma/client";
+import { Prisma, PrismaClient } from "@prisma/client";
 import random from "lodash/random";
 import matches from "lodash/matches";
 import { Parser } from "node-sql-parser";
+import { AsyncReturnType } from "type-fest";
 import { escapeLiteral } from "./escape";
 import { defineDmmfProperty } from "@prisma/client/runtime/library";
 import { jsonb_array_elements_text } from "./ast-fragments";
@@ -31,16 +32,23 @@ type Token = {
 };
 type Tokens = Record<string, Token>;
 
-export type Expression<ContextKeys extends string = string> =
+type FFMeta<M extends Prisma.ModelName> = PrismaClient[Uncapitalize<M>]["findFirst"];
+type ModelWhereArgs<M extends Prisma.ModelName> = Exclude<Parameters<FFMeta<M>>["0"], undefined>["where"];
+type ModelResult<M extends Prisma.ModelName> = AsyncReturnType<FFMeta<M>>;
+type NonNullableModelResult<M extends Prisma.ModelName> = Exclude<ModelResult<M>, null>;
+
+// The expression below explicitly excludes returning a client query for the model the expression is for, as this can create infinite loops as the access logic recurses
+export type Expression<ContextKeys extends string, M extends Prisma.ModelName> =
 	| string
 	| ((
 			client: PrismaClient,
 			// Explicitly return any, so that the prisma client doesn't error
-			row: (col: string) => any,
+			row: <K extends keyof NonNullableModelResult<M>>(col: K) => NonNullableModelResult<M>[K],
+			// TODO infer the return type of the context function automatically
 			context: (key: ContextKeys) => string,
-	  ) => Promise<any> | { [col: string]: any });
+	  ) => Promise<ModelResult<Exclude<Prisma.ModelName, M>>> | ModelWhereArgs<M>);
 
-const expressionRowName = (col: string) => `___yates_row_${col}`;
+const expressionRowName = (col: any) => `___yates_row_${col}`;
 const expressionContext = (context: string) => `___yates_context_${context}`;
 // Generate a big 32bit signed integer to use as an ID
 const getLargeRandomInt = () => random(1000000000, 2147483647);
@@ -265,7 +273,10 @@ const tokenizeWhereExpression = (
 	};
 };
 
-export const expressionToSQL = async (getExpression: Expression, table: string): Promise<string> => {
+export const expressionToSQL = async <ContextKeys extends string, YModel extends Prisma.ModelName>(
+	getExpression: Expression<ContextKeys, YModel>,
+	table: string,
+): Promise<string> => {
 	if (typeof getExpression === "string") {
 		return getExpression;
 	}
@@ -308,12 +319,15 @@ export const expressionToSQL = async (getExpression: Expression, table: string):
 		async (resolve, reject) => {
 			const rawExpression = getExpression(
 				expressionClient as any as PrismaClient,
-				expressionRowName,
+				expressionRowName as any,
 				expressionContext,
 			);
 			// If the raw expression is a promise, then this is a client subselect,
 			// as opposed to a plain SQL expression or "where" object
-			const isSubselect = typeof rawExpression === "object" && typeof rawExpression.then === "function";
+			const isSubselect =
+				typeof rawExpression === "object" &&
+				"then" in rawExpression &&
+				typeof (rawExpression as Promise<any>).then === "function";
 
 			baseClient.$on("query", (e: any) => {
 				try {
diff --git a/src/index.ts b/src/index.ts
index ca94d12..097ded5 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -8,7 +8,7 @@ import { Expression, expressionToSQL, RuntimeDataModel } from "./expressions";
 
 const VALID_OPERATIONS = ["SELECT", "UPDATE", "INSERT", "DELETE"] as const;
 
-type Operation = typeof VALID_OPERATIONS[number];
+type Operation = (typeof VALID_OPERATIONS)[number];
 export type Models = Prisma.ModelName;
 
 interface ClientOptions {
@@ -18,18 +18,26 @@ interface ClientOptions {
 	txTimeout?: number;
 }
 
-export interface Ability<ContextKeys extends string = string> {
+export interface Ability<ContextKeys extends string, M extends Models> {
 	description?: string;
-	expression?: Expression<ContextKeys>;
+	expression?: Expression<ContextKeys, M>;
 	operation: Operation;
-	model?: Models;
+	model?: M;
 	slug?: string;
 }
+
+// This creates a union type of all possible abilities for a given set of models
+export type AllAbilities<ContextKeys extends string, YModels extends Models> = {
+	[model in YModels]: Ability<ContextKeys, model>;
+}[YModels];
+
 type CRUDOperations = "read" | "create" | "update" | "delete";
-export type DefaultAbilities = { [Model in Models]: { [op in CRUDOperations]: Ability } };
-export type CustomAbilities<ContextKeys extends string = string> = {
-	[model in Models]?: {
-		[op in string]?: Ability<ContextKeys>;
+export type DefaultAbilities<ContextKeys extends string = string, YModels extends Models = Models> = {
+	[Model in YModels]: { [op in CRUDOperations]: Ability<ContextKeys, Model> };
+};
+export type CustomAbilities<ContextKeys extends string = string, YModels extends Models = Models> = {
+	[model in YModels]?: {
+		[op in string]?: Ability<ContextKeys, model>;
 	};
 };
 
@@ -69,7 +77,11 @@ const hashWithPrefix = (prefix: string, abilityName: string) => {
 };
 
 // Sanitize a single string by ensuring the it has only lowercase alpha characters and underscores
-const sanitizeSlug = (slug: string) => slug.toLowerCase().replace("-", "_").replace(/[^a-z0-9_]/gi, "");
+const sanitizeSlug = (slug: string) =>
+	slug
+		.toLowerCase()
+		.replace("-", "_")
+		.replace(/[^a-z0-9_]/gi, "");
 
 export const createAbilityName = (model: string, ability: string) => {
 	return sanitizeSlug(hashWithPrefix("yates_ability_", `${model}_${ability}`));
@@ -184,12 +196,12 @@ export const createClient = (prisma: PrismaClient, getContext: GetContextFn, opt
 	return client;
 };
 
-const setRLS = async (
+const setRLS = async <ContextKeys extends string, YModel extends Models>(
 	prisma: PrismaClient,
 	table: string,
 	roleName: string,
 	operation: Operation,
-	rawExpression: Expression,
+	rawExpression: Expression<ContextKeys, YModel>,
 ) => {
 	let expression = await expressionToSQL(rawExpression, table);
 
@@ -223,7 +235,12 @@ const setRLS = async (
 	}
 };
 
-export const createRoles = async <K extends CustomAbilities = CustomAbilities, T = DefaultAbilities & K>({
+export const createRoles = async <
+	ContextKeys extends string,
+	YModels extends Models,
+	K extends CustomAbilities = CustomAbilities,
+	T = DefaultAbilities<ContextKeys, YModels> & K,
+>({
 	prisma,
 	customAbilities,
 	getRoles,
@@ -231,7 +248,7 @@ export const createRoles = async <K extends CustomAbilities = CustomAbilities, T
 	prisma: PrismaClient;
 	customAbilities?: Partial<K>;
 	getRoles: (abilities: T) => {
-		[key: string]: Ability[] | "*";
+		[role: string]: AllAbilities<ContextKeys, YModels>[] | "*";
 	};
 }) => {
 	const abilities: Partial<DefaultAbilities> = {};
@@ -254,28 +271,28 @@ export const createRoles = async <K extends CustomAbilities = CustomAbilities, T
 				description: `Create ${model}`,
 				expression: "true",
 				operation: "INSERT",
-				model,
+				model: model as any,
 				slug: "create",
 			},
 			read: {
 				description: `Read ${model}`,
 				expression: "true",
 				operation: "SELECT",
-				model,
+				model: model as any,
 				slug: "read",
 			},
 			update: {
 				description: `Update ${model}`,
 				expression: "true",
 				operation: "UPDATE",
-				model,
+				model: model as any,
 				slug: "update",
 			},
 			delete: {
 				description: `Delete ${model}`,
 				expression: "true",
 				operation: "DELETE",
-				model,
+				model: model as any,
 				slug: "delete",
 			},
 		};
@@ -286,7 +303,7 @@ export const createRoles = async <K extends CustomAbilities = CustomAbilities, T
 				abilities[model]![ability as CRUDOperations] = {
 					...customAbilities[model]![ability],
 					operation,
-					model,
+					model: model as any,
 					slug: ability,
 				};
 			}
@@ -334,12 +351,12 @@ export const createRoles = async <K extends CustomAbilities = CustomAbilities, T
 			]);
 
 			if (ability.expression) {
-				await setRLS(prisma, table, roleName, ability.operation, ability.expression);
+				await setRLS(prisma, table, roleName, ability.operation, ability.expression as any);
 			}
 		}
 	}
 
-	// For each of the Cortex roles, create a role in the database and grant it the relevant permissions.
+	// For each of the given roles, create a role in the database and grant it the relevant permissions.
 	// By defining each permission as a seperate role, we can GRANT them to the user role here, re-using them.
 	// It's not possible to dynamically GRANT these to a shared user role, as the GRANT is not isolated per transaction and leads to broken permissions.
 	for (const key in roles) {
@@ -403,7 +420,8 @@ export const createRoles = async <K extends CustomAbilities = CustomAbilities, T
 
 export interface SetupParams<
 	ContextKeys extends string = string,
-	K extends CustomAbilities<ContextKeys> = CustomAbilities<ContextKeys>,
+	YModels extends Models = Models,
+	K extends CustomAbilities<ContextKeys, YModels> = CustomAbilities<ContextKeys, YModels>,
 > {
 	/**
 	 * The Prisma client instance. Used for database queries and model introspection.
@@ -417,8 +435,8 @@ export interface SetupParams<
 	 * A function that returns the roles for your application.
 	 * This is paramaterised by the abilities, so you can use it to create roles that are a combination of abilities.
 	 */
-	getRoles: (abilities: DefaultAbilities & K) => {
-		[key: string]: Ability[] | "*";
+	getRoles: (abilities: DefaultAbilities<ContextKeys, YModels> & K) => {
+		[role: string]: AllAbilities<ContextKeys, YModels>[] | "*";
 	};
 	/**
 	 * A function that returns the context for the current request.
@@ -435,12 +453,13 @@ export interface SetupParams<
  **/
 export const setup = async <
 	ContextKeys extends string = string,
-	K extends CustomAbilities<ContextKeys> = CustomAbilities<ContextKeys>,
+	YModels extends Models = Models,
+	K extends CustomAbilities<ContextKeys, YModels> = CustomAbilities<ContextKeys, YModels>,
 >(
-	params: SetupParams<ContextKeys, K>,
+	params: SetupParams<ContextKeys, YModels, K>,
 ) => {
 	const { prisma, customAbilities, getRoles, getContext } = params;
-	await createRoles<K>({ prisma, customAbilities, getRoles });
+	await createRoles<ContextKeys, YModels, K>({ prisma, customAbilities, getRoles });
 	const client = createClient(prisma, getContext, params.options);
 
 	return client;
diff --git a/test/integration/expressions.spec.ts b/test/integration/expressions.spec.ts
index 1ecf323..f4728e0 100644
--- a/test/integration/expressions.spec.ts
+++ b/test/integration/expressions.spec.ts
@@ -174,7 +174,7 @@ describe("expressions", () => {
 							operation: "SELECT",
 							expression: (_client: PrismaClient, _row, context) => {
 								return {
-									value: context("item.value"),
+									value: context("item.value") as any as number,
 								};
 							},
 						},
@@ -337,7 +337,8 @@ describe("expressions", () => {
 								operation: "SELECT",
 								expression: () => {
 									return {
-										stock: "escape'--",
+										// We're intentionally using the wrong type to run this test
+										stock: "escape'--" as any as number,
 									};
 								},
 							},
@@ -369,7 +370,8 @@ describe("expressions", () => {
 								operation: "SELECT",
 								expression: (_client, row) => {
 									return {
-										name: row(`escape"--`),
+										// We're intentionally using the wrong type to run this test
+										name: row(`escape"--` as any),
 									};
 								},
 							},
@@ -1125,7 +1127,12 @@ describe("expressions", () => {
 				},
 				getRoles(abilities) {
 					return {
-						[role]: [abilities.Post.customCreateAbility, abilities.Post.read, abilities.Tag.read, abilities.Tag.create],
+						[role]: [
+							abilities.Post.customCreateAbility!,
+							abilities.Post.read,
+							abilities.Tag.read,
+							abilities.Tag.create,
+						],
 					};
 				},
 				getContext: () => ({
diff --git a/test/types/index.js b/test/types/index.js
new file mode 100644
index 0000000..c337288
--- /dev/null
+++ b/test/types/index.js
@@ -0,0 +1,92 @@
+"use strict";
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+var __generator = (this && this.__generator) || function (thisArg, body) {
+    var _ = { label: 0, sent: function() { if (t[0] & 1) throw t[1]; return t[1]; }, trys: [], ops: [] }, f, y, t, g;
+    return g = { next: verb(0), "throw": verb(1), "return": verb(2) }, typeof Symbol === "function" && (g[Symbol.iterator] = function() { return this; }), g;
+    function verb(n) { return function (v) { return step([n, v]); }; }
+    function step(op) {
+        if (f) throw new TypeError("Generator is already executing.");
+        while (g && (g = 0, op[0] && (_ = 0)), _) try {
+            if (f = 1, y && (t = op[0] & 2 ? y["return"] : op[0] ? y["throw"] || ((t = y["return"]) && t.call(y), 0) : y.next) && !(t = t.call(y, op[1])).done) return t;
+            if (y = 0, t) op = [op[0] & 2, t.value];
+            switch (op[0]) {
+                case 0: case 1: t = op; break;
+                case 4: _.label++; return { value: op[1], done: false };
+                case 5: _.label++; y = op[1]; op = [0]; continue;
+                case 7: op = _.ops.pop(); _.trys.pop(); continue;
+                default:
+                    if (!(t = _.trys, t = t.length > 0 && t[t.length - 1]) && (op[0] === 6 || op[0] === 2)) { _ = 0; continue; }
+                    if (op[0] === 3 && (!t || (op[1] > t[0] && op[1] < t[3]))) { _.label = op[1]; break; }
+                    if (op[0] === 6 && _.label < t[1]) { _.label = t[1]; t = op; break; }
+                    if (t && _.label < t[2]) { _.label = t[2]; _.ops.push(op); break; }
+                    if (t[2]) _.ops.pop();
+                    _.trys.pop(); continue;
+            }
+            op = body.call(thisArg, _);
+        } catch (e) { op = [6, e]; y = 0; } finally { f = t = 0; }
+        if (op[0] & 5) throw op[1]; return { value: op[0] ? op[1] : void 0, done: true };
+    }
+};
+exports.__esModule = true;
+var client_1 = require("@prisma/client");
+var src_1 = require("../../src");
+// it should error if an invalid model is used
+var run = function () { return __awaiter(void 0, void 0, void 0, function () {
+    return __generator(this, function (_a) {
+        (0, src_1.setup)({
+            prisma: new client_1.PrismaClient(),
+            customAbilities: {
+                // @ts-expect-error
+                ThisModelDoesntExist: {
+                    readOwnUser: {
+                        description: "Read own user",
+                        operation: "SELECT",
+                        expression: "true"
+                    }
+                }
+            },
+            // @ts-expect-error
+            getRoles: function (abilities) {
+                return {
+                    User: [abilities.User.readOwnUser]
+                };
+            },
+            getContext: function () { return ({
+                role: "User",
+                context: {}
+            }); }
+        });
+        (0, src_1.setup)({
+            prisma: new client_1.PrismaClient(),
+            customAbilities: {
+                // @ts-expect-error
+                User: {
+                    readOwnUser: {
+                        description: "Read own user",
+                        operation: "SELECT",
+                        expression: "true"
+                    }
+                }
+            },
+            // @ts-expect-error
+            getRoles: function (abilities) {
+                return {
+                    User: [abilities.User.readOwnUser]
+                };
+            },
+            getContext: function () { return ({
+                role: "User",
+                context: {}
+            }); }
+        });
+        return [2 /*return*/];
+    });
+}); };
diff --git a/test/types/index.ts b/test/types/index.ts
new file mode 100644
index 0000000..42e1e94
--- /dev/null
+++ b/test/types/index.ts
@@ -0,0 +1,135 @@
+import { PrismaClient } from "@prisma/client";
+import _ from "lodash";
+import { setup } from "../../src";
+
+const run = async () => {
+	// it should error if an invalid model is used
+	setup({
+		prisma: new PrismaClient(),
+		customAbilities: {
+			// @ts-expect-error
+			ThisModelDoesntExist: {
+				readOwnUser: {
+					description: "Read own user",
+					operation: "SELECT",
+					expression: "true",
+				},
+			},
+		},
+		// @ts-expect-error
+		getRoles(abilities) {
+			return {
+				User: [abilities.User.readOwnUser],
+			};
+		},
+		getContext: () => ({
+			role: "User",
+			context: {},
+		}),
+	});
+
+	// It should error if an unknown custom ability is used
+	setup({
+		prisma: new PrismaClient(),
+		customAbilities: {},
+		getRoles(abilities) {
+			return {
+				// @ts-expect-error
+				User: [abilities.User.superCustomAbility],
+			};
+		},
+		getContext: () => ({
+			role: "User",
+			context: {},
+		}),
+	});
+
+	// It should error if an incorrect where clause is used for a custom ability
+	setup({
+		prisma: new PrismaClient(),
+		customAbilities: {
+			User: {
+				superCustomAbility: {
+					description: "Super custom ability",
+					operation: "SELECT",
+					// @ts-expect-error
+					expression: (_client, _row, _context) => {
+						return {
+							foo: "bar",
+						};
+					},
+				},
+			},
+		},
+		// @ts-expect-error
+		getRoles(abilities) {
+			return {
+				User: [abilities.User.superCustomAbility],
+			};
+		},
+		getContext: () => ({
+			role: "User",
+			context: {},
+		}),
+	});
+
+	// It should error if an incorrect row key is used for a custom ability
+	setup({
+		prisma: new PrismaClient(),
+		customAbilities: {
+			User: {
+				superCustomAbility: {
+					description: "Super custom ability",
+					operation: "SELECT",
+					// @ts-expect-error
+					expression: (_client, row, _context) => {
+						return {
+							// @ts-expect-error
+							id: row("foo"),
+						};
+					},
+				},
+			},
+		},
+		// @ts-expect-error
+		getRoles(abilities) {
+			return {
+				User: [abilities.User.superCustomAbility],
+			};
+		},
+		getContext: () => ({
+			role: "User",
+			context: {},
+		}),
+	});
+
+	// It should error if an incorrect context key is used for a custom ability
+	setup({
+		prisma: new PrismaClient(),
+		customAbilities: {
+			User: {
+				superCustomAbility: {
+					description: "Super custom ability",
+					operation: "SELECT",
+					expression: (_client, _row, context) => {
+						return {
+							// @ts-expect-error
+							id: context("foo"),
+						};
+					},
+				},
+			},
+		},
+		getRoles(abilities) {
+			return {
+				User: [abilities.User.superCustomAbility],
+			};
+		},
+		getContext: () => ({
+			role: "User",
+			context: {
+				"user.id": "123",
+			},
+		}),
+	});
+};
diff --git a/tsconfig.json b/tsconfig.json
index 8104fad..b7422fa 100644
--- a/tsconfig.json
+++ b/tsconfig.json
@@ -1,15 +1,15 @@
 {
-  "compilerOptions": {
-    "esModuleInterop": true,
-    "lib": ["es2020"],
-    "module": "CommonJS",
-    "outDir": "./dist",
-    "sourceMap": true,
-    "strict": true,
-    "useUnknownInCatchVariables": false,
-    "noUnusedParameters": true,
-    "declaration": true,
-    "downlevelIteration": true
-  },
-  "include": ["src/**/*"]
+	"compilerOptions": {
+		"esModuleInterop": true,
+		"lib": ["es2020"],
+		"module": "CommonJS",
+		"outDir": "./dist",
+		"sourceMap": true,
+		"strict": true,
+		"useUnknownInCatchVariables": false,
+		"noUnusedParameters": true,
+		"declaration": true,
+		"downlevelIteration": true
+	},
+	"include": ["src/**/*", "test/types/index.ts"]
 }
diff --git a/yarn.lock b/yarn.lock
index 432147d..178ddcb 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -2250,6 +2250,11 @@ type-fest@^0.21.3:
   resolved "https://registry.yarnpkg.com/type-fest/-/type-fest-0.21.3.tgz#d260a24b0198436e133fa26a524a6d65fa3b2e37"
   integrity sha512-t0rzBq87m3fVcduHDUFhKmyyX+9eo6WQjZvf51Ea/M0Q7+T374Jp1aUiyUl0GKxp8M/OETVHSDvmkyPgvX+X2w==
 
+type-fest@^4.10.3:
+  version "4.10.3"
+  resolved "https://registry.yarnpkg.com/type-fest/-/type-fest-4.10.3.tgz#ff01cb0a1209f59583d61e1312de9715e7ea4874"
+  integrity sha512-JLXyjizi072smKGGcZiAJDCNweT8J+AuRxmPZ1aG7TERg4ijx9REl8CNhbr36RV4qXqL1gO1FF9HL8OkVmmrsA==
+
 typescript@^4.9.4:
   version "4.9.5"
   resolved "https://registry.yarnpkg.com/typescript/-/typescript-4.9.5.tgz#095979f9bcc0d09da324d58d03ce8f8374cbe65a"