Skip to content

Commit c36cf8e

Browse files
marcsigmundymc9
marcsigmund
and
ymc9
authored
feat(zod): add ORM-style select/include/omit options to makeModelSchema (#2503)
Co-authored-by: ymc9 <104139426+ymc9@users.noreply.github.com>
1 parent 109171b commit c36cf8e

File tree

5 files changed

+859
-21
lines changed

5 files changed

+859
-21
lines changed

‎packages/zod/src/factory.ts‎

Lines changed: 235 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,10 @@ import { SchemaFactoryError } from './error';
1515
import type {
1616
GetModelCreateFieldsShape,
1717
GetModelFieldsShape,
18+
GetModelSchemaShapeWithOptions,
1819
GetModelUpdateFieldsShape,
1920
GetTypeDefFieldsShape,
21+
ModelSchemaOptions,
2022
} from './types';
2123
import {
2224
addBigIntValidation,
@@ -30,6 +32,44 @@ export function createSchemaFactory<Schema extends SchemaDef>(schema: Schema) {
3032
return new SchemaFactory(schema);
3133
}
3234

35+
/** Internal untyped representation of the options object used at runtime. */
36+
type RawOptions = {
37+
select?: Record<string, unknown>;
38+
include?: Record<string, unknown>;
39+
omit?: Record<string, unknown>;
40+
};
41+
42+
/**
43+
* Recursive Zod schema that validates a `RawOptions` object at runtime,
44+
* enforcing the same mutual-exclusion rules that the TypeScript union type
45+
* enforces at compile time:
46+
* - `select` and `include` cannot be used together.
47+
* - `select` and `omit` cannot be used together.
48+
* Nested relation options are validated with the same rules.
49+
*/
50+
const rawOptionsSchema: z.ZodType<RawOptions> = z.lazy(() =>
51+
z
52+
.object({
53+
select: z.record(z.string(), z.union([z.boolean(), rawOptionsSchema])).optional(),
54+
include: z.record(z.string(), z.union([z.boolean(), rawOptionsSchema])).optional(),
55+
omit: z.record(z.string(), z.boolean()).optional(),
56+
})
57+
.superRefine((val, ctx) => {
58+
if (val.select && val.include) {
59+
ctx.addIssue({
60+
code: 'custom',
61+
message: '`select` and `include` cannot be used together',
62+
});
63+
}
64+
if (val.select && val.omit) {
65+
ctx.addIssue({
66+
code: 'custom',
67+
message: '`select` and `omit` cannot be used together',
68+
});
69+
}
70+
}),
71+
);
72+
3373
class SchemaFactory<Schema extends SchemaDef> {
3474
private readonly schema: SchemaAccessor<Schema>;
3575

@@ -39,29 +79,64 @@ class SchemaFactory<Schema extends SchemaDef> {
3979

4080
makeModelSchema<Model extends GetModels<Schema>>(
4181
model: Model,
42-
): z.ZodObject<GetModelFieldsShape<Schema, Model>, z.core.$strict> {
82+
): z.ZodObject<GetModelFieldsShape<Schema, Model>, z.core.$strict>;
83+
84+
makeModelSchema<Model extends GetModels<Schema>, Options extends ModelSchemaOptions<Schema, Model>>(
85+
model: Model,
86+
options: Options,
87+
): z.ZodObject<GetModelSchemaShapeWithOptions<Schema, Model, Options>, z.core.$strict>;
88+
89+
makeModelSchema<Model extends GetModels<Schema>, Options extends ModelSchemaOptions<Schema, Model>>(
90+
model: Model,
91+
options?: Options,
92+
): z.ZodObject<Record<string, z.ZodType>, z.core.$strict> {
4393
const modelDef = this.schema.requireModel(model);
44-
const fields: Record<string, z.ZodType> = {};
4594

46-
for (const [fieldName, fieldDef] of Object.entries(modelDef.fields)) {
47-
if (fieldDef.relation) {
48-
const relatedModelName = fieldDef.type;
49-
const lazySchema: z.ZodType = z.lazy(() => this.makeModelSchema(relatedModelName as GetModels<Schema>));
50-
// relation fields are always optional
51-
fields[fieldName] = this.applyDescription(
52-
this.applyCardinality(lazySchema, fieldDef).optional(),
53-
fieldDef.attributes,
54-
);
55-
} else {
56-
fields[fieldName] = this.applyDescription(this.makeScalarFieldSchema(fieldDef), fieldDef.attributes);
95+
if (!options) {
96+
// ── No-options path (original behaviour) ─────────────────────────
97+
const fields: Record<string, z.ZodType> = {};
98+
99+
for (const [fieldName, fieldDef] of Object.entries(modelDef.fields)) {
100+
if (fieldDef.relation) {
101+
const relatedModelName = fieldDef.type;
102+
const lazySchema: z.ZodType = z.lazy(() =>
103+
this.makeModelSchema(relatedModelName as GetModels<Schema>),
104+
);
105+
// relation fields are always optional
106+
fields[fieldName] = this.applyDescription(
107+
this.applyCardinality(lazySchema, fieldDef).optional(),
108+
fieldDef.attributes,
109+
);
110+
} else {
111+
fields[fieldName] = this.applyDescription(
112+
this.makeScalarFieldSchema(fieldDef),
113+
fieldDef.attributes,
114+
);
115+
}
57116
}
117+
118+
const shape = z.strictObject(fields);
119+
return this.applyDescription(
120+
addCustomValidation(shape, modelDef.attributes),
121+
modelDef.attributes,
122+
) as unknown as z.ZodObject<GetModelFieldsShape<Schema, Model>, z.core.$strict>;
58123
}
59124

125+
// ── Options path ─────────────────────────────────────────────────────
126+
const rawOptions = rawOptionsSchema.parse(options);
127+
const fields = this.buildFieldsWithOptions(model as string, rawOptions);
60128
const shape = z.strictObject(fields);
61-
return this.applyDescription(
62-
addCustomValidation(shape, modelDef.attributes),
63-
modelDef.attributes,
64-
) as unknown as z.ZodObject<GetModelFieldsShape<Schema, Model>, z.core.$strict>;
129+
// @@validate conditions only reference scalar fields of the same model
130+
// (the ZModel compiler rejects relation fields). When `select` or `omit`
131+
// produces a partial shape some of those scalar fields may be absent;
132+
// we skip any rule that references a missing field so it can't produce
133+
// a false negative against a partial payload.
134+
const presentFields = this.buildPresentFields(model as string, rawOptions);
135+
const withValidation = addCustomValidation(shape, modelDef.attributes, presentFields);
136+
return this.applyDescription(withValidation, modelDef.attributes) as unknown as z.ZodObject<
137+
GetModelSchemaShapeWithOptions<Schema, Model, Options>,
138+
z.core.$strict
139+
>;
65140
}
66141

67142
makeModelCreateSchema<Model extends GetModels<Schema>>(
@@ -114,6 +189,149 @@ class SchemaFactory<Schema extends SchemaDef> {
114189
) as unknown as z.ZodObject<GetModelUpdateFieldsShape<Schema, Model>, z.core.$strict>;
115190
}
116191

192+
// -------------------------------------------------------------------------
193+
// Options-aware field building
194+
// -------------------------------------------------------------------------
195+
196+
/**
197+
* Internal loose options shape used at runtime (we've already validated the
198+
* type-level constraints via the overload signatures).
199+
*/
200+
private buildFieldsWithOptions(model: string, options: RawOptions): Record<string, z.ZodType> {
201+
const { select, include, omit } = options;
202+
const modelDef = this.schema.requireModel(model);
203+
const fields: Record<string, z.ZodType> = {};
204+
205+
if (select) {
206+
// ── select branch ────────────────────────────────────────────────
207+
// Only include fields that are explicitly listed with a truthy value.
208+
for (const [key, value] of Object.entries(select)) {
209+
if (!value) continue; // false → skip
210+
211+
const fieldDef = modelDef.fields[key];
212+
if (!fieldDef) {
213+
throw new SchemaFactoryError(`Field "${key}" does not exist on model "${model}"`);
214+
}
215+
216+
if (fieldDef.relation) {
217+
const subOptions = typeof value === 'object' ? (value as RawOptions) : undefined;
218+
const relSchema = this.makeRelationFieldSchema(fieldDef, subOptions);
219+
fields[key] = this.applyDescription(
220+
this.applyCardinality(relSchema, fieldDef).optional(),
221+
fieldDef.attributes,
222+
);
223+
} else {
224+
if (typeof value === 'object') {
225+
throw new SchemaFactoryError(
226+
`Field "${key}" on model "${model}" is a scalar field and cannot have nested options`,
227+
);
228+
}
229+
fields[key] = this.applyDescription(this.makeScalarFieldSchema(fieldDef), fieldDef.attributes);
230+
}
231+
}
232+
} else {
233+
// ── include + omit branch ────────────────────────────────────────
234+
// Validate omit keys up-front.
235+
if (omit) {
236+
for (const key of Object.keys(omit)) {
237+
const fieldDef = modelDef.fields[key];
238+
if (!fieldDef) {
239+
throw new SchemaFactoryError(`Field "${key}" does not exist on model "${model}"`);
240+
}
241+
if (fieldDef.relation) {
242+
throw new SchemaFactoryError(
243+
`Field "${key}" on model "${model}" is a relation field and cannot be used in "omit"`,
244+
);
245+
}
246+
}
247+
}
248+
249+
// Start with all scalar fields, applying omit exclusions.
250+
for (const [fieldName, fieldDef] of Object.entries(modelDef.fields)) {
251+
if (fieldDef.relation) continue;
252+
253+
if (omit?.[fieldName] === true) continue;
254+
fields[fieldName] = this.applyDescription(this.makeScalarFieldSchema(fieldDef), fieldDef.attributes);
255+
}
256+
257+
// Validate include keys and add relation fields.
258+
if (include) {
259+
for (const [key, value] of Object.entries(include)) {
260+
if (!value) continue; // false → skip
261+
262+
const fieldDef = modelDef.fields[key];
263+
if (!fieldDef) {
264+
throw new SchemaFactoryError(`Field "${key}" does not exist on model "${model}"`);
265+
}
266+
if (!fieldDef.relation) {
267+
throw new SchemaFactoryError(
268+
`Field "${key}" on model "${model}" is not a relation field and cannot be used in "include"`,
269+
);
270+
}
271+
272+
const subOptions = typeof value === 'object' ? (value as RawOptions) : undefined;
273+
const relSchema = this.makeRelationFieldSchema(fieldDef, subOptions);
274+
fields[key] = this.applyDescription(
275+
this.applyCardinality(relSchema, fieldDef).optional(),
276+
fieldDef.attributes,
277+
);
278+
}
279+
}
280+
}
281+
282+
return fields;
283+
}
284+
285+
/**
286+
* Returns the set of scalar field names that will be present in the
287+
* resulting schema after applying `options`. Used by `addCustomValidation`
288+
* to skip `@@validate` rules that reference an absent field.
289+
*
290+
* Only scalar fields matter here because `@@validate` conditions are
291+
* restricted by the ZModel compiler to scalar fields of the same model.
292+
*/
293+
private buildPresentFields(model: string, options: RawOptions): ReadonlySet<string> {
294+
const { select, omit } = options;
295+
const modelDef = this.schema.requireModel(model);
296+
const fields = new Set<string>();
297+
298+
if (select) {
299+
// Only scalar fields explicitly selected with a truthy value.
300+
for (const [key, value] of Object.entries(select)) {
301+
if (!value) continue;
302+
const fieldDef = modelDef.fields[key];
303+
if (fieldDef && !fieldDef.relation) {
304+
fields.add(key);
305+
}
306+
}
307+
} else {
308+
// All scalar fields minus explicitly omitted ones.
309+
for (const [fieldName, fieldDef] of Object.entries(modelDef.fields)) {
310+
if (fieldDef.relation) continue;
311+
if (omit?.[fieldName] === true) continue;
312+
fields.add(fieldName);
313+
}
314+
}
315+
316+
return fields;
317+
}
318+
319+
/**
320+
* Build the inner Zod schema for a relation field, optionally with nested
321+
* query options. Does NOT apply cardinality/optional wrappers — the caller
322+
* does that.
323+
*/
324+
private makeRelationFieldSchema(fieldDef: FieldDef, subOptions?: RawOptions): z.ZodType {
325+
const relatedModelName = fieldDef.type as GetModels<Schema>;
326+
if (subOptions) {
327+
// Recurse: build the related model's schema with its own options.
328+
return this.makeModelSchema(relatedModelName, subOptions as ModelSchemaOptions<Schema, GetModels<Schema>>);
329+
}
330+
// No sub-options: use a lazy reference to the default schema so that
331+
// circular models don't cause infinite recursion at build time.
332+
return z.lazy(() => this.makeModelSchema(relatedModelName));
333+
}
334+
117335
private makeScalarFieldSchema(fieldDef: FieldDef): z.ZodType {
118336
const { type, attributes } = fieldDef;
119337

‎packages/zod/src/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
export { createSchemaFactory } from './factory';
2+
export type { ModelSchemaOptions, GetModelSchemaShapeWithOptions } from './types';
23
export * as ZodUtils from './utils';

0 commit comments

Comments
 (0)