perf(orm): memoize implicit m2m join-table and model lookups (#2715)

packages/orm/src/client/query-utils.ts

@@ -20,8 +20,38 @@ export function hasModel(schema: SchemaDef, model: string) {
         .includes(model.toLowerCase());
 }
 
+/**
+ * Structural lookups derived purely from the (immutable) schema. The schema does not
+ * change for a client's lifetime, so these results are memoized per-schema and shared
+ * across every consumer (orm query building + plugins). Keyed by the schema object so
+ * the cache dies with the schema and never leaks across schemas. See issue #2715.
+ */
+interface SchemaLookupCache {
+    model: Map<string, ModelDef | undefined>;
+    m2mRelation: Map<string, ReturnType<typeof computeManyToManyRelation>>;
+    m2mJoinTable?: Map<string, ManyToManyJoinTableEndpoints | undefined>;
+}
+
+const schemaLookupCache = new WeakMap<SchemaDef, SchemaLookupCache>();
+
+function getSchemaLookupCache(schema: SchemaDef): SchemaLookupCache {
+    let cache = schemaLookupCache.get(schema);
+    if (!cache) {
+        cache = { model: new Map(), m2mRelation: new Map() };
+        schemaLookupCache.set(schema, cache);
+    }
+    return cache;
+}
+
 export function getModel(schema: SchemaDef, model: string) {
-    return Object.values(schema.models).find((m) => m.name.toLowerCase() === model.toLowerCase());
+    const cache = getSchemaLookupCache(schema);
+    const key = model.toLowerCase();
+    if (cache.model.has(key)) {
+        return cache.model.get(key);
+    }
+    const result = Object.values(schema.models).find((m) => m.name.toLowerCase() === key);
+    cache.model.set(key, result);
+    return result;
 }
 
 export function getTypeDef(schema: SchemaDef, type: string) {
@@ -260,6 +290,58 @@ export function makeDefaultOrderBy(schema: SchemaDef, model: string) {
 }
 
 export function getManyToManyRelation(schema: SchemaDef, model: string, field: string) {
+    const cache = getSchemaLookupCache(schema);
+    const key = `${model} ${field}`;
+    if (cache.m2mRelation.has(key)) {
+        return cache.m2mRelation.get(key);
+    }
+    const result = computeManyToManyRelation(schema, model, field);
+    cache.m2mRelation.set(key, result);
+    return result;
+}
+
+/**
+ * Endpoints of an implicit many-to-many relation, identified by its join table name.
+ */
+export interface ManyToManyJoinTableEndpoints {
+    model: string;
+    field: string;
+    otherModel: string;
+    otherField: string;
+}
+
+/**
+ * Resolve the relation endpoints for an implicit many-to-many join table by its table
+ * name, or `undefined` if the table is not an implicit m2m join table. The join-table
+ * index is built once per schema (single pass) and reused, making this an O(1) lookup.
+ * See issue #2715.
+ */
+export function getManyToManyJoinTable(
+    schema: SchemaDef,
+    joinTableName: string,
+): ManyToManyJoinTableEndpoints | undefined {
+    const cache = getSchemaLookupCache(schema);
+    if (!cache.m2mJoinTable) {
+        const map = new Map<string, ManyToManyJoinTableEndpoints>();
+        for (const model of Object.values(schema.models)) {
+            for (const field of Object.values(model.fields)) {
+                const m2m = getManyToManyRelation(schema, model.name, field.name);
+                if (m2m?.joinTable && !map.has(m2m.joinTable)) {
+                    map.set(m2m.joinTable, {
+                        model: model.name,
+                        field: field.name,
+                        otherModel: m2m.otherModel,
+                        otherField: m2m.otherField,
+                    });
+                }
+            }
+        }
+        cache.m2mJoinTable = map;
+    }
+    return cache.m2mJoinTable.get(joinTableName);
+}
+
+function computeManyToManyRelation(schema: SchemaDef, model: string, field: string) {
     const fieldDef = requireField(schema, model, field);
     if (!fieldDef.array || !fieldDef.relation?.opposite) {
         return undefined;

packages/plugins/policy/src/policy-handler.ts

@@ -1290,40 +1290,39 @@ export class PolicyHandler<Schema extends SchemaDef> extends OperationNodeTransf
     }
 
     private resolveManyToManyJoinTable(tableName: string) {
-        for (const model of Object.values(this.client.$schema.models)) {
-            for (const field of Object.values(model.fields)) {
-                const m2m = QueryUtils.getManyToManyRelation(this.client.$schema, model.name, field.name);
-                if (m2m?.joinTable === tableName) {
-                    const sortedRecord = [
-                        {
-                            model: model.name,
-                            field: field.name,
-                        },
-                        {
-                            model: m2m.otherModel,
-                            field: m2m.otherField,
-                        },
-                    ].sort(this.manyToManySorter);
-
-                    const firstIdFields = QueryUtils.requireIdFields(this.client.$schema, sortedRecord[0]!.model);
-                    const secondIdFields = QueryUtils.requireIdFields(this.client.$schema, sortedRecord[1]!.model);
-                    invariant(
-                        firstIdFields.length === 1 && secondIdFields.length === 1,
-                        'only single-field id is supported for implicit many-to-many join table',
-                    );
-
-                    return {
-                        firstModel: sortedRecord[0]!.model,
-                        firstField: sortedRecord[0]!.field,
-                        firstIdField: firstIdFields[0]!,
-                        secondModel: sortedRecord[1]!.model,
-                        secondField: sortedRecord[1]!.field,
-                        secondIdField: secondIdFields[0]!,
-                    };
-                }
-            }
+        // O(1) lookup backed by a per-schema index (built once); previously this scanned
+        // the entire schema on every call for every table. See issue #2715.
+        const endpoints = QueryUtils.getManyToManyJoinTable(this.client.$schema, tableName);
+        if (!endpoints) {
+            return undefined;
         }
-        return undefined;
+
+        const sortedRecord = [
+            {
+                model: endpoints.model,
+                field: endpoints.field,
+            },
+            {
+                model: endpoints.otherModel,
+                field: endpoints.otherField,
+            },
+        ].sort(this.manyToManySorter);
+
+        const firstIdFields = QueryUtils.requireIdFields(this.client.$schema, sortedRecord[0]!.model);
+        const secondIdFields = QueryUtils.requireIdFields(this.client.$schema, sortedRecord[1]!.model);
+        invariant(
+            firstIdFields.length === 1 && secondIdFields.length === 1,
+            'only single-field id is supported for implicit many-to-many join table',
+        );
+
+        return {
+            firstModel: sortedRecord[0]!.model,
+            firstField: sortedRecord[0]!.field,
+            firstIdField: firstIdFields[0]!,
+            secondModel: sortedRecord[1]!.model,
+            secondField: sortedRecord[1]!.field,
+            secondIdField: secondIdFields[0]!,
+        };
     }
 
     private manyToManySorter(a: { model: string; field: string }, b: { model: string; field: string }): number {

tests/regression/test/issue-2715.test.ts

@@ -0,0 +1,74 @@
+import { QueryUtils } from '@zenstackhq/orm';
+import { createPolicyTestClient } from '@zenstackhq/testtools';
+import { describe, expect, it } from 'vitest';
+
+// https://github.com/zenstackhq/zenstack/issues/2715
+// Resolving an implicit many-to-many join table is a pure function of (schema, tableName),
+// but it used to re-scan the entire schema on every call, for every table in every query.
+// The resolution is now memoized per (immutable) schema. These tests guard that the
+// memoization holds (so the per-query full-schema scan cannot silently regress) and that
+// the resolved endpoints are correct across the shapes implicit m2m relations can take:
+// plain two-model, self-relation, explicit @relation name, and multiple relations at once.
+describe('Regression for issue #2715', () => {
+    it('resolves implicit m2m join tables correctly and caches each per schema', async () => {
+        const db = await createPolicyTestClient(
+            `
+model User {
+    id     Int     @id @default(autoincrement())
+    groups Group[]
+    @@allow('all', true)
+}
+
+model Group {
+    id    Int    @id @default(autoincrement())
+    users User[]
+    @@allow('all', true)
+}
+
+// self-relation m2m with an explicit join-table name
+model Person {
+    id        Int      @id @default(autoincrement())
+    following Person[] @relation('follows')
+    followers Person[] @relation('follows')
+    @@allow('all', true)
+}
+            `,
+        );
+
+        const schema = db.$schema;
+
+        // plain two-model m2m: join table follows Prisma's `_<A>To<B>` (sorted) convention,
+        // endpoints come from the first declared side (User.groups), other side is Group.users
+        const groupToUser = QueryUtils.getManyToManyJoinTable(schema, '_GroupToUser');
+        expect(groupToUser).toEqual({
+            model: 'User',
+            field: 'groups',
+            otherModel: 'Group',
+            otherField: 'users',
+        });
+
+        // self-relation m2m with explicit name: join table is `_<name>`, both endpoints are
+        // the same model, fields are the two relation fields (first declared side first)
+        const follows = QueryUtils.getManyToManyJoinTable(schema, '_follows');
+        expect(follows).toEqual({
+            model: 'Person',
+            field: 'following',
+            otherModel: 'Person',
+            otherField: 'followers',
+        });
+
+        // memoized: repeated resolution returns the SAME object reference (the index is built
+        // once and reused). An un-memoized re-scan would construct a fresh descriptor each call,
+        // so these `toBe`s are what fail if the per-query scan regresses.
+        expect(QueryUtils.getManyToManyJoinTable(schema, '_GroupToUser')).toBe(groupToUser);
+        expect(QueryUtils.getManyToManyJoinTable(schema, '_follows')).toBe(follows);
+
+        // distinct relations are cached independently (no cross-contamination)
+        expect(groupToUser).not.toBe(follows);
+
+        // non-m2m table names and unknown tables resolve to undefined (negative results are
+        // part of the cached index, not an uncached miss)
+        expect(QueryUtils.getManyToManyJoinTable(schema, 'User')).toBeUndefined();
+        expect(QueryUtils.getManyToManyJoinTable(schema, '_DoesNotExist')).toBeUndefined();
+    });
+});