feat(zod): add ORM-style select/include/omit options to makeModelSchema

packages/zod/src/factory.ts

@@ -15,8 +15,10 @@ import { SchemaFactoryError } from './error';
 import type {
     GetModelCreateFieldsShape,
     GetModelFieldsShape,
+    GetModelSchemaShapeWithOptions,
     GetModelUpdateFieldsShape,
     GetTypeDefFieldsShape,
+    ModelSchemaOptions,
 } from './types';
 import {
     addBigIntValidation,
@@ -30,6 +32,44 @@ export function createSchemaFactory<Schema extends SchemaDef>(schema: Schema) {
     return new SchemaFactory(schema);
 }
 
+/** Internal untyped representation of the options object used at runtime. */
+type RawOptions = {
+    select?: Record<string, unknown>;
+    include?: Record<string, unknown>;
+    omit?: Record<string, unknown>;
+};
+
+/**
+ * Recursive Zod schema that validates a `RawOptions` object at runtime,
+ * enforcing the same mutual-exclusion rules that the TypeScript union type
+ * enforces at compile time:
+ *  - `select` and `include` cannot be used together.
+ *  - `select` and `omit` cannot be used together.
+ * Nested relation options are validated with the same rules.
+ */
+const rawOptionsSchema: z.ZodType<RawOptions> = z.lazy(() =>
+    z
+        .object({
+            select: z.record(z.string(), z.union([z.boolean(), rawOptionsSchema])).optional(),
+            include: z.record(z.string(), z.union([z.boolean(), rawOptionsSchema])).optional(),
+            omit: z.record(z.string(), z.boolean()).optional(),
+        })
+        .superRefine((val, ctx) => {
+            if (val.select && val.include) {
+                ctx.addIssue({
+                    code: 'custom',
+                    message: '`select` and `include` cannot be used together',
+                });
+            }
+            if (val.select && val.omit) {
+                ctx.addIssue({
+                    code: 'custom',
+                    message: '`select` and `omit` cannot be used together',
+                });
+            }
+        }),
+);
+
 class SchemaFactory<Schema extends SchemaDef> {
     private readonly schema: SchemaAccessor<Schema>;
 
@@ -39,29 +79,64 @@ class SchemaFactory<Schema extends SchemaDef> {
 
     makeModelSchema<Model extends GetModels<Schema>>(
         model: Model,
-    ): z.ZodObject<GetModelFieldsShape<Schema, Model>, z.core.$strict> {
+    ): z.ZodObject<GetModelFieldsShape<Schema, Model>, z.core.$strict>;
+
+    makeModelSchema<Model extends GetModels<Schema>, Options extends ModelSchemaOptions<Schema, Model>>(
+        model: Model,
+        options: Options,
+    ): z.ZodObject<GetModelSchemaShapeWithOptions<Schema, Model, Options>, z.core.$strict>;
+
+    makeModelSchema<Model extends GetModels<Schema>, Options extends ModelSchemaOptions<Schema, Model>>(
+        model: Model,
+        options?: Options,
+    ): z.ZodObject<Record<string, z.ZodType>, z.core.$strict> {
         const modelDef = this.schema.requireModel(model);
-        const fields: Record<string, z.ZodType> = {};
 
-        for (const [fieldName, fieldDef] of Object.entries(modelDef.fields)) {
-            if (fieldDef.relation) {
-                const relatedModelName = fieldDef.type;
-                const lazySchema: z.ZodType = z.lazy(() => this.makeModelSchema(relatedModelName as GetModels<Schema>));
-                // relation fields are always optional
-                fields[fieldName] = this.applyDescription(
-                    this.applyCardinality(lazySchema, fieldDef).optional(),
-                    fieldDef.attributes,
-                );
-            } else {
-                fields[fieldName] = this.applyDescription(this.makeScalarFieldSchema(fieldDef), fieldDef.attributes);
+        if (!options) {
+            // ── No-options path (original behaviour) ─────────────────────────
+            const fields: Record<string, z.ZodType> = {};
+
+            for (const [fieldName, fieldDef] of Object.entries(modelDef.fields)) {
+                if (fieldDef.relation) {
+                    const relatedModelName = fieldDef.type;
+                    const lazySchema: z.ZodType = z.lazy(() =>
+                        this.makeModelSchema(relatedModelName as GetModels<Schema>),
+                    );
+                    // relation fields are always optional
+                    fields[fieldName] = this.applyDescription(
+                        this.applyCardinality(lazySchema, fieldDef).optional(),
+                        fieldDef.attributes,
+                    );
+                } else {
+                    fields[fieldName] = this.applyDescription(
+                        this.makeScalarFieldSchema(fieldDef),
+                        fieldDef.attributes,
+                    );
+                }
             }
+
+            const shape = z.strictObject(fields);
+            return this.applyDescription(
+                addCustomValidation(shape, modelDef.attributes),
+                modelDef.attributes,
+            ) as unknown as z.ZodObject<GetModelFieldsShape<Schema, Model>, z.core.$strict>;
         }
 
+        // ── Options path ─────────────────────────────────────────────────────
+        const rawOptions = rawOptionsSchema.parse(options);
+        const fields = this.buildFieldsWithOptions(model as string, rawOptions);
         const shape = z.strictObject(fields);
-        return this.applyDescription(
-            addCustomValidation(shape, modelDef.attributes),
-            modelDef.attributes,
-        ) as unknown as z.ZodObject<GetModelFieldsShape<Schema, Model>, z.core.$strict>;
+        // @@validate conditions only reference scalar fields of the same model
+        // (the ZModel compiler rejects relation fields). When `select` or `omit`
+        // produces a partial shape some of those scalar fields may be absent;
+        // we skip any rule that references a missing field so it can't produce
+        // a false negative against a partial payload.
+        const presentFields = this.buildPresentFields(model as string, rawOptions);
+        const withValidation = addCustomValidation(shape, modelDef.attributes, presentFields);
+        return this.applyDescription(withValidation, modelDef.attributes) as unknown as z.ZodObject<
+            GetModelSchemaShapeWithOptions<Schema, Model, Options>,
+            z.core.$strict
+        >;
     }
 
     makeModelCreateSchema<Model extends GetModels<Schema>>(
@@ -114,6 +189,149 @@ class SchemaFactory<Schema extends SchemaDef> {
         ) as unknown as z.ZodObject<GetModelUpdateFieldsShape<Schema, Model>, z.core.$strict>;
     }
 
+    // -------------------------------------------------------------------------
+    // Options-aware field building
+    // -------------------------------------------------------------------------
+
+    /**
+     * Internal loose options shape used at runtime (we've already validated the
+     * type-level constraints via the overload signatures).
+     */
+    private buildFieldsWithOptions(model: string, options: RawOptions): Record<string, z.ZodType> {
+        const { select, include, omit } = options;
+        const modelDef = this.schema.requireModel(model);
+        const fields: Record<string, z.ZodType> = {};
+
+        if (select) {
+            // ── select branch ────────────────────────────────────────────────
+            // Only include fields that are explicitly listed with a truthy value.
+            for (const [key, value] of Object.entries(select)) {
+                if (!value) continue; // false → skip
+
+                const fieldDef = modelDef.fields[key];
+                if (!fieldDef) {
+                    throw new SchemaFactoryError(`Field "${key}" does not exist on model "${model}"`);
+                }
+
+                if (fieldDef.relation) {
+                    const subOptions = typeof value === 'object' ? (value as RawOptions) : undefined;
+                    const relSchema = this.makeRelationFieldSchema(fieldDef, subOptions);
+                    fields[key] = this.applyDescription(
+                        this.applyCardinality(relSchema, fieldDef).optional(),
+                        fieldDef.attributes,
+                    );
+                } else {
+                    if (typeof value === 'object') {
+                        throw new SchemaFactoryError(
+                            `Field "${key}" on model "${model}" is a scalar field and cannot have nested options`,
+                        );
+                    }
+                    fields[key] = this.applyDescription(this.makeScalarFieldSchema(fieldDef), fieldDef.attributes);
+                }
+            }
+        } else {
+            // ── include + omit branch ────────────────────────────────────────
+            // Validate omit keys up-front.
+            if (omit) {
+                for (const key of Object.keys(omit)) {
+                    const fieldDef = modelDef.fields[key];
+                    if (!fieldDef) {
+                        throw new SchemaFactoryError(`Field "${key}" does not exist on model "${model}"`);
+                    }
+                    if (fieldDef.relation) {
+                        throw new SchemaFactoryError(
+                            `Field "${key}" on model "${model}" is a relation field and cannot be used in "omit"`,
+                        );
+                    }
+                }
+            }
+
+            // Start with all scalar fields, applying omit exclusions.
+            for (const [fieldName, fieldDef] of Object.entries(modelDef.fields)) {
+                if (fieldDef.relation) continue;
+
+                if (omit?.[fieldName] === true) continue;
+                fields[fieldName] = this.applyDescription(this.makeScalarFieldSchema(fieldDef), fieldDef.attributes);
+            }
+
+            // Validate include keys and add relation fields.
+            if (include) {
+                for (const [key, value] of Object.entries(include)) {
+                    if (!value) continue; // false → skip
+
+                    const fieldDef = modelDef.fields[key];
+                    if (!fieldDef) {
+                        throw new SchemaFactoryError(`Field "${key}" does not exist on model "${model}"`);
+                    }
+                    if (!fieldDef.relation) {
+                        throw new SchemaFactoryError(
+                            `Field "${key}" on model "${model}" is not a relation field and cannot be used in "include"`,
+                        );
+                    }
+
+                    const subOptions = typeof value === 'object' ? (value as RawOptions) : undefined;
+                    const relSchema = this.makeRelationFieldSchema(fieldDef, subOptions);
+                    fields[key] = this.applyDescription(
+                        this.applyCardinality(relSchema, fieldDef).optional(),
+                        fieldDef.attributes,
+                    );
+                }
+            }
+        }
+
+        return fields;
+    }
+
+    /**
+     * Returns the set of scalar field names that will be present in the
+     * resulting schema after applying `options`. Used by `addCustomValidation`
+     * to skip `@@validate` rules that reference an absent field.
+     *
+     * Only scalar fields matter here because `@@validate` conditions are
+     * restricted by the ZModel compiler to scalar fields of the same model.
+     */
+    private buildPresentFields(model: string, options: RawOptions): ReadonlySet<string> {
+        const { select, omit } = options;
+        const modelDef = this.schema.requireModel(model);
+        const fields = new Set<string>();
+
+        if (select) {
+            // Only scalar fields explicitly selected with a truthy value.
+            for (const [key, value] of Object.entries(select)) {
+                if (!value) continue;
+                const fieldDef = modelDef.fields[key];
+                if (fieldDef && !fieldDef.relation) {
+                    fields.add(key);
+                }
+            }
+        } else {
+            // All scalar fields minus explicitly omitted ones.
+            for (const [fieldName, fieldDef] of Object.entries(modelDef.fields)) {
+                if (fieldDef.relation) continue;
+                if (omit?.[fieldName] === true) continue;
+                fields.add(fieldName);
+            }
+        }
+
+        return fields;
+    }
+
+    /**
+     * Build the inner Zod schema for a relation field, optionally with nested
+     * query options.  Does NOT apply cardinality/optional wrappers — the caller
+     * does that.
+     */
+    private makeRelationFieldSchema(fieldDef: FieldDef, subOptions?: RawOptions): z.ZodType {
+        const relatedModelName = fieldDef.type as GetModels<Schema>;
+        if (subOptions) {
+            // Recurse: build the related model's schema with its own options.
+            return this.makeModelSchema(relatedModelName, subOptions as ModelSchemaOptions<Schema, GetModels<Schema>>);
+        }
+        // No sub-options: use a lazy reference to the default schema so that
+        // circular models don't cause infinite recursion at build time.
+        return z.lazy(() => this.makeModelSchema(relatedModelName));
+    }
+
     private makeScalarFieldSchema(fieldDef: FieldDef): z.ZodType {
         const { type, attributes } = fieldDef;
 

packages/zod/src/index.ts

@@ -1,2 +1,3 @@
 export { createSchemaFactory } from './factory';
+export type { ModelSchemaOptions, GetModelSchemaShapeWithOptions } from './types';
 export * as ZodUtils from './utils';