fix(orm): share the zod schema factory across derived clients that need the same schemas

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

@@ -33,6 +33,7 @@ import { FindOperationHandler } from './crud/operations/find';
 import { GroupByOperationHandler } from './crud/operations/group-by';
 import { UpdateOperationHandler } from './crud/operations/update';
 import { InputValidator } from './crud/validator';
+import { haveSameSchemaInputs } from './zod/factory';
 import type { Diagnostics, QueryInfo } from './diagnostics';
 import { createConfigError, createNotFoundError, createNotSupportedError } from './errors';
 import { ZenStackDriver } from './executor/zenstack-driver';
@@ -518,11 +519,7 @@ export class ClientImpl {
             plugins: newPlugins,
         };
         const newClient = new ClientImpl(this.schema, newOptions, this);
-        // create a new validator to have a fresh schema cache, because plugins may extend the
-        // query args schemas
-        newClient.inputValidator = new InputValidator(newClient as any, {
-            enabled: newOptions.validateInput !== false,
-        });
+        newClient.inputValidator = this.deriveInputValidator(newClient);
         return newClient;
     }
 
@@ -539,11 +536,7 @@ export class ClientImpl {
             plugins: newPlugins,
         };
         const newClient = new ClientImpl(this.schema, newOptions, this);
-        // create a new validator to have a fresh schema cache, because plugins may
-        // extend the query args schemas
-        newClient.inputValidator = new InputValidator(newClient as any, {
-            enabled: newClient.$options.validateInput !== false,
-        });
+        newClient.inputValidator = this.deriveInputValidator(newClient);
         return newClient;
     }
 
@@ -554,14 +547,26 @@ export class ClientImpl {
             plugins: [] as AnyPlugin[],
         };
         const newClient = new ClientImpl(this.schema, newOptions, this);
-        // create a new validator to have a fresh schema cache, because plugins may
-        // extend the query args schemas
-        newClient.inputValidator = new InputValidator(newClient as any, {
-            enabled: newOptions.validateInput !== false,
-        });
+        newClient.inputValidator = this.deriveInputValidator(newClient);
         return newClient;
     }
 
