perf(orm): reuse QueryNameMapper across derived executors and clients

packages/orm/src/client/client-impl.ts

@@ -50,6 +50,13 @@ type ExtResultFieldDef = {
     compute: (data: Record<string, any>) => unknown;
 };
 
+/**
+ * Returns the name mapper held by a ZenStack executor, or undefined for a plain kysely one.
+ */
+function getExecutorNameMapper(executor: QueryExecutor | undefined) {
+    return executor instanceof ZenStackQueryExecutor ? executor.getNameMapper() : undefined;
+}
+
 /**
  * ZenStack ORM client.
  */
@@ -99,6 +106,16 @@ export class ClientImpl {
                         baseClient.kyselyProps.dialect.createQueryCompiler(),
                         baseClient.kyselyProps.dialect.createAdapter(),
                         new DefaultConnectionProvider(baseClient.kyselyProps.driver),
+                        [],
+                        false,
+                        // A name mapper is derived purely from `$schema` and `$options`, so it can be
+                        // reused when neither changed - which is the case for derived clients like the
+                        // one `$transaction` creates. Rebuilding it is O(models x fields). See #2773.
+                        // Deliberately an identity check: `$use`/`$setOptions` and friends pass a new
+                        // options object, and the mapper's dialect is built from those options.
+                        baseClient.$schema === schema && baseClient.$options === options
+                            ? getExecutorNameMapper(baseClient.kyselyProps.executor)
+                            : undefined,
                     ),
             };
             this.kyselyRaw = baseClient.kyselyRaw;

packages/orm/src/client/executor/zenstack-query-executor.ts

@@ -1,5 +1,5 @@
 import { invariant } from '@zenstackhq/common-helpers';
-import type { ModelDef, SchemaDef, TypeDefDef } from '@zenstackhq/schema';
+import type { SchemaDef } from '@zenstackhq/schema';
 import type { QueryId } from 'kysely';
 import {
     AndNode,
@@ -37,7 +37,7 @@ import { getCrudDialect } from '../crud/dialects';
 import type { BaseCrudDialect } from '../crud/dialects/base-dialect';
 import { createDBQueryError, createInternalError, ORMError } from '../errors';
 import type { AfterEntityMutationCallback, OnKyselyQueryCallback } from '../plugin';
-import { requireIdFields, stripAlias } from '../query-utils';
+import { requireIdFields, schemaHasMappedNames, stripAlias } from '../query-utils';
 import { QueryNameMapper } from './name-mapper';
 import { TempAliasTransformer } from './temp-alias-transformer';
 import type { ZenStackDriver } from './zenstack-driver';
@@ -86,28 +86,30 @@ export class ZenStackQueryExecutor extends DefaultQueryExecutor {
         private readonly connectionProvider: ConnectionProvider,
         plugins: KyselyPlugin[] = [],
         private suppressMutationHooks: boolean = false,
+        nameMapper?: QueryNameMapper,
     ) {
         super(compiler, adapter, connectionProvider, plugins);
 
-        if (
-            client.$schema.provider.type === 'postgresql' || // postgres queries need to be schema-qualified
-            this.schemaHasMappedNames(client.$schema)
-        ) {
-            this.nameMapper = new QueryNameMapper(client as unknown as ClientContract<SchemaDef>);
-        }
+        // A `QueryNameMapper` is derived purely from the client's `$schema` and `$options`, and building
+        // it is O(models x fields) (plus an O(models x relations) pass for postgres). Reuse the one from
+        // the executor/client we're derived from when it was built from the same schema and options,
+        // otherwise every derived executor rebuilds whole-schema state. See issue #2773.
+        this.nameMapper =
+            nameMapper ??
+            (client.$schema.provider.type === 'postgresql' || // postgres queries need to be schema-qualified
+            schemaHasMappedNames(client.$schema)
+                ? new QueryNameMapper(client as unknown as ClientContract<SchemaDef>)
+                : undefined);
 
         this.dialect = getCrudDialect(client.$schema, client.$options);
     }
 
-    private schemaHasMappedNames(schema: SchemaDef) {
-        const hasMapAttr = (decl: ModelDef | TypeDefDef) => {
-            if (decl.attributes?.some((attr) => attr.name === '@@map')) {
-                return true;
-            }
-            return Object.values(decl.fields).some((field) => field.attributes?.some((attr) => attr.name === '@map'));
-        };
-
-        return Object.values(schema.models).some(hasMapAttr) || Object.values(schema.typeDefs ?? []).some(hasMapAttr);
+    /**
+     * The name mapper built for this executor's schema, if the schema needs one. Exposed so that
+     * derived clients built from the same schema and options can reuse it instead of rebuilding it.
+     */
+    getNameMapper() {
+        return this.nameMapper;
     }
 
     private get kysely() {
@@ -770,6 +772,7 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
             this.connectionProvider,
             [...this.plugins, plugin],
             this.suppressMutationHooks,
+            this.nameMapper,
         );
     }
 
@@ -782,6 +785,7 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
             this.connectionProvider,
             [...this.plugins, ...plugins],
             this.suppressMutationHooks,
+            this.nameMapper,
         );
     }
 
