feat: add test for multiple Yates clients operating in the same DB

Author: LucianBuzzo

.gitignore

@@ -257,4 +257,5 @@ tags
 
 dist/
 .npmrc
-*.tgz
+*.tgz
+prisma/generated

biome.json

@@ -1,7 +1,7 @@
 {
 	"$schema": "https://biomejs.dev/schemas/1.5.3/schema.json",
 	"files": {
-		"ignore": ["dist", "coverage", "package.json"]
+		"ignore": ["dist", "coverage", "package.json", "prisma/generated"]
 	},
 	"organizeImports": {
 		"enabled": true

docker-compose.yml

@@ -17,14 +17,17 @@ services:
       - internal
     environment:
       - DATABASE_URL=postgresql://postgres:postgres@db:5432/yates?connection_limit=30
+      - DATABASE_URL_2=postgresql://postgres:postgres@db:5432/yates_2?connection_limit=30
 
   db:
     image: postgres:11
+    volumes:
+      - ./docker-init:/docker-entrypoint-initdb.d
     restart: always
     environment:
       POSTGRES_USER: postgres
       POSTGRES_PASSWORD: postgres
-      POSTGRES_DB: yates
+      POSTGRES_MULTIPLE_DATABASES: yates,yates_2
     ports:
       - 5432:5432
     networks:

docker-init/setup-databases.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+# Borrowed from https://github.com/mrts/docker-postgresql-multiple-databases
+# This script creates multiple Postgres databases, so we can test multi-tenant setups
+
+set -e
+set -u
+
+function create_user_and_database() {
+	local database=$1
+	echo "  Creating user and database '$database'"
+	psql -v ON_ERROR_STOP=1 --username "$POSTGRES_USER" <<-EOSQL
+	    CREATE USER $database;
+	    CREATE DATABASE $database;
+	    GRANT ALL PRIVILEGES ON DATABASE $database TO $database;
+EOSQL
+}
+
+if [ -n "$POSTGRES_MULTIPLE_DATABASES" ]; then
+	echo "Multiple database creation requested: $POSTGRES_MULTIPLE_DATABASES"
+	for db in $(echo $POSTGRES_MULTIPLE_DATABASES | tr ',' ' '); do
+		create_user_and_database $db
+	done
+	echo "Multiple databases created"
+fi

package.json

@@ -17,7 +17,7 @@
 		"test:types": "tsc --noEmit",
 		"test:integration": "jest --runInBand test/integration",
 		"test:compose:integration": "docker compose -f docker-compose.yml --profile with-sut up db sut --exit-code-from sut",
-		"setup": "prisma generate && prisma migrate dev",
+		"setup": "prisma generate --schema prisma/schema.prisma && prisma migrate dev --schema prisma/schema.prisma && prisma generate --schema prisma/schema2.prisma && prisma migrate dev --schema prisma/schema2.prisma",
 		"prepublishOnly": "npm run build"
 	},
 	"author": "Cerebrum <[email protected]> (https://cerebrum.com)",

prisma/migrations/20251115091853_init/migration.sql

@@ -0,0 +1,54 @@
+/*
+  Warnings:
+
+  - You are about to drop the `Account` table. If the table is not empty, all the data it contains will be lost.
+  - You are about to drop the `Hat` table. If the table is not empty, all the data it contains will be lost.
+  - You are about to drop the `Item` table. If the table is not empty, all the data it contains will be lost.
+  - You are about to drop the `Organization` table. If the table is not empty, all the data it contains will be lost.
+  - You are about to drop the `Role` table. If the table is not empty, all the data it contains will be lost.
+  - You are about to drop the `RoleAssignment` table. If the table is not empty, all the data it contains will be lost.
+  - You are about to drop the `Tag` table. If the table is not empty, all the data it contains will be lost.
+  - You are about to drop the `_PostToTag` table. If the table is not empty, all the data it contains will be lost.
+
+*/
+-- DropForeignKey
+ALTER TABLE "Hat" DROP CONSTRAINT "Hat_userId_fkey";
+
+-- DropForeignKey
+ALTER TABLE "RoleAssignment" DROP CONSTRAINT "RoleAssignment_organizationId_fkey";
+
+-- DropForeignKey
+ALTER TABLE "RoleAssignment" DROP CONSTRAINT "RoleAssignment_roleId_fkey";
+
+-- DropForeignKey
+ALTER TABLE "RoleAssignment" DROP CONSTRAINT "RoleAssignment_userId_fkey";
+
+-- DropForeignKey
+ALTER TABLE "_PostToTag" DROP CONSTRAINT "_PostToTag_A_fkey";
+
+-- DropForeignKey
+ALTER TABLE "_PostToTag" DROP CONSTRAINT "_PostToTag_B_fkey";
+
+-- DropTable
+DROP TABLE "Account";
+
+-- DropTable
+DROP TABLE "Hat";
+
+-- DropTable
+DROP TABLE "Item";
+
+-- DropTable
+DROP TABLE "Organization";
+
+-- DropTable
+DROP TABLE "Role";
+
+-- DropTable
+DROP TABLE "RoleAssignment";
+
+-- DropTable
+DROP TABLE "Tag";
+
+-- DropTable
+DROP TABLE "_PostToTag";

prisma/schema2.prisma

@@ -0,0 +1,29 @@
+// This particular schema has a subset of models compared to schema.prisma
+// It is used to test the behaviour of Yates when multiple schemas are being used on the same database server
+datasource db {
+    provider = "postgresql"
+    url      = env("DATABASE_URL_2")
+}
+
+generator client {
+    provider = "prisma-client-js"
+    output   = "./generated/client2"
+}
+
+model User {
+    id        String   @id @default(uuid())
+    createdAt DateTime @default(now())
+    email     String   @unique
+    name      String?
+    posts     Post[]
+}
+
+model Post {
+    id        Int      @id @default(autoincrement())
+    createdAt DateTime @default(now())
+    updatedAt DateTime @updatedAt
+    published Boolean  @default(false)
+    title     String   @db.VarChar(255)
+    author    User?    @relation(fields: [authorId], references: [id])
+    authorId  String?
+}

src/index.ts

@@ -185,7 +185,7 @@ export const sanitizeSlug = (slug: string) =>
 		.replace(/-/g, "_")
 		.replace(/[^a-z0-9_]/gi, "");
 
-class Yates {
+export class Yates {
 	private databaseScope: string | null = null;
 
 	constructor(private prisma: PrismaClient) {}
@@ -229,6 +229,8 @@ class Yates {
 
 		const currentDatabase = result[0]?.current_database;
 
+		debug("Current database for Yates:", currentDatabase);
+
 		if (!currentDatabase) {
 			throw new Error(
 				"Failed to determine the current database for scoping Yates roles.",
@@ -668,14 +670,7 @@ class Yates {
 	}) => {
 		await this.ensureDatabaseScope();
 
-		// 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
-		const runtimeDataModel = (this.prisma as any)
-			._runtimeDataModel as RuntimeDataModel;
+		const runtimeDataModel = this.inspectRunTimeDataModel();
 		const models = Object.keys(runtimeDataModel.models).map(
 			(m) => runtimeDataModel.models[m].dbName || m,
 		) as Models[];
@@ -903,6 +898,73 @@ class Yates {
 			}
 		}
 	};
+
+	inspectDBRoles = async (role: string) => {
+		await this.ensureDatabaseScope();
+		const hashedRoleName = this.createRoleName(role);
+
+		// Load all policies for the role
+		const roles = await this.prisma.$queryRawUnsafe<
+			{
+				tablename: string;
+				policyname: string;
+				cmd: string;
+				policy_roles: string[];
+				matched_role: string[];
+			}[]
+		>(`
+        WITH RECURSIVE role_tree AS(
+            --Start from your role
+            SELECT 
+                r.oid,
+                    r.rolname
+            FROM pg_roles r
+            WHERE r.rolname = '${hashedRoleName}'
+
+            UNION
+
+            --Walk "upwards": all parent roles granted to it
+            SELECT 
+                parent.oid,
+                    parent.rolname
+            FROM pg_auth_members m
+            JOIN role_tree rt
+            ON m.member = rt.oid
+            JOIN pg_roles parent
+            ON parent.oid = m.roleid
+                )
+                SELECT
+                p.tablename,
+                    p.policyname,
+                    p.cmd,
+                    p.roles:: text[]          AS policy_roles,
+                        rt.rolname                  AS matched_role
+        FROM pg_policies p
+        JOIN role_tree rt
+                ON(
+                    p.roles IS NULL
+            OR array_length(p.roles, 1) = 0
+            OR rt.rolname = ANY(p.roles:: text[])
+                )
+        WHERE p.schemaname = 'public'
+        ORDER BY p.policyname, matched_role;
+        `);
+
+		return roles;
+	};
+
+	inspectRunTimeDataModel = (): RuntimeDataModel => {
+		// 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
+		const runtimeDataModel = (this.prisma as any)
+			._runtimeDataModel as RuntimeDataModel;
+
+		return runtimeDataModel;
+	};
 }
 
 /**

test/integration/multi-tenant-db.spec.ts

@@ -0,0 +1,52 @@
+import { PrismaClient } from "@prisma/client";
+import { PrismaClient as PrismaClient2 } from "../../prisma/generated/client2";
+import { Yates, setup } from "../../src";
+
+describe("Multi-tenant database tests", () => {
+	it("should not overwrite data between tenants", async () => {
+		const rootClient1 = new PrismaClient();
+		const rootClient2 = new PrismaClient2();
+
+		const role = "ADMIN";
+
+		const _client1 = await setup({
+			prisma: rootClient1,
+			getRoles(_abilities) {
+				return {
+					[role]: "*",
+				};
+			},
+			getContext: () => ({
+				role,
+			}),
+		});
+
+		const yates1 = new Yates(rootClient1);
+		const roles1 = await yates1.inspectDBRoles(role);
+		const models1 = Object.keys(yates1.inspectRunTimeDataModel().models);
+		// Expect to see 4 roles (CRUD) for each defined model when using wildcard abilities (*)
+		expect(roles1.length).toBe(models1.length * 4);
+
+		const _client2 = await setup({
+			prisma: rootClient2 as PrismaClient,
+			getRoles(abilities) {
+				return {
+					[role]: [abilities.User.create, abilities.User.read],
+				};
+			},
+			getContext: () => ({
+				role,
+			}),
+		});
+
+		const yates2 = new Yates(rootClient2 as PrismaClient);
+		const roles2 = await yates2.inspectDBRoles(role);
+		// Expect to see 2 roles (CREATE, READ) for User model only
+		expect(roles2.length).toBe(2);
+
+		// The setup of client2 should not have affected client1
+		const rolesAfterSetup = await yates1.inspectDBRoles(role);
+		// Expect to see 4 roles (CRUD) for each defined model when using wildcard abilities (*)
+		expect(rolesAfterSetup.length).toBe(models1.length * 4);
+	});
+});