+    /**
+     * Validator for a client derived from this one. The derived client reuses this client's Zod
+     * schema factory, and so its schema cache, whenever its options shape the same query-args
+     * schemas; it only gets a fresh factory when a plugin that extends query args or results was
+     * added or removed, or `omit`/`slicing` changed. The validator itself is always new because
+     * `validateInput` may differ.
+     */
+    private deriveInputValidator(newClient: ClientImpl) {
+        return new InputValidator(newClient as any, {
+            enabled: newClient.$options.validateInput !== false,
+            zodFactory: haveSameSchemaInputs(this.$options, newClient.$options)
+                ? this.inputValidator.zodFactory
+                : undefined,
+        });
+    }
+
     $setAuth(auth: AuthType<SchemaDef> | undefined) {
         if (auth !== undefined && (typeof auth !== 'object' || auth === null || Array.isArray(auth))) {
             throw new Error('Invalid auth object');
@@ -588,10 +593,7 @@ export class ClientImpl {
 
     $setOptions<Options extends ClientOptions<SchemaDef>>(options: Options): ClientContract<SchemaDef, Options> {
         const newClient = new ClientImpl(this.schema, options as ClientOptions<SchemaDef>, this);
-        // create a new validator to have a fresh schema cache, because options may change validation settings
-        newClient.inputValidator = new InputValidator(newClient as any, {
-            enabled: newClient.$options.validateInput !== false,
-        });
+        newClient.inputValidator = this.deriveInputValidator(newClient);
         return newClient as unknown as ClientContract<SchemaDef, Options>;
     }
 

packages/orm/src/client/crud/validator/validator.ts

@@ -25,11 +25,17 @@ import { ZodSchemaFactory } from '../../zod/factory';
 
 type GetSchemaFunc<Schema extends SchemaDef> = (model: GetModels<Schema>) => ZodType;
 
-export type InputValidatorOptions = {
+export type InputValidatorOptions<Schema extends SchemaDef> = {
     /**
      * Whether validation is enabled. Defaults to `true`.
      */
     enabled?: boolean;
+
+    /**
+     * Schema factory to validate with. Defaults to a new factory (and so a new schema cache) bound
+     * to the client.
+     */
+    zodFactory?: ZodSchemaFactory<Schema>;
 };
 
 export class InputValidator<Schema extends SchemaDef> {
@@ -38,9 +44,9 @@ export class InputValidator<Schema extends SchemaDef> {
 
     constructor(
         private readonly client: ClientContract<Schema>,
-        options?: InputValidatorOptions,
+        options?: InputValidatorOptions<Schema>,
     ) {
-        this.zodFactory = new ZodSchemaFactory(client);
+        this.zodFactory = options?.zodFactory ?? new ZodSchemaFactory(client);
         this.enabled = options?.enabled !== false;
     }
 

packages/orm/src/client/zod/factory.ts

@@ -84,6 +84,26 @@ export function createQuerySchemaFactory(clientOrSchema: any, options?: any) {
     return new ZodSchemaFactory(clientOrSchema, options);
 }
 
+/**
+ * Whether two client options produce the same query-args schemas, i.e. a `ZodSchemaFactory`
+ * built for one can serve the other. Only `omit`, `slicing`, `allowQueryTimeOmitOverride` and
+ * the plugins that extend query args or results shape the schemas; dialect, auth, hooks and
+ * validation flags do not.
+ */
+export function haveSameSchemaInputs(a: ClientOptions<SchemaDef>, b: ClientOptions<SchemaDef>): boolean {
+    if (a === b) {
+        return true;
+    }
+    if (a.omit !== b.omit || a.slicing !== b.slicing || a.allowQueryTimeOmitOverride !== b.allowQueryTimeOmitOverride) {
+        return false;
+    }
+    const shaping = (options: ClientOptions<SchemaDef>) =>
+        (options.plugins ?? []).filter((plugin) => plugin.queryArgs || plugin.result);
+    const pa = shaping(a);
+    const pb = shaping(b);
+    return pa.length === pb.length && pa.every((plugin, index) => plugin === pb[index]);
+}
+
 /**
  * Builds a `DateTime` value schema that accepts a `Date` object or any string
  * the JS `Date` constructor parses, and coerces it to a `Date`. ISO datetime,

tests/e2e/orm/plugin-infra/zod-factory-sharing.test.ts

@@ -0,0 +1,54 @@
+import { definePlugin, type ClientContract } from '@zenstackhq/orm';
+import { createTestClient } from '@zenstackhq/testtools';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+import z from 'zod';
+import { schema } from '../schemas/basic';
+
+describe('Zod schema factory sharing across derived clients', () => {
+    let db: ClientContract<typeof schema>;
+
+    beforeEach(async () => {
+        db = await createTestClient(schema);
+    });
+
+    afterEach(async () => {
+        await db?.$disconnect();
+    });
+
+    it('shares the factory with a $use client whose plugin does not extend schemas', async () => {
+        const derived = db.$use(
+            definePlugin({
+                id: 'noop',
+                onQuery: ({ args, proceed }) => proceed(args),
+            }),
+        );
+        expect(derived.$zod).toBe(db.$zod);
+
+        await derived.user.findMany();
+        const base = await db.$diagnostics;
+        const ext = await derived.$diagnostics;
+        expect(base.zodCache.size).toBeGreaterThan(0);
+        expect(ext.zodCache).toEqual(base.zodCache);
+    });
+
+    it('gives a fresh factory to a $use client whose plugin extends query args', () => {
+        const derived = db.$use(
+            definePlugin({
+                id: 'cache',
+                queryArgs: { $read: z.object({ cache: z.boolean().optional() }) },
+                onQuery: ({ args, proceed }) => proceed(args),
+            }),
+        );
+        expect(derived.$zod).not.toBe(db.$zod);
+    });
+
+    it('shares the factory when only validation or non-schema options change', () => {
+        expect(db.$setInputValidation(false).$zod).toBe(db.$zod);
+        expect(db.$setOptions({ ...db.$options }).$zod).toBe(db.$zod);
+    });
+
+    it('gives a fresh factory when a schema-shaping option changes', () => {
+        const derived = db.$setOptions({ ...db.$options, allowQueryTimeOmitOverride: false });
+        expect(derived.$zod).not.toBe(db.$zod);
+    });
+});