Merge pull request #92 from cerebruminc/lucianbuzzo/improve-startup-perf

README.md

@@ -154,7 +154,7 @@ When defining an ability you need to provide the following properties:
 
 ### Debug
 
-To run Yates in debug mode, set the `YATES_DEBUG` environment variable to `true`.
+To run Yates in debug mode, use the environment variable `DEBUG=yates`.
 
 ## Known limitations
 

package.json

@@ -38,6 +38,8 @@
     "uuid": "^9.0.0"
   },
   "dependencies": {
+    "@types/debug": "^4.1.12",
+    "debug": "^4.3.4",
     "lodash": "^4.17.21",
     "node-sql-parser": "^4.12.0",
     "type-fest": "^4.10.3"
@@ -46,4 +48,4 @@
     "@prisma/client": "^5.0.0",
     "prisma": "^5.0.0"
   }
-}
+}

src/index.ts

@@ -1,5 +1,6 @@
 import * as crypto from "crypto";
 import { Prisma, PrismaClient } from "@prisma/client";
+import logger from "debug";
 import difference from "lodash/difference";
 import flatMap from "lodash/flatMap";
 import map from "lodash/map";
@@ -8,17 +9,22 @@ import { Expression, RuntimeDataModel, expressionToSQL } from "./expressions";
 
 const VALID_OPERATIONS = ["SELECT", "UPDATE", "INSERT", "DELETE"] as const;
 
-const DEBUG = process.env.YATES_DEBUG === "1";
-
-const debug = (...args: unknown[]) => {
-	if (DEBUG) {
-		console.log(...args);
-	}
-};
+const debug = logger("yates");
 
 type Operation = (typeof VALID_OPERATIONS)[number];
 export type Models = Prisma.ModelName;
 
+interface PgPolicy {
+	policyname: string;
+	tablename: string;
+	qual: string | null;
+	with_check: string | null;
+}
+
+interface PgRole {
+	rolname: string;
+}
+
 interface ClientOptions {
 	/** The maximum amount of time Yates will wait to acquire a transaction from the database. The default value is 30 seconds. */
 	txMaxWait?: number;
@@ -244,6 +250,7 @@ export const createClient = (
 
 const setRLS = async <ContextKeys extends string, YModel extends Models>(
 	prisma: PrismaClient,
+	pgPolicies: PgPolicy[],
 	table: string,
 	roleName: string,
 	operation: Operation,
@@ -255,18 +262,32 @@ const setRLS = async <ContextKeys extends string, YModel extends Models>(
 
 	// Check if RLS exists
 	const policyName = roleName;
-	// biome-ignore lint/suspicious/noExplicitAny: TODO fix this, by providing the correct type for the catalog
-	const rows: any[] = await prisma.$queryRawUnsafe(`
-		select * from pg_catalog.pg_policies where tablename = '${table}' AND policyname = '${policyName}';
-	`);
+	const rows = pgPolicies.filter(
+		(row) => row.tablename === table && row.policyname === policyName,
+	);
 
 	debug("Creating RLS policy", policyName);
 	debug("On table", table);
 	debug("For operation", operation);
 	debug("To role", roleName);
 	debug("With expression", expression);
 
+	// If the expression is a plain "true" it is not wrapped in parentheses
+	const normalizedExpression =
+		expression === "true"
+			? expression
+			: `(${expression.replace(/(\r\n|\n|\r)/gm, "")})`;
+
+	// If the op is INSERT, the expression is in the "with_check" column
+	const normalizedQual =
+		operation === "INSERT"
+			? rows?.[0]?.with_check?.replace(/(\r\n|\n|\r)/gm, "")
+			: rows?.[0]?.qual?.replace(/(\r\n|\n|\r)/gm, "");
+
 	// IF RLS doesn't exist or expression is different, set RLS
+	// Note that PG performs various optimizations and mods to the expression
+	// on write so we need to normalize it before comparing, and even then it
+	// might not be exactly the same
 	if (rows.length === 0) {
 		// If the operation is an insert or update, we need to use a different syntax as the "WITH CHECK" expression is used.
 		if (operation === "INSERT") {
@@ -278,7 +299,7 @@ const setRLS = async <ContextKeys extends string, YModel extends Models>(
         CREATE POLICY ${policyName} ON "public"."${table}" FOR ${operation} TO ${roleName} USING (${expression});
       `);
 		}
-	} else if (rows[0].qual !== expression) {
+	} else if (normalizedQual !== normalizedExpression) {
 		if (operation === "INSERT") {
 			await prisma.$queryRawUnsafe(`
         ALTER POLICY ${policyName} ON "public"."${table}" TO ${roleName} WITH CHECK (${expression});
@@ -381,6 +402,13 @@ export const createRoles = async <
 
 	const roles = getRoles(abilities as T);
 
+	const pgRoles: PgRole[] = await prisma.$queryRawUnsafe(`
+			select * from pg_catalog.pg_roles
+		`);
+	const pgPolicies: PgPolicy[] = await prisma.$queryRawUnsafe(`
+		select * from pg_catalog.pg_policies;
+	`);
+
 	// For each of the models and abilities, create a role and a corresponding RLS policy
 	// We can then mix & match these roles to create a user's permissions by granting them to a user role (like SUPER_ADMIN)
 	for (const model in abilities) {
@@ -405,9 +433,14 @@ export const createRoles = async <
 			const roleName = createAbilityName(model, slug);
 
 			// Check if role already exists
-			await prisma.$transaction([
-				takeLock(prisma),
-				prisma.$queryRawUnsafe(`
+			if (
+				pgRoles.find((role: { rolname: string }) => role.rolname === roleName)
+			) {
+				debug("Role already exists", roleName);
+			} else {
+				await prisma.$transaction([
+					takeLock(prisma),
+					prisma.$queryRawUnsafe(`
 					do
 					$$
 					begin
@@ -418,14 +451,16 @@ export const createRoles = async <
 					$$
 					;
 				`),
-				prisma.$queryRawUnsafe(`
+					prisma.$queryRawUnsafe(`
 					GRANT ${ability.operation} ON "${table}" TO ${roleName};
 				`),
-			]);
+				]);
+			}
 
 			if (ability.expression) {
 				await setRLS(
 					prisma,
+					pgPolicies,
 					table,
 					roleName,
 					ability.operation,
@@ -554,6 +589,8 @@ export const setup = async <
 >(
 	params: SetupParams<ContextKeys, YModels, K>,
 ) => {
+	const start = performance.now();
+
 	const { prisma, customAbilities, getRoles, getContext } = params;
 	await createRoles<ContextKeys, YModels, K>({
 		prisma,
@@ -562,5 +599,7 @@ export const setup = async <
 	});
 	const client = createClient(prisma, getContext, params.options);
 
+	debug("Setup completed in", performance.now() - start, "ms");
+
 	return client;
 };