@cerebruminc/yates 3.7.1 → 3.7.2-beta.dangerous.30d5b2f

package/dist/index.js

@@ -113,9 +113,10 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
     return (mod && mod.__esModule) ? mod : { "default": mod };
 };
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.setup = exports.createRoles = exports.createClient = exports.getBatchId = exports.createRoleName = exports.createAbilityName = exports.sanitizeSlug = void 0;
+exports.setup = exports.Yates = exports.sanitizeSlug = void 0;
 var crypto = __importStar(require("crypto"));
 var debug_1 = __importDefault(require("debug"));
+var cloneDeep_1 = __importDefault(require("lodash/cloneDeep"));
 var difference_1 = __importDefault(require("lodash/difference"));
 var flatMap_1 = __importDefault(require("lodash/flatMap"));
 var map_1 = __importDefault(require("lodash/map"));
@@ -130,23 +131,6 @@ var debug = (0, debug_1.default)("yates");
 var takeLock = function (prisma) {
     return prisma.$executeRawUnsafe("SELECT pg_advisory_xact_lock(2142616474639426746);");
 };
-/*
- * This function creates a table used to track the abilities that have been
- * defined in the system. We can use this to see if an ability needs to be updated.
- * We can't look up the pg policy table for this, as pg performs formatting on
- * the expression, making it very hard to check if the two expressions are equivalent.
- *
- * We also need to create a schema for this table, as we don't want to pollute the public schema.
- * If we use the public schema, we could potentially conflict with a user's table and we will
- * also cause issues for Prisma's migrate tooling, as it will detect a DB drift.
- */
-var setupAbilityTable = function (prisma) {
-    return prisma.$transaction([
-        takeLock(prisma),
-        prisma.$executeRawUnsafe("\n\t\tCREATE SCHEMA IF NOT EXISTS _yates;\n\t\t"),
-        prisma.$executeRawUnsafe("\n\t\tCREATE TABLE IF NOT EXISTS _yates._yates_abilities (\n\t\t\tid SERIAL PRIMARY KEY,\n\t\t\tability_model TEXT NOT NULL,\n\t\t\tability_name TEXT NOT NULL,\n\t\t\tability_policy_name TEXT NOT NULL UNIQUE,\n\t\t\tability_description TEXT NOT NULL,\n\t\t\tability_operation TEXT NOT NULL,\n\t\t\tability_expression TEXT NOT NULL,\n\t\t\tcreated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\n\t\t\tupdated_at TIMESTAMP\n\t\t);\n\t"),
-    ]);
-};
 var upsertAbility = function (prisma, ability) {
     var ability_model = ability.ability_model, ability_name = ability.ability_name, ability_policy_name = ability.ability_policy_name, ability_description = ability.ability_description, ability_operation = ability.ability_operation, ability_expression = ability.ability_expression;
     return prisma.$queryRaw(templateObject_1 || (templateObject_1 = __makeTemplateObject(["\n\t\tINSERT INTO _yates._yates_abilities (ability_model, ability_name, ability_policy_name, ability_description, ability_operation, ability_expression)\n\t\tVALUES (", ", ", ", ", ", ", ", ", ", ", ")\n\t\tON CONFLICT (ability_policy_name) DO UPDATE\n\t\tSET ability_model = EXCLUDED.ability_model, ability_name = EXCLUDED.ability_name, ability_description = EXCLUDED.ability_description, ability_operation = EXCLUDED.ability_operation, ability_expression = EXCLUDED.ability_expression, updated_at = now();\n\t"], ["\n\t\tINSERT INTO _yates._yates_abilities (ability_model, ability_name, ability_policy_name, ability_description, ability_operation, ability_expression)\n\t\tVALUES (", ", ", ", ", ", ", ", ", ", ", ")\n\t\tON CONFLICT (ability_policy_name) DO UPDATE\n\t\tSET ability_model = EXCLUDED.ability_model, ability_name = EXCLUDED.ability_name, ability_description = EXCLUDED.ability_description, ability_operation = EXCLUDED.ability_operation, ability_expression = EXCLUDED.ability_expression, updated_at = now();\n\t"])), ability_model, ability_name, ability_policy_name, ability_description, ability_operation, ability_expression);
@@ -173,660 +157,786 @@ var sanitizeSlug = function (slug) {
         .replace(/[^a-z0-9_]/gi, "");
 };
 exports.sanitizeSlug = sanitizeSlug;
-var createAbilityName = function (model, ability) {
-    return (0, exports.sanitizeSlug)(hashWithPrefix("yates_ability_", "".concat(model, "_").concat(ability)));
-};
-exports.createAbilityName = createAbilityName;
-var createRoleName = function (name) {
-    return (0, exports.sanitizeSlug)(hashWithPrefix("yates_role_", "".concat(name)));
-};
-exports.createRoleName = createRoleName;
-// @ts-ignore
-function getBatchId(query) {
-    if (query.action !== "findUnique" && query.action !== "findUniqueOrThrow") {
-        return undefined;
-    }
-    var parts = [];
-    if (query.modelName) {
-        parts.push(query.modelName);
-    }
-    if (query.query.arguments) {
-        parts.push(buildKeysString(query.query.arguments));
-    }
-    parts.push(buildKeysString(query.query.selection));
-    return parts.join("");
-}
-exports.getBatchId = getBatchId;
-function buildKeysString(obj) {
-    var keysArray = Object.keys(obj)
-        .sort()
-        .map(function (key) {
-        // @ts-ignore
-        var value = obj[key];
-        if (typeof value === "object" && value !== null) {
-            return "(".concat(key, " ").concat(buildKeysString(value), ")");
-        }
-        return key;
-    });
-    return "(".concat(keysArray.join(" "), ")");
-}
-// This uses client extensions to set the role and context for the current user so that RLS can be applied
-var createClient = function (prisma, getContext, options) {
-    if (options === void 0) { options = {}; }
-    // Set default options
-    var _a = options.txMaxWait, txMaxWait = _a === void 0 ? 30000 : _a, _b = options.txTimeout, txTimeout = _b === void 0 ? 30000 : _b;
-    // By default, Prisma will batch requests by the transaction ID if it is present.
-    // This behaviour prevents automatic batching from working when using Yates, since all queries are executed inside an interactive transaction.
-    // To get around this we by monkey patching the batching function to use the Yates ID as the batch ID.
-    // To get the batching to work we also need to ensure that all the requests we might want to batch together are generated inside the same tick.
-    // This means that all the requests per-tick that have the same role and context values will be batched together,
-    // allowing the in-built prisma batch optimizations to work for us.
-    // This is why we use process.nextTick and the tickActive flag to ensure we only tick once at a time.
-    // See:
-    // - https://github.com/prisma/prisma/blob/5.21.1/packages/client/src/runtime/RequestHandler.ts#L122
-    // - https://www.prisma.io/docs/orm/prisma-client/queries/query-optimization-performance
-    prisma._requestHandler.dataloader.options.batchBy = function (request) {
-        var _a;
-        var batchIdPQ = getBatchId(request.protocolQuery);
-        if ((_a = request.transaction) === null || _a === void 0 ? void 0 : _a.id) {
-            return "transaction-".concat(request.transaction.id).concat(batchIdPQ ? "-".concat(batchIdPQ) : "");
-        }
-        return getBatchId(request.protocolQuery);
-    };
-    var tickActive = false;
-    var batches = {};
-    // This function is called once per tick, and processes all the batches that have been created during that tick.
-    // Each batch represents a unique role + context combination, and contains all the requests that need to be executed with that role + context.
-    var dispatchBatches = function () {
-        var e_1, _a;
-        var _loop_1 = function (key, batch) {
-            delete batches[key];
-            // Because batch transactions inside a prisma client query extension can run out of order if used with async middleware,
-            // we need to run the logic inside an interactive transaction, however this brings a different set of problems in that the
-            // main query will no longer automatically run inside the transaction. We resolve this issue by manually executing the prisma request.
-            // See https://github.com/prisma/prisma/issues/18276
-            prisma
-                .$transaction(function (tx) { return __awaiter(void 0, void 0, void 0, function () {
-                var _a, _b, _c, key_1, value, e_2_1, txId, results;
-                var e_2, _d;
-                return __generator(this, function (_e) {
-                    switch (_e.label) {
-                        case 0: 
-                        // Switch to the user role, We can't use a prepared statement here, due to limitations in PG not allowing prepared statements to be used in SET LOCAL ROLE
-                        return [4 /*yield*/, tx.$queryRawUnsafe("SET LOCAL ROLE ".concat(batch.pgRole))];
-                        case 1:
-                            // Switch to the user role, We can't use a prepared statement here, due to limitations in PG not allowing prepared statements to be used in SET LOCAL ROLE
-                            _e.sent();
-                            _e.label = 2;
-                        case 2:
-                            _e.trys.push([2, 7, 8, 9]);
-                            _a = __values((0, toPairs_1.default)(batch.context)), _b = _a.next();
-                            _e.label = 3;
-                        case 3:
-                            if (!!_b.done) return [3 /*break*/, 6];
-                            _c = __read(_b.value, 2), key_1 = _c[0], value = _c[1];
-                            return [4 /*yield*/, tx.$queryRaw(templateObject_2 || (templateObject_2 = __makeTemplateObject(["SELECT set_config(", ", ", ", true);"], ["SELECT set_config(", ", ", ", true);"])), key_1, value.toString())];
-                        case 4:
-                            _e.sent();
-                            _e.label = 5;
-                        case 5:
-                            _b = _a.next();
-                            return [3 /*break*/, 3];
-                        case 6: return [3 /*break*/, 9];
-                        case 7:
-                            e_2_1 = _e.sent();
-                            e_2 = { error: e_2_1 };
-                            return [3 /*break*/, 9];
-                        case 8:
-                            try {
-                                if (_b && !_b.done && (_d = _a.return)) _d.call(_a);
-                            }
-                            finally { if (e_2) throw e_2.error; }
-                            return [7 /*endfinally*/];
-                        case 9:
-                            txId = tx[Symbol.for("prisma.client.transaction.id")];
-                            return [4 /*yield*/, Promise.all(batch.requests.map(function (request) {
-                                    return prisma._executeRequest(__assign(__assign({}, request.params), { transaction: {
-                                            kind: "itx",
-                                            id: txId,
-                                        } }));
-                                }))];
-                        case 10:
-                            results = _e.sent();
-                            return [2 /*return*/, results];
-                    }
-                });
-            }); }, {
-                maxWait: txMaxWait,
-                timeout: txTimeout,
-            })
-                .then(function (results) {
-                results.forEach(function (result, index) {
-                    batch.requests[index].resolve(result);
-                });
-            })
-                .catch(function (e) {
-                var e_3, _a;
-                try {
-                    for (var _b = (e_3 = void 0, __values(batch.requests)), _c = _b.next(); !_c.done; _c = _b.next()) {
-                        var request = _c.value;
-                        request.reject(e);
-                    }
+var Yates = /** @class */ (function () {
+    function Yates(prisma) {
+        var _this = this;
+        this.prisma = prisma;
+        this.databaseScope = null;
+        this.init = function () { return __awaiter(_this, void 0, void 0, function () {
+            return __generator(this, function (_a) {
+                switch (_a.label) {
+                    case 0: return [4 /*yield*/, this.ensureDatabaseScope()];
+                    case 1:
+                        _a.sent();
+                        debug("Setting up ability table");
+                        return [4 /*yield*/, this.setupAbilityTable()];
+                    case 2:
+                        _a.sent();
+                        return [2 /*return*/];
                 }
-                catch (e_3_1) { e_3 = { error: e_3_1 }; }
-                finally {
-                    try {
-                        if (_c && !_c.done && (_a = _b.return)) _a.call(_b);
-                    }
-                    finally { if (e_3) throw e_3.error; }
-                }
-                delete batches[key];
             });
+        }); };
+        this.createDatabaseScope = function (databaseName) {
+            var sanitizedName = (0, exports.sanitizeSlug)(databaseName);
+            if (sanitizedName.length > 0) {
+                return sanitizedName;
+            }
+            var hash = crypto.createHash("sha256");
+            hash.update(databaseName);
+            return hash.digest("hex").slice(0, 8);
         };
-        try {
-            for (var _b = __values(Object.entries(batches)), _c = _b.next(); !_c.done; _c = _b.next()) {
-                var _d = __read(_c.value, 2), key = _d[0], batch = _d[1];
-                _loop_1(key, batch);
+        this.getDatabaseScope = function () {
+            if (!_this.databaseScope) {
+                throw new Error("Yates database scope has not been initialised. Ensure setup() has been called before using the client.");
             }
-        }
-        catch (e_1_1) { e_1 = { error: e_1_1 }; }
-        finally {
+            return _this.databaseScope;
+        };
+        this.ensureDatabaseScope = function () { return __awaiter(_this, void 0, void 0, function () {
+            var result, currentDatabase;
+            var _a;
+            return __generator(this, function (_b) {
+                switch (_b.label) {
+                    case 0:
+                        if (this.databaseScope) {
+                            return [2 /*return*/, this.databaseScope];
+                        }
+                        return [4 /*yield*/, this.prisma.$queryRawUnsafe("select current_database() as current_database;")];
+                    case 1:
+                        result = _b.sent();
+                        currentDatabase = (_a = result[0]) === null || _a === void 0 ? void 0 : _a.current_database;
+                        debug("Current database for Yates:", currentDatabase);
+                        if (!currentDatabase) {
+                            throw new Error("Failed to determine the current database for scoping Yates roles.");
+                        }
+                        this.databaseScope = this.createDatabaseScope(currentDatabase);
+                        return [2 /*return*/, this.databaseScope];
+                }
+            });
+        }); };
+        /*
+         * This function creates a table used to track the abilities that have been
+         * defined in the system. We can use this to see if an ability needs to be updated.
+         * We can't look up the pg policy table for this, as pg performs formatting on
+         * the expression, making it very hard to check if the two expressions are equivalent.
+         *
+         * We also need to create a schema for this table, as we don't want to pollute the public schema.
+         * If we use the public schema, we could potentially conflict with a user's table and we will
+         * also cause issues for Prisma's migrate tooling, as it will detect a DB drift.
+         */
+        this.setupAbilityTable = function () {
+            return _this.prisma.$transaction([
+                takeLock(_this.prisma),
+                _this.prisma.$executeRawUnsafe("\n\t\t\t\tCREATE SCHEMA IF NOT EXISTS _yates;\n\t\t\t"),
+                _this.prisma.$executeRawUnsafe("\n\t\t\t\tCREATE TABLE IF NOT EXISTS _yates._yates_abilities (\n\t\t\t\t\tid SERIAL PRIMARY KEY,\n\t\t\t\t\tability_model TEXT NOT NULL,\n\t\t\t\t\tability_name TEXT NOT NULL,\n\t\t\t\t\tability_policy_name TEXT NOT NULL UNIQUE,\n\t\t\t\t\tability_description TEXT NOT NULL,\n\t\t\t\t\tability_operation TEXT NOT NULL,\n\t\t\t\t\tability_expression TEXT NOT NULL,\n\t\t\t\t\tcreated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,\n\t\t\t\t\tupdated_at TIMESTAMP\n\t\t\t\t);\n\t\t\t"),
+            ]);
+        };
+        this.createAbilityName = function (model, ability) {
+            var scope = _this.getDatabaseScope();
+            return (0, exports.sanitizeSlug)(hashWithPrefix("yates_ability_", "".concat(scope, "_").concat(model, "_").concat(ability)));
+        };
+        this.createRoleName = function (name) {
+            var scope = _this.getDatabaseScope();
+            return (0, exports.sanitizeSlug)(hashWithPrefix("yates_role_", "".concat(scope, "_").concat(name)));
+        };
+        this.getDefaultAbilities = function (models) {
+            var e_1, _a;
+            var abilities = {};
             try {
-                if (_c && !_c.done && (_a = _b.return)) _a.call(_b);
+                for (var models_1 = __values(models), models_1_1 = models_1.next(); !models_1_1.done; models_1_1 = models_1.next()) {
+                    var model = models_1_1.value;
+                    abilities[model] = {
+                        create: {
+                            description: "Create ".concat(model),
+                            expression: "true",
+                            operation: "INSERT",
+                            model: model,
+                            slug: "create",
+                        },
+                        read: {
+                            description: "Read ".concat(model),
+                            expression: "true",
+                            operation: "SELECT",
+                            model: model,
+                            slug: "read",
+                        },
+                        update: {
+                            description: "Update ".concat(model),
+                            expression: "true",
+                            operation: "UPDATE",
+                            model: model,
+                            slug: "update",
+                        },
+                        delete: {
+                            description: "Delete ".concat(model),
+                            expression: "true",
+                            operation: "DELETE",
+                            model: model,
+                            slug: "delete",
+                        },
+                    };
+                }
             }
-            finally { if (e_1) throw e_1.error; }
-        }
-    };
-    var client = prisma.$extends({
-        name: "Yates client",
-        query: {
-            $allModels: {
-                $allOperations: function (params) {
-                    return __awaiter(this, void 0, void 0, function () {
-                        var model, args, query, operation, ctx, role, context, pgRole, _a, _b, k, _c, _d, v, txId, hash, __internalParams;
-                        var e_4, _e, e_5, _f;
-                        return __generator(this, function (_g) {
-                            model = params.model, args = params.args, query = params.query, operation = params.operation;
-                            if (!model) {
-                                // If the model is not defined, we can't apply RLS
-                                // This can occur when you are making a call with Prisma's $queryRaw method
-                                return [2 /*return*/, query(args)];
-                            }
-                            ctx = getContext();
-                            // If ctx is null, the middleware is explicitly skipped
-                            if (ctx === null) {
-                                return [2 /*return*/, query(args)];
-                            }
-                            role = ctx.role, context = ctx.context;
-                            pgRole = (0, exports.createRoleName)(role);
-                            if (context) {
-                                try {
-                                    for (_a = __values(Object.keys(context)), _b = _a.next(); !_b.done; _b = _a.next()) {
-                                        k = _b.value;
-                                        if (!k.match(/^[a-z_\.]+$/)) {
-                                            throw new Error("Context variable \"".concat(k, "\" contains invalid characters. Context variables must only contain lowercase letters, numbers, periods and underscores."));
-                                        }
-                                        if (typeof context[k] !== "number" &&
-                                            typeof context[k] !== "string" &&
-                                            !Array.isArray(context[k])) {
-                                            throw new Error("Context variable \"".concat(k, "\" must be a string, number or array. Got ").concat(typeof context[k]));
-                                        }
-                                        if (Array.isArray(context[k])) {
-                                            try {
-                                                for (_c = (e_5 = void 0, __values(context[k])), _d = _c.next(); !_d.done; _d = _c.next()) {
-                                                    v = _d.value;
-                                                    if (typeof v !== "string") {
-                                                        throw new Error("Context variable \"".concat(k, "\" must be an array of strings. Got ").concat(typeof v));
-                                                    }
-                                                }
-                                            }
-                                            catch (e_5_1) { e_5 = { error: e_5_1 }; }
-                                            finally {
-                                                try {
-                                                    if (_d && !_d.done && (_f = _c.return)) _f.call(_c);
-                                                }
-                                                finally { if (e_5) throw e_5.error; }
-                                            }
-                                            // Cast to a JSON string so that it can be used in RLS expressions
-                                            context[k] = JSON.stringify(context[k]);
-                                        }
-                                    }
-                                }
-                                catch (e_4_1) { e_4 = { error: e_4_1 }; }
-                                finally {
-                                    try {
-                                        if (_b && !_b.done && (_e = _a.return)) _e.call(_a);
-                                    }
-                                    finally { if (e_4) throw e_4.error; }
-                                }
-                            }
-                            txId = hashWithPrefix("yates_tx_", JSON.stringify(ctx));
-                            hash = txId;
-                            if (!batches[hash]) {
-                                // Create a new batch for this role + context combination
-                                batches[hash] = {
-                                    pgRole: pgRole,
-                                    context: context,
-                                    requests: [],
-                                };
-                                // make sure, that we only tick once at a time
-                                if (!tickActive) {
-                                    tickActive = true;
-                                    process.nextTick(function () {
-                                        dispatchBatches();
-                                        tickActive = false;
-                                    });
-                                }
-                            }
-                            __internalParams = params.__internalParams;
-                            // Add the request to the batch, and return a promise that will be resolved or rejected in dispatchBatches
-                            return [2 /*return*/, new Promise(function (resolve, reject) {
-                                    batches[hash].requests.push({
-                                        params: __internalParams,
-                                        query: query,
-                                        args: args,
-                                        resolve: resolve,
-                                        reject: reject,
-                                    });
-                                }).catch(function (e) {
-                                    var _a;
-                                    // Normalize RLS errors to make them a bit more readable.
-                                    if ((_a = e.message) === null || _a === void 0 ? void 0 : _a.includes("new row violates row-level security policy for table")) {
-                                        throw new Error("You do not have permission to perform this action: ".concat(model, ".").concat(operation, "(...)"));
-                                    }
-                                    throw e;
-                                })];
-                        });
-                    });
-                },
-            },
-        },
-    });
-    return client;
-};
-exports.createClient = createClient;
-var setRLS = function (prisma, table, roleName, slug, ability) { return __awaiter(void 0, void 0, void 0, function () {
-    var operation, rawExpression, description;
-    return __generator(this, function (_a) {
-        switch (_a.label) {
-            case 0:
-                operation = ability.operation, rawExpression = ability.expression, description = ability.description;
-                if (!rawExpression) {
-                    throw new Error("Expression must be defined for RLS abilities");
+            catch (e_1_1) { e_1 = { error: e_1_1 }; }
+            finally {
+                try {
+                    if (models_1_1 && !models_1_1.done && (_a = models_1.return)) _a.call(models_1);
                 }
-                // Take a lock and run the RLS setup in a transaction to prevent conflicts
-                // in a multi-server environment
-                return [4 /*yield*/, prisma.$transaction(function (tx) { return __awaiter(void 0, void 0, void 0, function () {
-                        var policyName, existingAbilities, existingAbility, shouldUpdateAbilityTable, expression, expression;
-                        return __generator(this, function (_a) {
-                            switch (_a.label) {
-                                case 0: return [4 /*yield*/, takeLock(tx)];
+                finally { if (e_1) throw e_1.error; }
+            }
+            return abilities;
+        };
+        // This uses client extensions to set the role and context for the current user so that RLS can be applied
+        this.createClient = function (getContext, options) {
+            if (options === void 0) { options = {}; }
+            var prisma = _this.prisma;
+            // Set default options
+            var _a = options.txMaxWait, txMaxWait = _a === void 0 ? 30000 : _a, _b = options.txTimeout, txTimeout = _b === void 0 ? 30000 : _b;
+            // By default, Prisma will batch requests by the transaction ID if it is present.
+            // This behaviour prevents automatic batching from working when using Yates, since all queries are executed inside an interactive transaction.
+            // To get around this we by monkey patching the batching function to use the Yates ID as the batch ID.
+            // To get the batching to work we also need to ensure that all the requests we might want to batch together are generated inside the same tick.
+            // This means that all the requests per-tick that have the same role and context values will be batched together,
+            // allowing the in-built prisma batch optimizations to work for us.
+            // This is why we use process.nextTick and the tickActive flag to ensure we only tick once at a time.
+            // See:
+            // - https://github.com/prisma/prisma/blob/5.21.1/packages/client/src/runtime/RequestHandler.ts#L122
+            // - https://www.prisma.io/docs/orm/prisma-client/queries/query-optimization-performance
+            prisma._requestHandler.dataloader.options.batchBy = function (request) {
+                var _a;
+                var batchIdPQ = _this.getBatchId(request.protocolQuery);
+                if ((_a = request.transaction) === null || _a === void 0 ? void 0 : _a.id) {
+                    return "transaction-".concat(request.transaction.id).concat(batchIdPQ ? "-".concat(batchIdPQ) : "");
+                }
+                return _this.getBatchId(request.protocolQuery);
+            };
+            var tickActive = false;
+            var batches = {};
+            // This function is called once per tick, and processes all the batches that have been created during that tick.
+            // Each batch represents a unique role + context combination, and contains all the requests that need to be executed with that role + context.
+            var dispatchBatches = function () {
+                var e_2, _a;
+                var _loop_1 = function (key, batch) {
+                    delete batches[key];
+                    // Because batch transactions inside a prisma client query extension can run out of order if used with async middleware,
+                    // we need to run the logic inside an interactive transaction, however this brings a different set of problems in that the
+                    // main query will no longer automatically run inside the transaction. We resolve this issue by manually executing the prisma request.
+                    // See https://github.com/prisma/prisma/issues/18276
+                    prisma
+                        .$transaction(function (tx) { return __awaiter(_this, void 0, void 0, function () {
+                        var _a, _b, _c, key_1, value, e_3_1, txId, results;
+                        var e_3, _d;
+                        return __generator(this, function (_e) {
+                            switch (_e.label) {
+                                case 0: 
+                                // Switch to the user role, We can't use a prepared statement here, due to limitations in PG not allowing prepared statements to be used in SET LOCAL ROLE
+                                return [4 /*yield*/, tx.$queryRawUnsafe("SET LOCAL ROLE ".concat(batch.pgRole))];
                                 case 1:
-                                    _a.sent();
-                                    policyName = roleName;
-                                    return [4 /*yield*/, tx.$queryRaw(templateObject_3 || (templateObject_3 = __makeTemplateObject(["\n\t\t\tselect * from _yates._yates_abilities where ability_model = ", " and ability_policy_name = ", "\n\t\t"], ["\n\t\t\tselect * from _yates._yates_abilities where ability_model = ", " and ability_policy_name = ", "\n\t\t"])), table, policyName)];
+                                    // Switch to the user role, We can't use a prepared statement here, due to limitations in PG not allowing prepared statements to be used in SET LOCAL ROLE
+                                    _e.sent();
+                                    _e.label = 2;
                                 case 2:
-                                    existingAbilities = _a.sent();
-                                    existingAbility = existingAbilities[0];
-                                    shouldUpdateAbilityTable = false;
-                                    if (!!existingAbility) return [3 /*break*/, 8];
-                                    debug("Creating RLS policy for", roleName, "on", table, "for", operation);
-                                    return [4 /*yield*/, (0, expressions_1.expressionToSQL)(rawExpression, table)];
+                                    _e.trys.push([2, 7, 8, 9]);
+                                    _a = __values((0, toPairs_1.default)(batch.context)), _b = _a.next();
+                                    _e.label = 3;
                                 case 3:
-                                    expression = _a.sent();
-                                    if (!(operation === "INSERT")) return [3 /*break*/, 5];
-                                    return [4 /*yield*/, tx.$queryRawUnsafe("\n\t\t\t\tCREATE POLICY ".concat(policyName, " ON \"public\".\"").concat(table, "\" FOR ").concat(operation, " TO ").concat(roleName, " WITH CHECK (").concat(expression, ");\n\t\t\t"))];
+                                    if (!!_b.done) return [3 /*break*/, 6];
+                                    _c = __read(_b.value, 2), key_1 = _c[0], value = _c[1];
+                                    return [4 /*yield*/, tx.$queryRaw(templateObject_2 || (templateObject_2 = __makeTemplateObject(["SELECT set_config(", ", ", ", true);"], ["SELECT set_config(", ", ", ", true);"])), key_1, value.toString())];
                                 case 4:
-                                    _a.sent();
-                                    return [3 /*break*/, 7];
-                                case 5: return [4 /*yield*/, tx.$queryRawUnsafe("\n\t\t\t\tCREATE POLICY ".concat(policyName, " ON \"public\".\"").concat(table, "\" FOR ").concat(operation, " TO ").concat(roleName, " USING (").concat(expression, ");\n\t\t\t"))];
-                                case 6:
-                                    _a.sent();
-                                    _a.label = 7;
+                                    _e.sent();
+                                    _e.label = 5;
+                                case 5:
+                                    _b = _a.next();
+                                    return [3 /*break*/, 3];
+                                case 6: return [3 /*break*/, 9];
                                 case 7:
-                                    shouldUpdateAbilityTable = true;
-                                    return [3 /*break*/, 14];
+                                    e_3_1 = _e.sent();
+                                    e_3 = { error: e_3_1 };
+                                    return [3 /*break*/, 9];
                                 case 8:
-                                    if (!(existingAbility.ability_expression !== rawExpression.toString())) return [3 /*break*/, 14];
-                                    debug("Updating RLS policy for", roleName, "on", table, "for", operation);
-                                    return [4 /*yield*/, (0, expressions_1.expressionToSQL)(rawExpression, table)];
+                                    try {
+                                        if (_b && !_b.done && (_d = _a.return)) _d.call(_a);
+                                    }
+                                    finally { if (e_3) throw e_3.error; }
+                                    return [7 /*endfinally*/];
                                 case 9:
-                                    expression = _a.sent();
-                                    if (!(operation === "INSERT")) return [3 /*break*/, 11];
-                                    return [4 /*yield*/, tx.$queryRawUnsafe("\n\t\t\t\tALTER POLICY ".concat(policyName, " ON \"public\".\"").concat(table, "\" TO ").concat(roleName, " WITH CHECK (").concat(expression, ");\n\t\t\t"))];
+                                    txId = tx[Symbol.for("prisma.client.transaction.id")];
+                                    return [4 /*yield*/, Promise.all(batch.requests.map(function (request) {
+                                            return prisma._executeRequest(__assign(__assign({}, request.params), { transaction: {
+                                                    kind: "itx",
+                                                    id: txId,
+                                                } }));
+                                        }))];
                                 case 10:
-                                    _a.sent();
-                                    return [3 /*break*/, 13];
-                                case 11: return [4 /*yield*/, tx.$queryRawUnsafe("\n\t\t\t\tALTER POLICY ".concat(policyName, " ON \"public\".\"").concat(table, "\" TO ").concat(roleName, " USING (").concat(expression, ");\n\t\t\t"))];
-                                case 12:
-                                    _a.sent();
-                                    _a.label = 13;
-                                case 13:
-                                    shouldUpdateAbilityTable = true;
-                                    _a.label = 14;
-                                case 14:
-                                    if (!shouldUpdateAbilityTable) return [3 /*break*/, 16];
-                                    return [4 /*yield*/, upsertAbility(tx, {
-                                            ability_model: table,
-                                            ability_name: slug,
-                                            ability_policy_name: policyName,
-                                            ability_description: description !== null && description !== void 0 ? description : "",
-                                            ability_operation: operation,
-                                            // We store the string representation of the expression so that
-                                            // we can compare it later without having to recompute the SQL
-                                            ability_expression: rawExpression.toString(),
-                                        })];
-                                case 15:
-                                    _a.sent();
-                                    _a.label = 16;
-                                case 16: return [2 /*return*/];
+                                    results = _e.sent();
+                                    return [2 /*return*/, results];
                             }
                         });
-                    }); })];
-            case 1:
-                // Take a lock and run the RLS setup in a transaction to prevent conflicts
-                // in a multi-server environment
-                _a.sent();
-                return [2 /*return*/];
-        }
-    });
-}); };
-var createRoles = function (_a) { return __awaiter(void 0, [_a], void 0, function (_b) {
-    var abilities, runtimeDataModel, models, diff, models_1, models_1_1, model, ability, operation, roles, pgRoles, existingAbilities, pgPolicies, migratedAbilities, _c, _d, _e, _i, model, table, _loop_2, _f, _g, _h, _j, slug, _loop_3, _k, _l, _m, _o, key;
-    var e_6, _p;
-    var _q;
-    var prisma = _b.prisma, customAbilities = _b.customAbilities, getRoles = _b.getRoles;
-    return __generator(this, function (_r) {
-        switch (_r.label) {
-            case 0:
-                abilities = {};
-                runtimeDataModel = prisma
-                    ._runtimeDataModel;
-                models = Object.keys(runtimeDataModel.models).map(function (m) { return runtimeDataModel.models[m].dbName || m; });
-                if (customAbilities) {
-                    diff = (0, difference_1.default)(Object.keys(customAbilities), models);
-                    if (diff.length) {
-                        throw new Error("Invalid models in custom abilities: ".concat(diff.join(", ")));
-                    }
-                }
-                try {
-                    for (models_1 = __values(models), models_1_1 = models_1.next(); !models_1_1.done; models_1_1 = models_1.next()) {
-                        model = models_1_1.value;
-                        abilities[model] = {
-                            create: {
-                                description: "Create ".concat(model),
-                                expression: "true",
-                                operation: "INSERT",
-                                model: model,
-                                slug: "create",
-                            },
-                            read: {
-                                description: "Read ".concat(model),
-                                expression: "true",
-                                operation: "SELECT",
-                                model: model,
-                                slug: "read",
-                            },
-                            update: {
-                                description: "Update ".concat(model),
-                                expression: "true",
-                                operation: "UPDATE",
-                                model: model,
-                                slug: "update",
-                            },
-                            delete: {
-                                description: "Delete ".concat(model),
-                                expression: "true",
-                                operation: "DELETE",
-                                model: model,
-                                slug: "delete",
-                            },
-                        };
-                        if (customAbilities === null || customAbilities === void 0 ? void 0 : customAbilities[model]) {
-                            for (ability in customAbilities[model]) {
-                                operation = 
-                                // biome-ignore lint/style/noNonNullAssertion: TODO fix this
-                                (_q = customAbilities[model][ability]) === null || _q === void 0 ? void 0 : _q.operation;
-                                if (!operation)
-                                    continue;
-                                // biome-ignore lint/style/noNonNullAssertion: TODO fix this
-                                abilities[model][ability] = __assign(__assign({}, customAbilities[model][ability]), { operation: operation, model: model, slug: ability });
+                    }); }, {
+                        maxWait: txMaxWait,
+                        timeout: txTimeout,
+                    })
+                        .then(function (results) {
+                        results.forEach(function (result, index) {
+                            batch.requests[index].resolve(result);
+                        });
+                    })
+                        .catch(function (e) {
+                        var e_4, _a;
+                        try {
+                            for (var _b = (e_4 = void 0, __values(batch.requests)), _c = _b.next(); !_c.done; _c = _b.next()) {
+                                var request = _c.value;
+                                request.reject(e);
+                            }
+                        }
+                        catch (e_4_1) { e_4 = { error: e_4_1 }; }
+                        finally {
+                            try {
+                                if (_c && !_c.done && (_a = _b.return)) _a.call(_b);
                             }
+                            finally { if (e_4) throw e_4.error; }
                         }
+                        delete batches[key];
+                    });
+                };
+                try {
+                    for (var _b = __values(Object.entries(batches)), _c = _b.next(); !_c.done; _c = _b.next()) {
+                        var _d = __read(_c.value, 2), key = _d[0], batch = _d[1];
+                        _loop_1(key, batch);
                     }
                 }
-                catch (e_6_1) { e_6 = { error: e_6_1 }; }
+                catch (e_2_1) { e_2 = { error: e_2_1 }; }
                 finally {
                     try {
-                        if (models_1_1 && !models_1_1.done && (_p = models_1.return)) _p.call(models_1);
+                        if (_c && !_c.done && (_a = _b.return)) _a.call(_b);
                     }
-                    finally { if (e_6) throw e_6.error; }
+                    finally { if (e_2) throw e_2.error; }
                 }
-                debug("Setting up ability table");
-                return [4 /*yield*/, setupAbilityTable(prisma)];
-            case 1:
-                _r.sent();
-                roles = getRoles(abilities);
-                return [4 /*yield*/, prisma.$queryRawUnsafe("\n\t\tselect * from pg_catalog.pg_roles where rolname like 'yates%'\n\t")];
-            case 2:
-                pgRoles = _r.sent();
-                return [4 /*yield*/, prisma.$queryRawUnsafe("\n\t\tselect * from _yates._yates_abilities;\n\t")];
-            case 3:
-                existingAbilities = _r.sent();
-                if (!(existingAbilities.length === 0)) return [3 /*break*/, 6];
-                debug('No existing abilities found, importing from "pg_policies" table');
-                return [4 /*yield*/, prisma.$queryRawUnsafe("\n\t\t\tselect * from pg_catalog.pg_policies where policyname like 'yates%'\n\t\t")];
-            case 4:
-                pgPolicies = _r.sent();
-                if (!pgPolicies.length) return [3 /*break*/, 6];
-                migratedAbilities = pgPolicies.map(function (policy) {
-                    var _a, _b;
-                    return ({
-                        ability_model: policy.tablename,
-                        ability_name: policy.policyname,
-                        ability_policy_name: policy.policyname,
-                        ability_description: "",
-                        ability_operation: policy.cmd,
-                        ability_expression: (_b = (_a = policy.qual) !== null && _a !== void 0 ? _a : policy.with_check) !== null && _b !== void 0 ? _b : "",
-                    });
-                });
-                return [4 /*yield*/, prisma.$transaction(__spreadArray([
-                        takeLock(prisma)
-                    ], __read(migratedAbilities.map(function (ma) { return upsertAbility(prisma, ma); })), false))];
-            case 5:
-                _r.sent();
-                existingAbilities.push.apply(existingAbilities, __spreadArray([], __read(migratedAbilities), false));
-                _r.label = 6;
-            case 6:
-                _c = abilities;
-                _d = [];
-                for (_e in _c)
-                    _d.push(_e);
-                _i = 0;
-                _r.label = 7;
-            case 7:
-                if (!(_i < _d.length)) return [3 /*break*/, 13];
-                _e = _d[_i];
-                if (!(_e in _c)) return [3 /*break*/, 12];
-                model = _e;
-                table = model;
-                return [4 /*yield*/, prisma.$transaction([
-                        takeLock(prisma),
-                        prisma.$queryRawUnsafe("ALTER table \"".concat(table, "\" enable row level security;")),
-                    ])];
-            case 8:
-                _r.sent();
-                _loop_2 = function (slug) {
-                    var ability, roleName;
-                    return __generator(this, function (_s) {
-                        switch (_s.label) {
-                            case 0:
-                                ability = 
-                                // biome-ignore lint/style/noNonNullAssertion: TODO fix this
-                                abilities[model][slug];
-                                if (!VALID_OPERATIONS.includes(ability.operation)) {
-                                    throw new Error("Invalid operation: ".concat(ability.operation));
-                                }
-                                roleName = (0, exports.createAbilityName)(model, slug);
-                                if (!pgRoles.find(function (role) { return role.rolname === roleName; })) return [3 /*break*/, 1];
-                                debug("Role already exists", roleName);
-                                return [3 /*break*/, 3];
-                            case 1: return [4 /*yield*/, prisma.$transaction([
-                                    takeLock(prisma),
-                                    prisma.$queryRawUnsafe("\n\t\t\t\t\tdo\n\t\t\t\t\t$$\n\t\t\t\t\tbegin\n\t\t\t\t\tif not exists (select * from pg_catalog.pg_roles where rolname = '".concat(roleName, "') then \n\t\t\t\t\t\tcreate role ").concat(roleName, ";\n\t\t\t\t\tend if;\n\t\t\t\t\tend\n\t\t\t\t\t$$\n\t\t\t\t\t;\n\t\t\t\t")),
-                                    prisma.$queryRawUnsafe("\n\t\t\t\t\tGRANT ".concat(ability.operation, " ON \"").concat(table, "\" TO ").concat(roleName, ";\n\t\t\t\t")),
-                                ])];
-                            case 2:
-                                _s.sent();
-                                _s.label = 3;
-                            case 3:
-                                if (!ability.expression) return [3 /*break*/, 5];
-                                return [4 /*yield*/, setRLS(prisma, table, roleName, slug, ability)];
-                            case 4:
-                                _s.sent();
-                                _s.label = 5;
-                            case 5: return [2 /*return*/];
+            };
+            var createRoleName = _this.createRoleName.bind(_this);
+            var client = prisma.$extends({
+                name: "Yates client",
+                query: {
+                    $allModels: {
+                        $allOperations: function (params) {
+                            return __awaiter(this, void 0, void 0, function () {
+                                var model, args, query, operation, ctx, role, context, pgRole, _a, _b, k, _c, _d, v, txId, hash, __internalParams;
+                                var e_5, _e, e_6, _f;
+                                return __generator(this, function (_g) {
+                                    model = params.model, args = params.args, query = params.query, operation = params.operation;
+                                    if (!model) {
+                                        // If the model is not defined, we can't apply RLS
+                                        // This can occur when you are making a call with Prisma's $queryRaw method
+                                        return [2 /*return*/, query(args)];
+                                    }
+                                    ctx = getContext();
+                                    // If ctx is null, the middleware is explicitly skipped
+                                    if (ctx === null) {
+                                        return [2 /*return*/, query(args)];
+                                    }
+                                    role = ctx.role, context = ctx.context;
+                                    pgRole = createRoleName(role);
+                                    if (context) {
+                                        try {
+                                            for (_a = __values(Object.keys(context)), _b = _a.next(); !_b.done; _b = _a.next()) {
+                                                k = _b.value;
+                                                if (!k.match(/^[a-z_\.]+$/)) {
+                                                    throw new Error("Context variable \"".concat(k, "\" contains invalid characters. Context variables must only contain lowercase letters, numbers, periods and underscores."));
+                                                }
+                                                if (typeof context[k] !== "number" &&
+                                                    typeof context[k] !== "string" &&
+                                                    !Array.isArray(context[k])) {
+                                                    throw new Error("Context variable \"".concat(k, "\" must be a string, number or array. Got ").concat(typeof context[k]));
+                                                }
+                                                if (Array.isArray(context[k])) {
+                                                    try {
+                                                        for (_c = (e_6 = void 0, __values(context[k])), _d = _c.next(); !_d.done; _d = _c.next()) {
+                                                            v = _d.value;
+                                                            if (typeof v !== "string") {
+                                                                throw new Error("Context variable \"".concat(k, "\" must be an array of strings. Got ").concat(typeof v));
+                                                            }
+                                                        }
+                                                    }
+                                                    catch (e_6_1) { e_6 = { error: e_6_1 }; }
+                                                    finally {
+                                                        try {
+                                                            if (_d && !_d.done && (_f = _c.return)) _f.call(_c);
+                                                        }
+                                                        finally { if (e_6) throw e_6.error; }
+                                                    }
+                                                    // Cast to a JSON string so that it can be used in RLS expressions
+                                                    context[k] = JSON.stringify(context[k]);
+                                                }
+                                            }
+                                        }
+                                        catch (e_5_1) { e_5 = { error: e_5_1 }; }
+                                        finally {
+                                            try {
+                                                if (_b && !_b.done && (_e = _a.return)) _e.call(_a);
+                                            }
+                                            finally { if (e_5) throw e_5.error; }
+                                        }
+                                    }
+                                    txId = hashWithPrefix("yates_tx_", JSON.stringify(ctx));
+                                    hash = txId;
+                                    if (!batches[hash]) {
+                                        // Create a new batch for this role + context combination
+                                        batches[hash] = {
+                                            pgRole: pgRole,
+                                            context: context,
+                                            requests: [],
+                                        };
+                                        // make sure, that we only tick once at a time
+                                        if (!tickActive) {
+                                            tickActive = true;
+                                            process.nextTick(function () {
+                                                dispatchBatches();
+                                                tickActive = false;
+                                            });
+                                        }
+                                    }
+                                    __internalParams = params.__internalParams;
+                                    // Add the request to the batch, and return a promise that will be resolved or rejected in dispatchBatches
+                                    return [2 /*return*/, new Promise(function (resolve, reject) {
+                                            batches[hash].requests.push({
+                                                params: __internalParams,
+                                                query: query,
+                                                args: args,
+                                                resolve: resolve,
+                                                reject: reject,
+                                            });
+                                        }).catch(function (e) {
+                                            var _a;
+                                            // Normalize RLS errors to make them a bit more readable.
+                                            if ((_a = e.message) === null || _a === void 0 ? void 0 : _a.includes("new row violates row-level security policy for table")) {
+                                                throw new Error("You do not have permission to perform this action: ".concat(model, ".").concat(operation, "(...)"));
+                                            }
+                                            throw e;
+                                        })];
+                                });
+                            });
+                        },
+                    },
+                },
+            });
+            return client;
+        };
+        this.setRLS = function (prisma, table, roleName, slug, ability) { return __awaiter(_this, void 0, void 0, function () {
+            var operation, rawExpression, description;
+            var _this = this;
+            return __generator(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        operation = ability.operation, rawExpression = ability.expression, description = ability.description;
+                        if (!rawExpression) {
+                            throw new Error("Expression must be defined for RLS abilities");
                         }
-                    });
-                };
-                _f = abilities[model];
-                _g = [];
-                for (_h in _f)
-                    _g.push(_h);
-                _j = 0;
-                _r.label = 9;
-            case 9:
-                if (!(_j < _g.length)) return [3 /*break*/, 12];
-                _h = _g[_j];
-                if (!(_h in _f)) return [3 /*break*/, 11];
-                slug = _h;
-                return [5 /*yield**/, _loop_2(slug)];
-            case 10:
-                _r.sent();
-                _r.label = 11;
-            case 11:
-                _j++;
-                return [3 /*break*/, 9];
-            case 12:
-                _i++;
-                return [3 /*break*/, 7];
-            case 13:
-                _loop_3 = function (key) {
-                    var role, wildCardAbilities, roleAbilities, rlsRoles, userRoles, oldRoles, policies;
-                    return __generator(this, function (_t) {
-                        switch (_t.label) {
-                            case 0:
-                                role = (0, exports.createRoleName)(key);
-                                return [4 /*yield*/, prisma.$executeRawUnsafe("\n\t\t\tdo\n\t\t\t$$\n\t\t\tbegin\n\t\t\tif not exists (select * from pg_catalog.pg_roles where rolname = '".concat(role, "') then \n\t\t\t\tcreate role ").concat(role, ";\n\t\t\tend if;\n\t\t\tend\n\t\t\t$$\n\t\t\t;\n\t\t"))];
-                            case 1:
-                                _t.sent();
-                                wildCardAbilities = (0, flatMap_1.default)(abilities, function (model, modelName) {
-                                    return (0, map_1.default)(model, function (_params, slug) {
-                                        return (0, exports.createAbilityName)(modelName, slug);
-                                    });
+                        // Take a lock and run the RLS setup in a transaction to prevent conflicts
+                        // in a multi-server environment
+                        return [4 /*yield*/, prisma.$transaction(function (tx) { return __awaiter(_this, void 0, void 0, function () {
+                                var policyName, existingAbilities, existingAbility, shouldUpdateAbilityTable, expression, expression;
+                                return __generator(this, function (_a) {
+                                    switch (_a.label) {
+                                        case 0: return [4 /*yield*/, takeLock(tx)];
+                                        case 1:
+                                            _a.sent();
+                                            policyName = roleName;
+                                            return [4 /*yield*/, tx.$queryRaw(templateObject_3 || (templateObject_3 = __makeTemplateObject(["\n\t\t\t\tselect * from _yates._yates_abilities where ability_model = ", " and ability_policy_name = ", "\n\t\t\t"], ["\n\t\t\t\tselect * from _yates._yates_abilities where ability_model = ", " and ability_policy_name = ", "\n\t\t\t"])), table, policyName)];
+                                        case 2:
+                                            existingAbilities = _a.sent();
+                                            existingAbility = existingAbilities[0];
+                                            shouldUpdateAbilityTable = false;
+                                            if (!!existingAbility) return [3 /*break*/, 8];
+                                            debug("Creating RLS policy for", roleName, "on", table, "for", operation);
+                                            return [4 /*yield*/, (0, expressions_1.expressionToSQL)(rawExpression, table)];
+                                        case 3:
+                                            expression = _a.sent();
+                                            if (!(operation === "INSERT")) return [3 /*break*/, 5];
+                                            return [4 /*yield*/, tx.$queryRawUnsafe("\n\t\t\t\t\tCREATE POLICY ".concat(policyName, " ON \"public\".\"").concat(table, "\" FOR ").concat(operation, " TO ").concat(roleName, " WITH CHECK (").concat(expression, ");\n\t\t\t\t"))];
+                                        case 4:
+                                            _a.sent();
+                                            return [3 /*break*/, 7];
+                                        case 5: return [4 /*yield*/, tx.$queryRawUnsafe("\n\t\t\t\t\tCREATE POLICY ".concat(policyName, " ON \"public\".\"").concat(table, "\" FOR ").concat(operation, " TO ").concat(roleName, " USING (").concat(expression, ");\n\t\t\t\t"))];
+                                        case 6:
+                                            _a.sent();
+                                            _a.label = 7;
+                                        case 7:
+                                            shouldUpdateAbilityTable = true;
+                                            return [3 /*break*/, 14];
+                                        case 8:
+                                            if (!(existingAbility.ability_expression !== rawExpression.toString())) return [3 /*break*/, 14];
+                                            debug("Updating RLS policy for", roleName, "on", table, "for", operation);
+                                            return [4 /*yield*/, (0, expressions_1.expressionToSQL)(rawExpression, table)];
+                                        case 9:
+                                            expression = _a.sent();
+                                            if (!(operation === "INSERT")) return [3 /*break*/, 11];
+                                            return [4 /*yield*/, tx.$queryRawUnsafe("\n\t\t\t\t\tALTER POLICY ".concat(policyName, " ON \"public\".\"").concat(table, "\" TO ").concat(roleName, " WITH CHECK (").concat(expression, ");\n\t\t\t\t"))];
+                                        case 10:
+                                            _a.sent();
+                                            return [3 /*break*/, 13];
+                                        case 11: return [4 /*yield*/, tx.$queryRawUnsafe("\n\t\t\t\t\tALTER POLICY ".concat(policyName, " ON \"public\".\"").concat(table, "\" TO ").concat(roleName, " USING (").concat(expression, ");\n\t\t\t\t"))];
+                                        case 12:
+                                            _a.sent();
+                                            _a.label = 13;
+                                        case 13:
+                                            shouldUpdateAbilityTable = true;
+                                            _a.label = 14;
+                                        case 14:
+                                            if (!shouldUpdateAbilityTable) return [3 /*break*/, 16];
+                                            return [4 /*yield*/, upsertAbility(tx, {
+                                                    ability_model: table,
+                                                    ability_name: slug,
+                                                    ability_policy_name: policyName,
+                                                    ability_description: description !== null && description !== void 0 ? description : "",
+                                                    ability_operation: operation,
+                                                    // We store the string representation of the expression so that
+                                                    // we can compare it later without having to recompute the SQL
+                                                    ability_expression: rawExpression.toString(),
+                                                })];
+                                        case 15:
+                                            _a.sent();
+                                            _a.label = 16;
+                                        case 16: return [2 /*return*/];
+                                    }
                                 });
-                                roleAbilities = roles[key];
-                                rlsRoles = roleAbilities === "*"
-                                    ? wildCardAbilities
-                                    : roleAbilities.map(function (ability) {
+                            }); })];
+                    case 1:
+                        // Take a lock and run the RLS setup in a transaction to prevent conflicts
+                        // in a multi-server environment
+                        _a.sent();
+                        return [2 /*return*/];
+                }
+            });
+        }); };
+        this.createRoles = function (_a) { return __awaiter(_this, [_a], void 0, function (_b) {
+            var runtimeDataModel, models, diff, defaultAbilities, abilities, models_2, models_2_1, model, ability, operation, roles, pgRoles, existingAbilities, pgPolicies, migratedAbilities, _c, _d, _e, _i, model, table, _loop_2, this_1, _f, _g, _h, _j, slug, _loop_3, this_2, _k, _l, _m, _o, key;
+            var e_7, _p;
+            var _this = this;
+            var _q;
+            var customAbilities = _b.customAbilities, getRoles = _b.getRoles;
+            return __generator(this, function (_r) {
+                switch (_r.label) {
+                    case 0: return [4 /*yield*/, this.ensureDatabaseScope()];
+                    case 1:
+                        _r.sent();
+                        runtimeDataModel = this.inspectRunTimeDataModel();
+                        models = Object.keys(runtimeDataModel.models).map(function (m) { return runtimeDataModel.models[m].dbName || m; });
+                        if (customAbilities) {
+                            diff = (0, difference_1.default)(Object.keys(customAbilities), models);
+                            if (diff.length) {
+                                throw new Error("Invalid models in custom abilities: ".concat(diff.join(", ")));
+                            }
+                        }
+                        defaultAbilities = this.getDefaultAbilities(models);
+                        abilities = (0, cloneDeep_1.default)(defaultAbilities);
+                        try {
+                            for (models_2 = __values(models), models_2_1 = models_2.next(); !models_2_1.done; models_2_1 = models_2.next()) {
+                                model = models_2_1.value;
+                                if (customAbilities === null || customAbilities === void 0 ? void 0 : customAbilities[model]) {
+                                    for (ability in customAbilities[model]) {
+                                        operation = 
                                         // biome-ignore lint/style/noNonNullAssertion: TODO fix this
-                                        return (0, exports.createAbilityName)(ability.model, ability.slug);
-                                    });
-                                // Note: We need to GRANT all on schema public so that we can resolve relation queries with prisma, as they will sometimes use a join table.
-                                // This is not ideal, but because we are using RLS, it's not a security risk. Any table with RLS also needs a corresponding policy for the role to have access.
-                                return [4 /*yield*/, prisma.$transaction([
-                                        takeLock(prisma),
-                                        prisma.$executeRawUnsafe("GRANT ALL ON ALL TABLES IN SCHEMA public TO ".concat(role, ";")),
-                                        prisma.$executeRawUnsafe("\n\t\t\t\tGRANT ALL ON ALL SEQUENCES IN SCHEMA public TO ".concat(role, ";\n\t\t\t")),
-                                        prisma.$executeRawUnsafe("\n\t\t\t\tGRANT ALL ON SCHEMA public TO ".concat(role, ";\n\t\t\t")),
-                                        prisma.$queryRawUnsafe("GRANT ".concat(rlsRoles.join(", "), " TO ").concat(role)),
-                                    ])];
-                            case 2:
-                                // Note: We need to GRANT all on schema public so that we can resolve relation queries with prisma, as they will sometimes use a join table.
-                                // This is not ideal, but because we are using RLS, it's not a security risk. Any table with RLS also needs a corresponding policy for the role to have access.
-                                _t.sent();
-                                return [4 /*yield*/, prisma.$queryRawUnsafe("\n\t\t\tWITH RECURSIVE cte AS (\n\t\t\t\tSELECT oid FROM pg_roles where rolname = '".concat(role, "'\n\t\t\t\tUNION ALL\n\t\t\t\tSELECT m.roleid\n\t\t\t\tFROM   cte\n\t\t\t\tJOIN   pg_auth_members m ON m.member = cte.oid\n\t\t\t\t)\n\t\t\tSELECT oid, oid::regrole::text AS rolename FROM cte where oid::regrole::text != '").concat(role, "'; \n\t "))];
-                            case 3:
-                                userRoles = _t.sent();
-                                oldRoles = userRoles
-                                    .filter(function (_a) {
-                                    var rolename = _a.rolename;
-                                    return !rlsRoles.includes(rolename);
-                                })
-                                    .map(function (_a) {
-                                    var rolename = _a.rolename;
-                                    return rolename;
-                                });
-                                if (!oldRoles.length) return [3 /*break*/, 8];
-                                // Now revoke old roles from the user role
-                                debug("Revoking old roles", oldRoles.join(", "));
-                                return [4 /*yield*/, prisma.$executeRawUnsafe("REVOKE ".concat(oldRoles.join(", "), " FROM ").concat(role))];
-                            case 4:
-                                _t.sent();
-                                return [4 /*yield*/, prisma.$queryRawUnsafe("SELECT * FROM pg_catalog.pg_policies WHERE policyname IN (".concat(oldRoles
-                                        .map(function (or) { return "'".concat(or, "'"); })
-                                        .join(", "), ")"))];
-                            case 5:
-                                policies = _t.sent();
-                                return [4 /*yield*/, prisma.$transaction(__spreadArray([
-                                        takeLock(prisma)
-                                    ], __read(policies.map(function (oldPolicy) {
-                                        return prisma.$executeRawUnsafe("DROP POLICY ".concat(oldPolicy.policyname, " ON \"").concat(oldPolicy.tablename, "\""));
-                                    })), false))];
-                            case 6:
-                                _t.sent();
-                                debug("Revoked old rows from ability table", oldRoles.join(", "));
-                                return [4 /*yield*/, prisma.$executeRawUnsafe("DELETE FROM _yates._yates_abilities WHERE ability_policy_name IN (".concat(oldRoles
-                                        .map(function (or) { return "'".concat(or, "'"); })
-                                        .join(", "), ")"))];
-                            case 7:
-                                _t.sent();
-                                _t.label = 8;
-                            case 8: return [2 /*return*/];
+                                        (_q = customAbilities[model][ability]) === null || _q === void 0 ? void 0 : _q.operation;
+                                        if (!operation)
+                                            continue;
+                                        // biome-ignore lint/style/noNonNullAssertion: TODO fix this
+                                        abilities[model][ability] = __assign(__assign({}, customAbilities[model][ability]), { operation: operation, model: model, slug: ability });
+                                    }
+                                }
+                            }
                         }
-                    });
-                };
-                _k = roles;
-                _l = [];
-                for (_m in _k)
-                    _l.push(_m);
-                _o = 0;
-                _r.label = 14;
-            case 14:
-                if (!(_o < _l.length)) return [3 /*break*/, 17];
-                _m = _l[_o];
-                if (!(_m in _k)) return [3 /*break*/, 16];
-                key = _m;
-                return [5 /*yield**/, _loop_3(key)];
-            case 15:
-                _r.sent();
-                _r.label = 16;
-            case 16:
-                _o++;
-                return [3 /*break*/, 14];
-            case 17: return [2 /*return*/];
+                        catch (e_7_1) { e_7 = { error: e_7_1 }; }
+                        finally {
+                            try {
+                                if (models_2_1 && !models_2_1.done && (_p = models_2.return)) _p.call(models_2);
+                            }
+                            finally { if (e_7) throw e_7.error; }
+                        }
+                        roles = getRoles(abilities);
+                        return [4 /*yield*/, this.prisma.$queryRawUnsafe("\n\t\t\tselect * from pg_catalog.pg_roles where rolname like 'yates%'\n\t\t")];
+                    case 2:
+                        pgRoles = _r.sent();
+                        return [4 /*yield*/, this.prisma.$queryRawUnsafe("\n\t\t\tselect * from _yates._yates_abilities;\n\t\t")];
+                    case 3:
+                        existingAbilities = _r.sent();
+                        if (!(existingAbilities.length === 0)) return [3 /*break*/, 6];
+                        debug('No existing abilities found, importing from "pg_policies" table');
+                        return [4 /*yield*/, this.prisma.$queryRawUnsafe("\n\t\t\t\tselect * from pg_catalog.pg_policies where policyname like 'yates%'\n\t\t\t")];
+                    case 4:
+                        pgPolicies = _r.sent();
+                        if (!pgPolicies.length) return [3 /*break*/, 6];
+                        migratedAbilities = pgPolicies.map(function (policy) {
+                            var _a, _b;
+                            return ({
+                                ability_model: policy.tablename,
+                                ability_name: policy.policyname,
+                                ability_policy_name: policy.policyname,
+                                ability_description: "",
+                                ability_operation: policy.cmd,
+                                ability_expression: (_b = (_a = policy.qual) !== null && _a !== void 0 ? _a : policy.with_check) !== null && _b !== void 0 ? _b : "",
+                            });
+                        });
+                        return [4 /*yield*/, this.prisma.$transaction(__spreadArray([
+                                takeLock(this.prisma)
+                            ], __read(migratedAbilities.map(function (ma) { return upsertAbility(_this.prisma, ma); })), false))];
+                    case 5:
+                        _r.sent();
+                        existingAbilities.push.apply(existingAbilities, __spreadArray([], __read(migratedAbilities), false));
+                        _r.label = 6;
+                    case 6:
+                        _c = abilities;
+                        _d = [];
+                        for (_e in _c)
+                            _d.push(_e);
+                        _i = 0;
+                        _r.label = 7;
+                    case 7:
+                        if (!(_i < _d.length)) return [3 /*break*/, 13];
+                        _e = _d[_i];
+                        if (!(_e in _c)) return [3 /*break*/, 12];
+                        model = _e;
+                        table = model;
+                        return [4 /*yield*/, this.prisma.$transaction([
+                                takeLock(this.prisma),
+                                this.prisma.$queryRawUnsafe("ALTER table \"".concat(table, "\" enable row level security;")),
+                            ])];
+                    case 8:
+                        _r.sent();
+                        _loop_2 = function (slug) {
+                            var ability, roleName;
+                            return __generator(this, function (_s) {
+                                switch (_s.label) {
+                                    case 0:
+                                        ability = 
+                                        // biome-ignore lint/style/noNonNullAssertion: TODO fix this
+                                        abilities[model][slug];
+                                        if (!VALID_OPERATIONS.includes(ability.operation)) {
+                                            throw new Error("Invalid operation: ".concat(ability.operation));
+                                        }
+                                        roleName = this_1.createAbilityName(model, slug);
+                                        if (!pgRoles.find(function (role) { return role.rolname === roleName; })) return [3 /*break*/, 1];
+                                        debug("Role already exists", roleName, model, slug);
+                                        return [3 /*break*/, 3];
+                                    case 1: return [4 /*yield*/, this_1.prisma.$transaction([
+                                            takeLock(this_1.prisma),
+                                            this_1.prisma.$queryRawUnsafe("\n\t\t\t\t\t\tdo\n\t\t\t\t\t\t$$\n\t\t\t\t\t\tbegin\n\t\t\t\t\t\tif not exists (select * from pg_catalog.pg_roles where rolname = '".concat(roleName, "') then \n\t\t\t\t\t\t\tcreate role ").concat(roleName, ";\n\t\t\t\t\t\tend if;\n\t\t\t\t\t\tend\n\t\t\t\t\t\t$$\n\t\t\t\t\t\t;\n\t\t\t\t\t")),
+                                            this_1.prisma.$queryRawUnsafe("\n\t\t\t\t\t\tGRANT ".concat(ability.operation, " ON \"").concat(table, "\" TO ").concat(roleName, ";\n\t\t\t\t\t")),
+                                        ])];
+                                    case 2:
+                                        _s.sent();
+                                        _s.label = 3;
+                                    case 3:
+                                        if (!ability.expression) return [3 /*break*/, 5];
+                                        return [4 /*yield*/, this_1.setRLS(this_1.prisma, table, roleName, slug, ability)];
+                                    case 4:
+                                        _s.sent();
+                                        _s.label = 5;
+                                    case 5: return [2 /*return*/];
+                                }
+                            });
+                        };
+                        this_1 = this;
+                        _f = abilities[model];
+                        _g = [];
+                        for (_h in _f)
+                            _g.push(_h);
+                        _j = 0;
+                        _r.label = 9;
+                    case 9:
+                        if (!(_j < _g.length)) return [3 /*break*/, 12];
+                        _h = _g[_j];
+                        if (!(_h in _f)) return [3 /*break*/, 11];
+                        slug = _h;
+                        return [5 /*yield**/, _loop_2(slug)];
+                    case 10:
+                        _r.sent();
+                        _r.label = 11;
+                    case 11:
+                        _j++;
+                        return [3 /*break*/, 9];
+                    case 12:
+                        _i++;
+                        return [3 /*break*/, 7];
+                    case 13:
+                        _loop_3 = function (key) {
+                            var role, wildCardAbilities, roleAbilities, rlsRoles, userRoles, oldRoles, policies;
+                            return __generator(this, function (_t) {
+                                switch (_t.label) {
+                                    case 0:
+                                        role = this_2.createRoleName(key);
+                                        return [4 /*yield*/, this_2.prisma.$executeRawUnsafe("\n\t\t\t\tdo\n\t\t\t\t$$\n\t\t\t\tbegin\n\t\t\t\tif not exists (select * from pg_catalog.pg_roles where rolname = '".concat(role, "') then \n\t\t\t\t\tcreate role ").concat(role, ";\n\t\t\t\tend if;\n\t\t\t\tend\n\t\t\t\t$$\n\t\t\t\t;\n\t\t\t"))];
+                                    case 1:
+                                        _t.sent();
+                                        wildCardAbilities = (0, flatMap_1.default)(defaultAbilities, function (model, modelName) {
+                                            return (0, map_1.default)(model, function (_params, slug) {
+                                                return _this.createAbilityName(modelName, slug);
+                                            });
+                                        });
+                                        roleAbilities = roles[key];
+                                        rlsRoles = roleAbilities === "*"
+                                            ? wildCardAbilities
+                                            : roleAbilities.map(function (ability) {
+                                                // biome-ignore lint/style/noNonNullAssertion: TODO fix this
+                                                return _this.createAbilityName(ability.model, ability.slug);
+                                            });
+                                        debug("Setting up role", key, role, "with abilities", rlsRoles.join(", "));
+                                        // Note: We need to GRANT all on schema public so that we can resolve relation queries with prisma, as they will sometimes use a join table.
+                                        // This is not ideal, but because we are using RLS, it's not a security risk. Any table with RLS also needs a corresponding policy for the role to have access.
+                                        return [4 /*yield*/, this_2.prisma.$transaction([
+                                                takeLock(this_2.prisma),
+                                                this_2.prisma.$executeRawUnsafe("GRANT ALL ON ALL TABLES IN SCHEMA public TO ".concat(role, ";")),
+                                                this_2.prisma.$executeRawUnsafe("\n\t\t\t\t\tGRANT ALL ON ALL SEQUENCES IN SCHEMA public TO ".concat(role, ";\n\t\t\t\t")),
+                                                this_2.prisma.$executeRawUnsafe("\n\t\t\t\t\tGRANT ALL ON SCHEMA public TO ".concat(role, ";\n\t\t\t\t")),
+                                                this_2.prisma.$queryRawUnsafe("GRANT ".concat(rlsRoles.join(", "), " TO ").concat(role)),
+                                            ])];
+                                    case 2:
+                                        // Note: We need to GRANT all on schema public so that we can resolve relation queries with prisma, as they will sometimes use a join table.
+                                        // This is not ideal, but because we are using RLS, it's not a security risk. Any table with RLS also needs a corresponding policy for the role to have access.
+                                        _t.sent();
+                                        return [4 /*yield*/, this_2.prisma.$queryRawUnsafe("\n\t\t\t\tWITH RECURSIVE cte AS (\n\t\t\t\t\tSELECT oid FROM pg_roles where rolname = '".concat(role, "'\n\t\t\t\t\tUNION ALL\n\t\t\t\t\tSELECT m.roleid\n\t\t\t\t\tFROM   cte\n\t\t\t\t\tJOIN   pg_auth_members m ON m.member = cte.oid\n\t\t\t\t\t)\n\t\t\t\tSELECT oid, oid::regrole::text AS rolename FROM cte where oid::regrole::text != '").concat(role, "'; \n\t\t"))];
+                                    case 3:
+                                        userRoles = _t.sent();
+                                        oldRoles = userRoles
+                                            .filter(function (_a) {
+                                            var rolename = _a.rolename;
+                                            return !rlsRoles.includes(rolename);
+                                        })
+                                            .map(function (_a) {
+                                            var rolename = _a.rolename;
+                                            return rolename;
+                                        });
+                                        if (!oldRoles.length) return [3 /*break*/, 8];
+                                        // Now revoke old roles from the user role
+                                        debug("Revoking old roles", key, role, oldRoles.join(", "));
+                                        return [4 /*yield*/, this_2.prisma.$executeRawUnsafe("REVOKE ".concat(oldRoles.join(", "), " FROM ").concat(role))];
+                                    case 4:
+                                        _t.sent();
+                                        return [4 /*yield*/, this_2.prisma.$queryRawUnsafe("SELECT * FROM pg_catalog.pg_policies WHERE policyname IN (".concat(oldRoles
+                                                .map(function (or) { return "'".concat(or, "'"); })
+                                                .join(", "), ")"))];
+                                    case 5:
+                                        policies = _t.sent();
+                                        return [4 /*yield*/, this_2.prisma.$transaction(__spreadArray([
+                                                takeLock(this_2.prisma)
+                                            ], __read(policies.map(function (oldPolicy) {
+                                                return _this.prisma.$executeRawUnsafe("DROP POLICY ".concat(oldPolicy.policyname, " ON \"").concat(oldPolicy.tablename, "\""));
+                                            })), false))];
+                                    case 6:
+                                        _t.sent();
+                                        debug("Revoked old rows from ability table", oldRoles.join(", "));
+                                        return [4 /*yield*/, this_2.prisma.$executeRawUnsafe("DELETE FROM _yates._yates_abilities WHERE ability_policy_name IN (".concat(oldRoles
+                                                .map(function (or) { return "'".concat(or, "'"); })
+                                                .join(", "), ")"))];
+                                    case 7:
+                                        _t.sent();
+                                        _t.label = 8;
+                                    case 8: return [2 /*return*/];
+                                }
+                            });
+                        };
+                        this_2 = this;
+                        _k = roles;
+                        _l = [];
+                        for (_m in _k)
+                            _l.push(_m);
+                        _o = 0;
+                        _r.label = 14;
+                    case 14:
+                        if (!(_o < _l.length)) return [3 /*break*/, 17];
+                        _m = _l[_o];
+                        if (!(_m in _k)) return [3 /*break*/, 16];
+                        key = _m;
+                        return [5 /*yield**/, _loop_3(key)];
+                    case 15:
+                        _r.sent();
+                        _r.label = 16;
+                    case 16:
+                        _o++;
+                        return [3 /*break*/, 14];
+                    case 17: return [2 /*return*/];
+                }
+            });
+        }); };
+        this.inspectDBRoles = function (role) { return __awaiter(_this, void 0, void 0, function () {
+            var hashedRoleName, roles;
+            return __generator(this, function (_a) {
+                switch (_a.label) {
+                    case 0: return [4 /*yield*/, this.ensureDatabaseScope()];
+                    case 1:
+                        _a.sent();
+                        hashedRoleName = this.createRoleName(role);
+                        return [4 /*yield*/, this.prisma.$queryRawUnsafe("\n        WITH RECURSIVE role_tree AS(\n            --Start from your role\n            SELECT \n                r.oid,\n                    r.rolname\n            FROM pg_roles r\n            WHERE r.rolname = '".concat(hashedRoleName, "'\n\n            UNION\n\n            --Walk \"upwards\": all parent roles granted to it\n            SELECT \n                parent.oid,\n                    parent.rolname\n            FROM pg_auth_members m\n            JOIN role_tree rt\n            ON m.member = rt.oid\n            JOIN pg_roles parent\n            ON parent.oid = m.roleid\n                )\n                SELECT\n                p.tablename,\n                    p.policyname,\n                    p.cmd,\n                    p.roles:: text[]          AS policy_roles,\n                        rt.rolname                  AS matched_role\n        FROM pg_policies p\n        JOIN role_tree rt\n                ON(\n                    p.roles IS NULL\n            OR array_length(p.roles, 1) = 0\n            OR rt.rolname = ANY(p.roles:: text[])\n                )\n        WHERE p.schemaname = 'public'\n        ORDER BY p.policyname, matched_role;\n        "))];
+                    case 2:
+                        roles = _a.sent();
+                        return [2 /*return*/, roles];
+                }
+            });
+        }); };
+        this.inspectRunTimeDataModel = function () {
+            // See https://github.com/prisma/prisma/discussions/14777
+            // We are reaching into the prisma internals to get the data model.
+            // This is a bit sketchy, but we can get the internal type definition from the runtime library
+            // and there is even a test case in prisma that checks that this value is exported
+            // See https://github.com/prisma/prisma/blob/5.1.0/packages/client/tests/functional/extensions/pdp.ts#L51
+            // This is a private API, so not much we can do about the cast
+            var runtimeDataModel = _this.prisma
+                ._runtimeDataModel;
+            return runtimeDataModel;
+        };
+    }
+    // @ts-ignore
+    Yates.prototype.getBatchId = function (query) {
+        if (query.action !== "findUnique" && query.action !== "findUniqueOrThrow") {
+            return undefined;
         }
-    });
-}); };
-exports.createRoles = createRoles;
+        var parts = [];
+        if (query.modelName) {
+            parts.push(query.modelName);
+        }
+        if (query.query.arguments) {
+            parts.push(this.buildKeysString(query.query.arguments));
+        }
+        parts.push(this.buildKeysString(query.query.selection));
+        return parts.join("");
+    };
+    Yates.prototype.buildKeysString = function (obj) {
+        var _this = this;
+        var keysArray = Object.keys(obj)
+            .sort()
+            .map(function (key) {
+            // @ts-ignore
+            var value = obj[key];
+            if (typeof value === "object" && value !== null) {
+                return "(".concat(key, " ").concat(_this.buildKeysString(value), ")");
+            }
+            return key;
+        });
+        return "(".concat(keysArray.join(" "), ")");
+    };
+    return Yates;
+}());
+exports.Yates = Yates;
 /**
  * Creates an extended client that sets contextual parameters and user role on every query
  **/
 var setup = function (params) { return __awaiter(void 0, void 0, void 0, function () {
-    var start, prisma, customAbilities, getRoles, getContext, client;
+    var start, prisma, customAbilities, getRoles, getContext, yates, client;
     return __generator(this, function (_a) {
         switch (_a.label) {
             case 0:
                 start = performance.now();
                 prisma = params.prisma, customAbilities = params.customAbilities, getRoles = params.getRoles, getContext = params.getContext;
-                return [4 /*yield*/, (0, exports.createRoles)({
-                        prisma: prisma,
+                yates = new Yates(prisma);
+                return [4 /*yield*/, yates.init()];
+            case 1:
+                _a.sent();
+                return [4 /*yield*/, yates.createRoles({
                         customAbilities: customAbilities,
                         getRoles: getRoles,
                     })];
-            case 1:
+            case 2:
                 _a.sent();
-                client = (0, exports.createClient)(prisma, getContext, params.options);
+                client = yates.createClient(getContext, params.options);
                 debug("Setup completed in", performance.now() - start, "ms");
                 return [2 /*return*/, client];
         }