Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor(zod): clean up optionality of generated schema fields #1139

Merged
merged 3 commits into from
Mar 14, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
153 changes: 116 additions & 37 deletions packages/schema/src/plugins/zod/generator.ts
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'
)}; }`
)};
}
ymc9 marked this conversation as resolved.
Show resolved Hide resolved
`
);
}

////////////////////////////////////////////////
// 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
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();
ymc9 marked this conversation as resolved.
Show resolved Hide resolved
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();
ymc9 marked this conversation as resolved.
Show resolved Hide resolved
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();

ymc9 marked this conversation as resolved.
Show resolved Hide resolved
// 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
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);
});
});
Loading