Skip to content

Commit

Permalink
refactor(zod): clean up optionality of generated schema fields (#1139)
Browse files Browse the repository at this point in the history
  • Loading branch information
@ymc9
ymc9 authored Mar 14, 2024
1 parent bbea0a9 commit 8099793
Show file tree
Hide file tree
Showing 3 changed files with 171 additions and 40 deletions.
153 changes: 116 additions & 37 deletions packages/schema/src/plugins/zod/generator.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import { name } from '.';
import { getDefaultOutputFolder } from '../plugin-utils';
import Transformer from './transformer';
import removeDir from './utils/removeDir';
import { getFieldSchemaDefault, makeFieldSchema, makeValidationRefinements } from './utils/schema-gen';
import { makeFieldSchema, makeValidationRefinements } from './utils/schema-gen';

export class ZodSchemaGenerator {
private readonly sourceFiles: SourceFile[] = [];
Expand Down Expand Up @@ -311,7 +311,7 @@ export class ZodSchemaGenerator {
writer.writeLine(`import { Decimal } from 'decimal.js';`);
}

// base schema
// base schema - including all scalar fields, with optionality following the schema
writer.write(`const baseSchema = z.object(`);
writer.inlineBlock(() => {
scalarFields.forEach((field) => {
Expand All @@ -325,11 +325,11 @@ export class ZodSchemaGenerator {
let relationSchema: string | undefined;
let fkSchema: string | undefined;

if (relations.length > 0 || fkFields.length > 0) {
if (relations.length > 0) {
relationSchema = 'relationSchema';
writer.write(`const ${relationSchema} = z.object(`);
writer.inlineBlock(() => {
[...relations, ...fkFields].forEach((field) => {
[...relations].forEach((field) => {
writer.writeLine(`${field.name}: ${makeFieldSchema(field)},`);
});
});
Expand All @@ -353,21 +353,23 @@ export class ZodSchemaGenerator {
if (refinements.length > 0) {
refineFuncName = `refine${upperCaseFirst(model.name)}`;
writer.writeLine(
`export function ${refineFuncName}<T, D extends z.ZodTypeDef>(schema: z.ZodType<T, D, T>) { return schema${refinements.join(
`
/**
* Schema refinement function for applying \`@@validate\` rules.
*/
export function ${refineFuncName}<T, D extends z.ZodTypeDef>(schema: z.ZodType<T, D, T>) { return schema${refinements.join(
'\n'
)}; }`
)};
}
`
);
}

////////////////////////////////////////////////
// 1. Model schema
////////////////////////////////////////////////
const fieldsWithoutDefault = scalarFields.filter((f) => !getFieldSchemaDefault(f));
// mark fields without default value as optional
let modelSchema = this.makePartial(
'baseSchema',
fieldsWithoutDefault.length < scalarFields.length ? fieldsWithoutDefault.map((f) => f.name) : undefined
);

let modelSchema = 'baseSchema';

// omit fields
const fieldsToOmit = scalarFields.filter((field) => hasAttribute(field, '@omit'));
Expand All @@ -378,23 +380,46 @@ export class ZodSchemaGenerator {
);
}

if (relationSchema) {
// export schema with only scalar fields
const modelScalarSchema = `${upperCaseFirst(model.name)}ScalarSchema`;
writer.writeLine(`export const ${modelScalarSchema} = ${modelSchema};`);
modelSchema = modelScalarSchema;
// export schema with only scalar fields: `[Model]ScalarSchema`
const modelScalarSchema = `${upperCaseFirst(model.name)}ScalarSchema`;
writer.writeLine(`
/**
* \`${model.name}\` schema excluding foreign keys and relations.
*/
export const ${modelScalarSchema} = ${modelSchema};
`);
modelSchema = modelScalarSchema;

// merge fk fields
if (fkSchema) {
modelSchema = this.makeMerge(modelSchema, fkSchema);
}

// merge relations
// merge relation fields (all optional)
if (relationSchema) {
modelSchema = this.makeMerge(modelSchema, this.makePartial(relationSchema));
}

// refine
if (refineFuncName) {
// export a schema without refinement for extensibility: `[Model]WithoutRefineSchema`
const noRefineSchema = `${upperCaseFirst(model.name)}WithoutRefineSchema`;
writer.writeLine(`export const ${noRefineSchema} = ${modelSchema};`);
writer.writeLine(`
/**
* \`${model.name}\` schema prior to calling \`.refine()\` for extensibility.
*/
export const ${noRefineSchema} = ${modelSchema};
`);
modelSchema = `${refineFuncName}(${noRefineSchema})`;
}
writer.writeLine(`export const ${upperCaseFirst(model.name)}Schema = ${modelSchema};`);

// export the final model schema: `[Model]Schema`
writer.writeLine(`
/**
* \`${model.name}\` schema including all fields (scalar, foreign key, and relations) and validations.
*/
export const ${upperCaseFirst(model.name)}Schema = ${modelSchema};
`);

////////////////////////////////////////////////
// 2. Prisma create & update
Expand All @@ -405,7 +430,13 @@ export class ZodSchemaGenerator {
if (refineFuncName) {
prismaCreateSchema = `${refineFuncName}(${prismaCreateSchema})`;
}
writer.writeLine(`export const ${upperCaseFirst(model.name)}PrismaCreateSchema = ${prismaCreateSchema};`);
writer.writeLine(`
/**
* Schema used for validating Prisma create input. For internal use only.
* @private
*/
export const ${upperCaseFirst(model.name)}PrismaCreateSchema = ${prismaCreateSchema};
`);

// schema for validating prisma update input (all fields optional)
// note numeric fields can be simple update or atomic operations
Expand All @@ -424,61 +455,109 @@ export class ZodSchemaGenerator {
if (refineFuncName) {
prismaUpdateSchema = `${refineFuncName}(${prismaUpdateSchema})`;
}
writer.writeLine(`export const ${upperCaseFirst(model.name)}PrismaUpdateSchema = ${prismaUpdateSchema};`);
writer.writeLine(
`
/**
* Schema used for validating Prisma update input. For internal use only.
* @private
*/
export const ${upperCaseFirst(model.name)}PrismaUpdateSchema = ${prismaUpdateSchema};
`
);

////////////////////////////////////////////////
// 3. Create schema
////////////////////////////////////////////////

let createSchema = 'baseSchema';
const fieldsWithDefault = scalarFields.filter(
(field) => hasAttribute(field, '@default') || hasAttribute(field, '@updatedAt') || field.type.array
);

// mark fields with default as optional
if (fieldsWithDefault.length > 0) {
createSchema = this.makePartial(
createSchema,
fieldsWithDefault.map((f) => f.name)
);
}

if (fkSchema) {
// export schema with only scalar fields
const createScalarSchema = `${upperCaseFirst(model.name)}CreateScalarSchema`;
writer.writeLine(`export const ${createScalarSchema} = ${createSchema};`);
// export schema with only scalar fields: `[Model]CreateScalarSchema`
const createScalarSchema = `${upperCaseFirst(model.name)}CreateScalarSchema`;
writer.writeLine(`
/**
* \`${model.name}\` schema for create operations excluding foreign keys and relations.
*/
export const ${createScalarSchema} = ${createSchema};
`);

if (fkSchema) {
// merge fk fields
createSchema = this.makeMerge(createScalarSchema, fkSchema);
}

if (refineFuncName) {
// export a schema without refinement for extensibility
// export a schema without refinement for extensibility: `[Model]CreateWithoutRefineSchema`
const noRefineSchema = `${upperCaseFirst(model.name)}CreateWithoutRefineSchema`;
writer.writeLine(`export const ${noRefineSchema} = ${createSchema};`);
writer.writeLine(`
/**
* \`${model.name}\` schema for create operations prior to calling \`.refine()\` for extensibility.
*/
export const ${noRefineSchema} = ${createSchema};
`);
createSchema = `${refineFuncName}(${noRefineSchema})`;
}
writer.writeLine(`export const ${upperCaseFirst(model.name)}CreateSchema = ${createSchema};`);

// export the final create schema: `[Model]CreateSchema`
writer.writeLine(`
/**
* \`${model.name}\` schema for create operations including scalar fields, foreign key fields, and validations.
*/
export const ${upperCaseFirst(model.name)}CreateSchema = ${createSchema};
`);

////////////////////////////////////////////////
// 3. Update schema
////////////////////////////////////////////////

// for update all fields are optional
let updateSchema = this.makePartial('baseSchema');

if (fkSchema) {
// export schema with only scalar fields
const updateScalarSchema = `${upperCaseFirst(model.name)}UpdateScalarSchema`;
writer.writeLine(`export const ${updateScalarSchema} = ${updateSchema};`);
updateSchema = updateScalarSchema;
// export schema with only scalar fields: `[Model]UpdateScalarSchema`
const updateScalarSchema = `${upperCaseFirst(model.name)}UpdateScalarSchema`;
writer.writeLine(`
/**
* \`${model.name}\` schema for update operations excluding foreign keys and relations.
*/
export const ${updateScalarSchema} = ${updateSchema};
`);

updateSchema = updateScalarSchema;

if (fkSchema) {
// merge fk fields
updateSchema = this.makeMerge(updateSchema, this.makePartial(fkSchema));
}

if (refineFuncName) {
// export a schema without refinement for extensibility
// export a schema without refinement for extensibility: `[Model]UpdateWithoutRefineSchema`
const noRefineSchema = `${upperCaseFirst(model.name)}UpdateWithoutRefineSchema`;
writer.writeLine(`export const ${noRefineSchema} = ${updateSchema};`);
writer.writeLine(`
/**
* \`${model.name}\` schema for update operations prior to calling \`.refine()\` for extensibility.
*/
export const ${noRefineSchema} = ${updateSchema};
`);
updateSchema = `${refineFuncName}(${noRefineSchema})`;
}
writer.writeLine(`export const ${upperCaseFirst(model.name)}UpdateSchema = ${updateSchema};`);

// export the final update schema: `[Model]UpdateSchema`
writer.writeLine(`
/**
* \`${model.name}\` schema for update operations including scalar fields, foreign key fields, and validations.
*/
export const ${upperCaseFirst(model.name)}UpdateSchema = ${updateSchema};
`);
});

return schemaName;
Expand Down
54 changes: 53 additions & 1 deletion tests/integration/tests/plugins/zod.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,9 @@ describe('Zod plugin tests', () => {
password String @omit
role Role @default(USER)
posts Post[]
age Int?
@@validate(length(password, 6, 20))
}
model Post {
Expand All @@ -61,8 +64,14 @@ describe('Zod plugin tests', () => {
{ addPrelude: false, pushDb: false }
);
const schemas = zodSchemas.models;
expect(schemas.UserScalarSchema).toBeTruthy();
expect(schemas.UserWithoutRefineSchema).toBeTruthy();
expect(schemas.UserSchema).toBeTruthy();
expect(schemas.UserCreateScalarSchema).toBeTruthy();
expect(schemas.UserCreateWithoutRefineSchema).toBeTruthy();
expect(schemas.UserCreateSchema).toBeTruthy();
expect(schemas.UserUpdateScalarSchema).toBeTruthy();
expect(schemas.UserUpdateWithoutRefineSchema).toBeTruthy();
expect(schemas.UserUpdateSchema).toBeTruthy();
expect(schemas.UserPrismaCreateSchema).toBeTruthy();
expect(schemas.UserPrismaUpdateSchema).toBeTruthy();
Expand All @@ -75,6 +84,16 @@ describe('Zod plugin tests', () => {
expect(
schemas.UserCreateSchema.safeParse({ email: '[email protected]', password: 'abc123' }).success
).toBeTruthy();
expect(
schemas.UserCreateSchema.safeParse({ email: '[email protected]', role: 'ADMIN', password: 'abc' }).success
).toBeFalsy();
expect(
schemas.UserCreateWithoutRefineSchema.safeParse({
email: '[email protected]',
role: 'ADMIN',
password: 'abc',
}).success
).toBeTruthy();
expect(
schemas.UserCreateSchema.safeParse({ email: '[email protected]', role: 'ADMIN', password: 'abc123' }).success
).toBeTruthy();
Expand All @@ -90,6 +109,8 @@ describe('Zod plugin tests', () => {
expect(schemas.UserUpdateSchema.safeParse({}).success).toBeTruthy();
expect(schemas.UserUpdateSchema.safeParse({ email: '[email protected]' }).success).toBeFalsy();
expect(schemas.UserUpdateSchema.safeParse({ email: '[email protected]' }).success).toBeTruthy();
expect(schemas.UserUpdateSchema.safeParse({ password: 'pas' }).success).toBeFalsy();
expect(schemas.UserUpdateWithoutRefineSchema.safeParse({ password: 'pas' }).success).toBeTruthy();
expect(schemas.UserUpdateSchema.safeParse({ password: 'password456' }).success).toBeTruthy();

// update unchecked
Expand All @@ -98,7 +119,25 @@ describe('Zod plugin tests', () => {
).toBeTruthy();

// model schema
expect(schemas.UserSchema.safeParse({ email: '[email protected]', role: 'ADMIN' }).success).toBeTruthy();

// missing fields
expect(
schemas.UserSchema.safeParse({
id: 1,
email: '[email protected]',
}).success
).toBeFalsy();

expect(
schemas.UserSchema.safeParse({
id: 1,
createdAt: new Date(),
updatedAt: new Date(),
email: '[email protected]',
role: 'ADMIN',
}).success
).toBeTruthy();

// without omitted field
expect(
schemas.UserSchema.safeParse({
Expand All @@ -109,6 +148,19 @@ describe('Zod plugin tests', () => {
updatedAt: new Date(),
}).success
).toBeTruthy();

// with optional field
expect(
schemas.UserSchema.safeParse({
id: 1,
email: '[email protected]',
role: 'ADMIN',
createdAt: new Date(),
updatedAt: new Date(),
age: 18,
}).success
).toBeTruthy();

// with omitted field
const withPwd = schemas.UserSchema.safeParse({
id: 1,
Expand Down
4 changes: 2 additions & 2 deletions tests/integration/tests/regression/issue-886.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,10 @@ describe('Regression: issue 886', () => {
`
);

const r = zodSchemas.models.ModelSchema.parse({});
const r = zodSchemas.models.ModelSchema.parse({ id: 1 });
expect(r.a).toBe(100);
expect(r.b).toBe('');
expect(r.c).toBeInstanceOf(Date);
expect(r.id).toBeUndefined();
expect(r.id).toBe(1);
});
});

0 comments on commit 8099793

Please sign in to comment.