@@ -794,6 +798,7 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
             this.connectionProvider,
             [plugin, ...this.plugins],
             this.suppressMutationHooks,
+            this.nameMapper,
         );
     }
 
@@ -806,6 +811,7 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
             this.connectionProvider,
             [],
             this.suppressMutationHooks,
+            this.nameMapper,
         );
     }
 
@@ -818,6 +824,7 @@ In such cases, ZenStack cannot reliably determine the IDs of the mutated entitie
             connectionProvider,
             this.plugins as KyselyPlugin[],
             this.suppressMutationHooks,
+            this.nameMapper,
         );
         // replace client with a new one associated with the new executor
         newExecutor.client = this.client.withExecutor(newExecutor);

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

@@ -1,5 +1,13 @@
 import { invariant } from '@zenstackhq/common-helpers';
-import { ExpressionUtils, type FieldDef, type GetModels, type ModelDef, type SchemaDef } from '@zenstackhq/schema';
+import {
+    type EnumDef,
+    ExpressionUtils,
+    type FieldDef,
+    type GetModels,
+    type ModelDef,
+    type SchemaDef,
+    type TypeDefDef,
+} from '@zenstackhq/schema';
 import {
     AliasNode,
     ColumnNode,
@@ -30,6 +38,7 @@ interface SchemaLookupCache {
     model: Map<string, ModelDef | undefined>;
     m2mRelation: Map<string, ReturnType<typeof computeManyToManyRelation>>;
     m2mJoinTable?: Map<string, ManyToManyJoinTableEndpoints | undefined>;
+    hasMappedNames?: boolean;
 }
 
 const schemaLookupCache = new WeakMap<SchemaDef, SchemaLookupCache>();
@@ -58,6 +67,32 @@ export function getTypeDef(schema: SchemaDef, type: string) {
     return schema.typeDefs?.[type];
 }
 
+/**
+ * Whether any model, type def, or field in the schema carries `@@map`/`@map`. Answering it walks
+ * every model and field, and it is asked once per query-executor construction, so the (immutable)
+ * answer is memoized per schema alongside the other structural lookups. See issue #2773.
+ */
+export function schemaHasMappedNames(schema: SchemaDef) {
+    const cache = getSchemaLookupCache(schema);
+    if (cache.hasMappedNames === undefined) {
+        // `fields` is optional on `EnumDef` (required on the other two), hence the `?? {}`.
+        const hasMapAttr = (decl: ModelDef | TypeDefDef | EnumDef) => {
+            if (decl.attributes?.some((attr) => attr.name === '@@map')) {
+                return true;
+            }
+            return Object.values(decl.fields ?? {}).some((field) =>
+                field.attributes?.some((attr) => attr.name === '@map'),
+            );
+        };
+        cache.hasMappedNames =
+            Object.values(schema.models).some(hasMapAttr) ||
+            Object.values(schema.typeDefs ?? {}).some(hasMapAttr) ||
+            // Enums carry name mapping too — `@@map` on the enum and `@map` on its members.
+            Object.values(schema.enums ?? {}).some(hasMapAttr);
+    }
+    return cache.hasMappedNames;
+}
+
 export function requireModel(schema: SchemaDef, model: string) {
     const modelDef = getModel(schema, model);
     if (!modelDef) {

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

@@ -0,0 +1,161 @@
+import { createTestClient } from '@zenstackhq/testtools';
+import { describe, expect, it } from 'vitest';
+
+// https://github.com/zenstackhq/zenstack/issues/2773
+//
+// `QueryNameMapper` is derived purely from the client's `$schema` and `$options`, but it was rebuilt
+// for every derived executor/client. Building it is O(models x fields), plus an O(models x relations)
+// pass on postgres, so on a large schema every `$transaction` paid that cost twice - once when
+// `$transaction` derives a client, and once when the query derives a connection-scoped executor.
+//
+// These assert the invariant (the mapper is not rebuilt) rather than elapsed time, which would be
+// both flaky and only an indirect proxy for it.
+
+const schema = `
+model Post {
+    id    Int    @id @default(autoincrement())
+    title String @map("post_title")
+
+    @@map("posts_table")
+}
+`;
+
+// Note: inside a transaction `$qb.getExecutor()` is kysely's own wrapping executor, so we read the
+// ZenStack executor the client itself was built with. We read the underlying field rather than the
+// accessor so that these assertions fail on a mapper-identity mismatch - the actual defect - rather
+// than on a missing method.
+function nameMapperOf(client: any) {
+    return client.kyselyProps.executor.nameMapper;
+}
+
+describe('Regression for issue #2773', () => {
+    it('reuses the name mapper for the client derived by $transaction', async () => {
+        const db = await createTestClient(schema, { provider: 'postgresql' });
+        const mapper = nameMapperOf(db);
+        expect(mapper).toBeDefined();
+
+        await db.$transaction(async (tx: any) => {
+            expect(nameMapperOf(tx)).toBe(mapper);
+        });
+    });
+
+    it('reuses the name mapper for connection-scoped and plugin-derived executors', async () => {
+        const db = await createTestClient(schema, { provider: 'postgresql' });
+        const executor = (db as any).kyselyProps.executor;
+        const mapper = executor.nameMapper;
+        expect(mapper).toBeDefined();
+
+        expect(executor.withoutPlugins().nameMapper).toBe(mapper);
+        expect(executor.withPlugin({}).nameMapper).toBe(mapper);
+        expect(executor.withPluginAtFront({}).nameMapper).toBe(mapper);
+        expect(executor.withPlugins([{}]).nameMapper).toBe(mapper);
+        expect(executor.withConnectionProvider(executor.connectionProvider).nameMapper).toBe(mapper);
+    });
+
+    it('still builds a mapper for a client derived with different options', async () => {
+        const db = await createTestClient(schema, { provider: 'postgresql' });
+        const mapper = nameMapperOf(db);
+
+        // `$use` derives a client with a NEW options object, and the mapper's dialect is built from
+        // options - so this one must not inherit the existing mapper.
+        const derived = db.$use({ id: 'noop', name: 'noop' } as any);
+        expect(nameMapperOf(derived)).toBeDefined();
+        expect(nameMapperOf(derived)).not.toBe(mapper);
+    });
+
+    it('keeps working when the schema needs no mapper at all', async () => {
+        // sqlite + no @@map/@map means nameMapper stays undefined. Deciding that walks every model
+        // and field, and it is asked on every executor construction, so it is memoized per schema -
+        // this guards the path where nothing is threaded through.
+        const db = await createTestClient(
+            `
+model Item {
+    id   Int    @id @default(autoincrement())
+    name String
+}
+            `,
+            { provider: 'sqlite' },
+        );
+        expect(nameMapperOf(db)).toBeUndefined();
+
+        await db.item.create({ data: { name: 'a' } });
+        const outside = await db.item.findMany();
+        expect(outside).toHaveLength(1);
+
+        await db.$transaction(async (tx: any) => {
+            expect(nameMapperOf(tx)).toBeUndefined();
+            expect(await tx.item.findMany()).toHaveLength(1);
+        });
+    });
+
+    // Enums carry name mapping too, and `schemaHasMappedNames` originally checked only models
+    // and type defs. On postgres the mapper is built unconditionally, which masked it; on any
+    // other provider a schema whose ONLY mapped name is on an enum got no mapper at all.
+    it('builds a mapper for a schema whose only mapped name is on an enum', async () => {
+        const db = await createTestClient(
+            `
+enum Status {
+    ACTIVE
+    ARCHIVED
+
+    @@map('status_enum')
+}
+
+model Item {
+    id     Int    @id @default(autoincrement())
+    name   String
+    status Status @default(ACTIVE)
+}
+            `,
+            { provider: 'sqlite' },
+        );
+
+        expect(nameMapperOf(db)).toBeDefined();
+
+        await db.item.create({ data: { name: 'a' } });
+        expect(await db.item.findMany()).toHaveLength(1);
+    });
+
+    it('builds a mapper for a schema whose only mapped name is on an enum MEMBER', async () => {
+        const db = await createTestClient(
+            `
+enum Status {
+    ACTIVE   @map('is_active')
+    ARCHIVED
+}
+
+model Item {
+    id     Int    @id @default(autoincrement())
+    name   String
+    status Status @default(ACTIVE)
+}
+            `,
+            { provider: 'sqlite' },
+        );
+
+        expect(nameMapperOf(db)).toBeDefined();
+
+        await db.item.create({ data: { name: 'a' } });
+        expect(await db.item.findMany()).toHaveLength(1);
+    });
+
+    it('applies @@map and @map identically inside and outside a transaction', async () => {
+        const db = await createTestClient(schema, { provider: 'postgresql' });
+        await db.post.create({ data: { title: 'hello' } });
+
+        const outside = await db.post.findMany();
+        expect(outside).toHaveLength(1);
+        expect(outside[0].title).toBe('hello');
+
+        const inside = await db.$transaction(async (tx: any) => tx.post.findMany());
+        expect(inside).toEqual(outside);
+
+        // the mapped table/column really are what the mapper produced
+        const raw = await db.$qb
+            .selectFrom('posts_table' as any)
+            .selectAll()
+            .execute();
+        expect(raw).toHaveLength(1);
+        expect((raw[0] as any).post_title).toBe('hello');
+    });
+});