Skip to content

Commit

Permalink
Reimplement createInsertSchema in drizzle-zod
Browse files Browse the repository at this point in the history
  • Loading branch information
@L-Mario564
L-Mario564 committed Nov 15, 2024
1 parent aa76163 commit 2579db6
Show file tree
Hide file tree
Showing 3 changed files with 172 additions and 188 deletions.
66 changes: 59 additions & 7 deletions drizzle-zod/src/create-schema.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,18 +3,27 @@ import { Column, getTableColumns, getViewSelectedFields, is, isTable, SQL } from
import { columnToSchema } from './column';
import { isPgEnum, PgEnum } from 'drizzle-orm/pg-core';
import type { Table, View } from 'drizzle-orm';
import type { CreateSchemaFactoryOptions, CreateSelectSchema } from './types';
import type { CreateInsertSchema, CreateSchemaFactoryOptions, CreateSelectSchema } from './types';

function getColumns(tableLike: Table | View) {
return isTable(tableLike) ? getTableColumns(tableLike) : getViewSelectedFields(tableLike);
}

function handleColumns(columns: Record<string, any>, refinements: Record<string, any>, factory?: CreateSchemaFactoryOptions): z.ZodTypeAny {
function handleColumns(
columns: Record<string, any>,
refinements: Record<string, any>,
conditions: {
never: (column?: Column) => boolean;
optional: (column: Column) => boolean;
nullable: (column: Column) => boolean;
},
factory?: CreateSchemaFactoryOptions
): z.ZodTypeAny {
const columnSchemas: Record<string, z.ZodTypeAny> = {};

for (const [key, selected] of Object.entries(columns)) {
if (!is(selected, Column) && !is(selected, SQL) && !is(selected, SQL.Aliased) && typeof selected === 'object') {
columnSchemas[key] = handleColumns(selected, refinements[key] ?? {});
columnSchemas[key] = handleColumns(selected, refinements[key] ?? {}, conditions, factory);
continue;
}

Expand All @@ -27,7 +36,22 @@ function handleColumns(columns: Record<string, any>, refinements: Record<string,
const column = is(selected, Column) ? selected : undefined;
const schema = !!column ? columnToSchema(column, factory?.zodInstance ?? z) : z.any();
const refined = typeof refinement === 'function' ? refinement(schema) : schema;
columnSchemas[key] = !!column && !column.notNull ? refined.nullable() : refined;

if (conditions.never(column)) {
continue;
} else {
columnSchemas[key] = refined;
};

if (column) {
if (conditions.nullable(column)) {
columnSchemas[key] = columnSchemas[key]!.nullable();
}

if (conditions.optional(column)) {
columnSchemas[key] = columnSchemas[key]!.optional();
}
}
}

return z.object(columnSchemas) as any;
Expand All @@ -38,6 +62,18 @@ function handleEnum(enum_: PgEnum<any>, factory?: CreateSchemaFactoryOptions) {
return zod.enum(enum_.enumValues);
}

const selectConditions = {
never: () => false,
optional: () => false,
nullable: (column: Column) => !column.notNull
};

const insertConditions = {
never: (column?: Column) => column?.generated?.type === 'always' || column?.generatedIdentity?.type === 'always',
optional: (column: Column) => !column.notNull || (column.notNull && column.hasDefault),
nullable: (column: Column) => !column.notNull
}

export const createSelectSchema: CreateSelectSchema = (
entity: Table | View | PgEnum<[string, ...string[]]>,
refine?: Record<string, any>
Expand All @@ -46,7 +82,15 @@ export const createSelectSchema: CreateSelectSchema = (
return handleEnum(entity);
}
const columns = getColumns(entity);
return handleColumns(columns, refine ?? {}) as any;
return handleColumns(columns, refine ?? {}, selectConditions) as any;
}

export const createInsertSchema: CreateInsertSchema = (
entity: Table,
refine?: Record<string, any>
) => {
const columns = getColumns(entity);
return handleColumns(columns, refine ?? {}, insertConditions) as any;
}

export function createSchemaFactory(options?: CreateSchemaFactoryOptions) {
Expand All @@ -58,8 +102,16 @@ export function createSchemaFactory(options?: CreateSchemaFactoryOptions) {
return handleEnum(entity, options);
}
const columns = getColumns(entity);
return handleColumns(columns, refine ?? {}, options) as any;
return handleColumns(columns, refine ?? {}, selectConditions, options) as any;
}

const createInsertSchema: CreateInsertSchema = (
entity: Table,
refine?: Record<string, any>
) => {
const columns = getColumns(entity);
return handleColumns(columns, refine ?? {}, insertConditions, options) as any;
}

return { createSelectSchema };
return { createSelectSchema, createInsertSchema };
}
122 changes: 100 additions & 22 deletions drizzle-zod/src/types.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,17 @@ type EnumHasAtLeastOneValue<TEnum extends [string, ...string[]] | undefined> =
: false
: false;

type ColumnIsGeneratedAlwaysAs<TColumn extends Column> =
TColumn['_']['generated'] extends infer TGenerated extends { type: string }
? TGenerated['type'] extends 'always'
? true
: false
: false;

type RemoveNever<T> = {
[K in keyof T as T[K] extends never ? never : K]: T[K];
};

export type GetZodType<
TData,
TDataType extends string,
Expand Down Expand Up @@ -69,55 +80,122 @@ export type BuildRefine<TColumns extends Record<string, any>> = BuildRefineColum
}
: never;

export type BuildSelectSchema<
type HandleRefinement<
TRefinement extends z.ZodTypeAny | ((schema: z.ZodTypeAny) => z.ZodTypeAny),
TColumn extends Column
> = TRefinement extends (schema: z.ZodTypeAny) => z.ZodTypeAny
? TColumn['_']['notNull'] extends true
? ReturnType<TRefinement>
: z.ZodNullable<ReturnType<TRefinement>>
: TRefinement;

type HandleColumn<
TType extends 'select' | 'insert' | 'update',
TColumn extends Column
> = GetZodType<
TColumn['_']['data'],
TColumn['_']['dataType'],
TColumn['_'] extends { enumValues: [string, ...string[]] } ? TColumn['_']['enumValues'] : undefined
> extends infer TSchema extends z.ZodTypeAny
? TSchema extends z.ZodAny
? z.ZodAny
: TType extends 'select'
? HandleSelectColumn<TSchema, TColumn>
: TType extends 'insert'
? HandleInsertColumn<TSchema, TColumn>
: TType extends 'update'
? HandleUpdateColumn<TSchema, TColumn>
: TSchema
: z.ZodAny

type HandleSelectColumn<
TSchema extends z.ZodTypeAny,
TColumn extends Column
> = TColumn['_']['notNull'] extends true
? TSchema
: z.ZodNullable<TSchema>;

type HandleInsertColumn<
TSchema extends z.ZodTypeAny,
TColumn extends Column
> = ColumnIsGeneratedAlwaysAs<TColumn> extends true
? never
: TColumn['_']['notNull'] extends true
? TColumn['_']['hasDefault'] extends true
? z.ZodOptional<TSchema>
: TSchema
: z.ZodOptional<z.ZodNullable<TSchema>>;

type HandleUpdateColumn<
TSchema extends z.ZodTypeAny,
TColumn extends Column
> = ColumnIsGeneratedAlwaysAs<TColumn> extends true
? never
: TColumn['_']['notNull'] extends true
? z.ZodOptional<TSchema>
: z.ZodOptional<z.ZodNullable<TSchema>>;

export type BuildSchema<
TType extends 'select' | 'insert' | 'update',
TColumns extends Record<string, any>,
TRefinements extends Record<string, ((schema: z.ZodTypeAny) => z.ZodTypeAny | z.ZodTypeAny)> | undefined
TRefinements extends Record<string, any> | undefined
> = z.ZodObject<
Simplify<{
Simplify<RemoveNever<{
[K in keyof TColumns]:
TColumns[K] extends infer TColumn extends Column
? TRefinements[Assume<K, keyof TRefinements>] extends infer TRefinement extends z.ZodTypeAny | ((schema: z.ZodTypeAny) => z.ZodTypeAny)
? TRefinement extends (schema: z.ZodTypeAny) => z.ZodTypeAny
? TColumn['_']['notNull'] extends true ? ReturnType<TRefinement> : z.ZodNullable<ReturnType<TRefinement>>
: TRefinement
: GetZodType<
TColumn['_']['data'],
TColumn['_']['dataType'],
TColumn['_'] extends { enumValues: [string, ...string[]] } ? TColumn['_']['enumValues'] : undefined
> extends infer TSchema extends z.ZodTypeAny
? TColumn['_']['notNull'] extends true
? TSchema
: z.ZodNullable<TSchema>
: z.ZodAny
? TRefinements extends object
? TRefinements[Assume<K, keyof TRefinements>] extends infer TRefinement extends z.ZodTypeAny | ((schema: z.ZodTypeAny) => z.ZodTypeAny)
? HandleRefinement<TRefinement, TColumn>
: HandleColumn<TType, TColumn>
: HandleColumn<TType, TColumn>
: TColumns[K] extends infer TObject extends SelectedFieldsFlat<Column>
? BuildSelectSchema<TObject, Assume<TRefinements[Assume<K, keyof TRefinements>], Record<string, any>>>
? BuildSchema<
TType,
TObject,
TRefinements extends object
? TRefinements[Assume<K, keyof TRefinements>] extends infer TNestedRefinements extends object
? TNestedRefinements
: undefined
: undefined
>
: z.ZodAny
}>,
}>>,
'strip'
>;

export interface CreateSelectSchema {
<TView extends View>(view: TView): BuildSelectSchema<TView['_']['selectedFields'], never>;
<TView extends View>(view: TView): BuildSchema<'select', TView['_']['selectedFields'], undefined>;
<
TView extends View,
TRefine extends BuildRefine<TView['_']['selectedFields']>
>(
view: TView,
refine: TRefine
): BuildSelectSchema<TView['_']['selectedFields'], TRefine>;
): BuildSchema<'select', TView['_']['selectedFields'], TRefine>;

<TTable extends Table>(table: TTable): BuildSelectSchema<TTable['_']['columns'], never>;
<TTable extends Table>(table: TTable): BuildSchema<'select', TTable['_']['columns'], undefined>;
<
TTable extends Table,
TRefine extends BuildRefine<TTable['_']['columns']>
>(
table: TTable,
refine?: TRefine
): BuildSelectSchema<TTable['_']['columns'], TRefine>;
): BuildSchema<'select', TTable['_']['columns'], TRefine>;

<TEnum extends PgEnum<any>>(enum_: TEnum): z.ZodEnum<TEnum['enumValues']>;
}

export interface CreateInsertSchema {
<TTable extends Table>(table: TTable): BuildSchema<'insert', TTable['_']['columns'], undefined>;
<
TTable extends Table,
TRefine extends BuildRefine<TTable['_']['columns']>
>(
table: TTable,
refine?: TRefine
): BuildSchema<'insert', TTable['_']['columns'], TRefine>;
}

export interface CreateSchemaFactoryOptions {
zodInstance?: any;
}
Loading

0 comments on commit 2579db6

Please sign in to comment.