From 016000a4588226f8117b72bbafeff7f0c6a8605b Mon Sep 17 00:00:00 2001 From: Mario564 Date: Thu, 26 Dec 2024 10:33:49 -0800 Subject: [PATCH 1/2] Implement type coercion in drizzle-zod --- drizzle-zod/src/column.ts | 43 +++++++++++++------ drizzle-zod/src/schema.ts | 2 +- drizzle-zod/src/schema.types.ts | 1 + drizzle-zod/tests/mysql.test.ts | 61 ++++++++++++++++++++++++++- drizzle-zod/tests/pg.test.ts | 61 ++++++++++++++++++++++++++- drizzle-zod/tests/singlestore.test.ts | 61 ++++++++++++++++++++++++++- drizzle-zod/tests/sqlite.test.ts | 58 ++++++++++++++++++++++++- drizzle-zod/tests/utils.ts | 1 + 8 files changed, 270 insertions(+), 18 deletions(-) diff --git a/drizzle-zod/src/column.ts b/drizzle-zod/src/column.ts index 23bc3c142..b6c50b85a 100644 --- a/drizzle-zod/src/column.ts +++ b/drizzle-zod/src/column.ts @@ -54,8 +54,9 @@ import type { } from 'drizzle-orm/singlestore-core'; import type { SQLiteInteger, SQLiteReal, SQLiteText } from 'drizzle-orm/sqlite-core'; import { z } from 'zod'; -import type { z as zod } from 'zod'; +import { z as zod } from 'zod'; import { CONSTANTS } from './constants.ts'; +import type { CreateSchemaFactoryOptions } from './schema.types.ts'; import { isColumnType, isWithEnum } from './utils.ts'; import type { Json } from './utils.ts'; @@ -65,7 +66,9 @@ export const jsonSchema: z.ZodType = z.lazy(() => ); export const bufferSchema: z.ZodType = z.custom((v) => v instanceof Buffer); // eslint-disable-line no-instanceof/no-instanceof -export function columnToSchema(column: Column, z: typeof zod): z.ZodTypeAny { +export function columnToSchema(column: Column, factory: CreateSchemaFactoryOptions | undefined): z.ZodTypeAny { + const z = factory?.zodInstance ?? zod; + const coerce = factory?.coerce ?? {}; let schema!: z.ZodTypeAny; if (isWithEnum(column)) { @@ -98,15 +101,15 @@ export function columnToSchema(column: Column, z: typeof zod): z.ZodTypeAny { } else if (column.dataType === 'array') { schema = z.array(z.any()); } else if (column.dataType === 'number') { - schema = numberColumnToSchema(column, z); + schema = numberColumnToSchema(column, z, coerce); } else if (column.dataType === 'bigint') { - schema = bigintColumnToSchema(column, z); + schema = bigintColumnToSchema(column, z, coerce); } else if (column.dataType === 'boolean') { - schema = z.boolean(); + schema = coerce.boolean ? z.coerce.boolean() : z.boolean(); } else if (column.dataType === 'date') { - schema = z.date(); + schema = coerce.date ? z.coerce.date() : z.date(); } else if (column.dataType === 'string') { - schema = stringColumnToSchema(column, z); + schema = stringColumnToSchema(column, z, coerce); } else if (column.dataType === 'json') { schema = jsonSchema; } else if (column.dataType === 'custom') { @@ -123,7 +126,11 @@ export function columnToSchema(column: Column, z: typeof zod): z.ZodTypeAny { return schema; } -function numberColumnToSchema(column: Column, z: typeof zod): z.ZodTypeAny { +function numberColumnToSchema( + column: Column, + z: typeof zod, + coerce: CreateSchemaFactoryOptions['coerce'], +): z.ZodTypeAny { let unsigned = column.getSQLType().includes('unsigned'); let min!: number; let max!: number; @@ -223,19 +230,29 @@ function numberColumnToSchema(column: Column, z: typeof zod): z.ZodTypeAny { max = Number.MAX_SAFE_INTEGER; } - const schema = z.number().min(min).max(max); + let schema = coerce?.number ? z.coerce.number() : z.number(); + schema = schema.min(min).max(max); return integer ? schema.int() : schema; } -function bigintColumnToSchema(column: Column, z: typeof zod): z.ZodTypeAny { +function bigintColumnToSchema( + column: Column, + z: typeof zod, + coerce: CreateSchemaFactoryOptions['coerce'], +): z.ZodTypeAny { const unsigned = column.getSQLType().includes('unsigned'); const min = unsigned ? 0n : CONSTANTS.INT64_MIN; const max = unsigned ? CONSTANTS.INT64_UNSIGNED_MAX : CONSTANTS.INT64_MAX; - return z.bigint().min(min).max(max); + const schema = coerce?.bigint ? z.coerce.bigint() : z.bigint(); + return schema.min(min).max(max); } -function stringColumnToSchema(column: Column, z: typeof zod): z.ZodTypeAny { +function stringColumnToSchema( + column: Column, + z: typeof zod, + coerce: CreateSchemaFactoryOptions['coerce'], +): z.ZodTypeAny { if (isColumnType>>(column, ['PgUUID'])) { return z.string().uuid(); } @@ -278,7 +295,7 @@ function stringColumnToSchema(column: Column, z: typeof zod): z.ZodTypeAny { max = column.dimensions; } - let schema = z.string(); + let schema = coerce?.string ? z.coerce.string() : z.string(); schema = regex ? schema.regex(regex) : schema; return max && fixed ? schema.length(max) : max ? schema.max(max) : schema; } diff --git a/drizzle-zod/src/schema.ts b/drizzle-zod/src/schema.ts index 67a9cb733..40c7e891c 100644 --- a/drizzle-zod/src/schema.ts +++ b/drizzle-zod/src/schema.ts @@ -38,7 +38,7 @@ function handleColumns( } const column = is(selected, Column) ? selected : undefined; - const schema = column ? columnToSchema(column, factory?.zodInstance ?? z) : z.any(); + const schema = column ? columnToSchema(column, factory) : z.any(); const refined = typeof refinement === 'function' ? refinement(schema) : schema; if (conditions.never(column)) { diff --git a/drizzle-zod/src/schema.types.ts b/drizzle-zod/src/schema.types.ts index 5873cd2a3..4d67d1d6f 100644 --- a/drizzle-zod/src/schema.types.ts +++ b/drizzle-zod/src/schema.types.ts @@ -49,4 +49,5 @@ export interface CreateUpdateSchema { export interface CreateSchemaFactoryOptions { zodInstance?: any; + coerce?: Partial>; } diff --git a/drizzle-zod/tests/mysql.test.ts b/drizzle-zod/tests/mysql.test.ts index 73ba48dae..18f8b3d50 100644 --- a/drizzle-zod/tests/mysql.test.ts +++ b/drizzle-zod/tests/mysql.test.ts @@ -4,7 +4,7 @@ import { test } from 'vitest'; import { z } from 'zod'; import { jsonSchema } from '~/column.ts'; import { CONSTANTS } from '~/constants.ts'; -import { createInsertSchema, createSelectSchema, createUpdateSchema } from '../src'; +import { createInsertSchema, createSchemaFactory, createSelectSchema, createUpdateSchema } from '../src'; import { Expect, expectSchemaShape } from './utils.ts'; const intSchema = z.number().min(CONSTANTS.INT32_MIN).max(CONSTANTS.INT32_MAX).int(); @@ -454,6 +454,65 @@ test('all data types', (t) => { Expect>(); }); +test('type coercion - all', (t) => { + const table = mysqlTable('test', ({ + bigint, + boolean, + timestamp, + int, + text, + }) => ({ + bigint: bigint({ mode: 'bigint' }).notNull(), + boolean: boolean().notNull(), + timestamp: timestamp().notNull(), + int: int().notNull(), + text: text().notNull(), + })); + + const { createSelectSchema } = createSchemaFactory({ + coerce: { + bigint: true, + boolean: true, + date: true, + number: true, + string: true, + }, + }); + const result = createSelectSchema(table); + const expected = z.object({ + bigint: z.coerce.bigint().min(CONSTANTS.INT64_MIN).max(CONSTANTS.INT64_MAX), + boolean: z.coerce.boolean(), + timestamp: z.coerce.date(), + int: z.coerce.number().min(CONSTANTS.INT32_MIN).max(CONSTANTS.INT32_MAX).int(), + text: z.coerce.string().max(CONSTANTS.INT16_UNSIGNED_MAX), + }); + expectSchemaShape(t, expected).from(result); + Expect>(); +}); + +test('type coercion - mixed', (t) => { + const table = mysqlTable('test', ({ + timestamp, + int, + }) => ({ + timestamp: timestamp().notNull(), + int: int().notNull(), + })); + + const { createSelectSchema } = createSchemaFactory({ + coerce: { + date: true, + }, + }); + const result = createSelectSchema(table); + const expected = z.object({ + timestamp: z.coerce.date(), + int: z.number().min(CONSTANTS.INT32_MIN).max(CONSTANTS.INT32_MAX).int(), + }); + expectSchemaShape(t, expected).from(result); + Expect>(); +}); + /* Disallow unknown keys in table refinement - select */ { const table = mysqlTable('test', { id: int() }); // @ts-expect-error diff --git a/drizzle-zod/tests/pg.test.ts b/drizzle-zod/tests/pg.test.ts index 7964f65d6..68642f073 100644 --- a/drizzle-zod/tests/pg.test.ts +++ b/drizzle-zod/tests/pg.test.ts @@ -14,7 +14,7 @@ import { test } from 'vitest'; import { z } from 'zod'; import { jsonSchema } from '~/column.ts'; import { CONSTANTS } from '~/constants.ts'; -import { createInsertSchema, createSelectSchema, createUpdateSchema } from '../src'; +import { createInsertSchema, createSchemaFactory, createSelectSchema, createUpdateSchema } from '../src'; import { Expect, expectEnumValues, expectSchemaShape } from './utils.ts'; const integerSchema = z.number().min(CONSTANTS.INT32_MIN).max(CONSTANTS.INT32_MAX).int(); @@ -500,6 +500,65 @@ test('all data types', (t) => { Expect>(); }); +test('type coercion - all', (t) => { + const table = pgTable('test', ({ + bigint, + boolean, + timestamp, + integer, + text, + }) => ({ + bigint: bigint({ mode: 'bigint' }).notNull(), + boolean: boolean().notNull(), + timestamp: timestamp().notNull(), + integer: integer().notNull(), + text: text().notNull(), + })); + + const { createSelectSchema } = createSchemaFactory({ + coerce: { + bigint: true, + boolean: true, + date: true, + number: true, + string: true, + }, + }); + const result = createSelectSchema(table); + const expected = z.object({ + bigint: z.coerce.bigint().min(CONSTANTS.INT64_MIN).max(CONSTANTS.INT64_MAX), + boolean: z.coerce.boolean(), + timestamp: z.coerce.date(), + integer: z.coerce.number().min(CONSTANTS.INT32_MIN).max(CONSTANTS.INT32_MAX).int(), + text: z.coerce.string(), + }); + expectSchemaShape(t, expected).from(result); + Expect>(); +}); + +test('type coercion - mixed', (t) => { + const table = pgTable('test', ({ + timestamp, + integer, + }) => ({ + timestamp: timestamp().notNull(), + integer: integer().notNull(), + })); + + const { createSelectSchema } = createSchemaFactory({ + coerce: { + date: true, + }, + }); + const result = createSelectSchema(table); + const expected = z.object({ + timestamp: z.coerce.date(), + integer: z.number().min(CONSTANTS.INT32_MIN).max(CONSTANTS.INT32_MAX).int(), + }); + expectSchemaShape(t, expected).from(result); + Expect>(); +}); + /* Disallow unknown keys in table refinement - select */ { const table = pgTable('test', { id: integer() }); // @ts-expect-error diff --git a/drizzle-zod/tests/singlestore.test.ts b/drizzle-zod/tests/singlestore.test.ts index b91c74be8..f4824e67e 100644 --- a/drizzle-zod/tests/singlestore.test.ts +++ b/drizzle-zod/tests/singlestore.test.ts @@ -4,7 +4,7 @@ import { test } from 'vitest'; import { z } from 'zod'; import { jsonSchema } from '~/column.ts'; import { CONSTANTS } from '~/constants.ts'; -import { createInsertSchema, createSelectSchema, createUpdateSchema } from '../src'; +import { createInsertSchema, createSchemaFactory, createSelectSchema, createUpdateSchema } from '../src'; import { Expect, expectSchemaShape } from './utils.ts'; const intSchema = z.number().min(CONSTANTS.INT32_MIN).max(CONSTANTS.INT32_MAX).int(); @@ -456,6 +456,65 @@ test('all data types', (t) => { Expect>(); }); +test('type coercion - all', (t) => { + const table = singlestoreTable('test', ({ + bigint, + boolean, + timestamp, + int, + text, + }) => ({ + bigint: bigint({ mode: 'bigint' }).notNull(), + boolean: boolean().notNull(), + timestamp: timestamp().notNull(), + int: int().notNull(), + text: text().notNull(), + })); + + const { createSelectSchema } = createSchemaFactory({ + coerce: { + bigint: true, + boolean: true, + date: true, + number: true, + string: true, + }, + }); + const result = createSelectSchema(table); + const expected = z.object({ + bigint: z.coerce.bigint().min(CONSTANTS.INT64_MIN).max(CONSTANTS.INT64_MAX), + boolean: z.coerce.boolean(), + timestamp: z.coerce.date(), + int: z.coerce.number().min(CONSTANTS.INT32_MIN).max(CONSTANTS.INT32_MAX).int(), + text: z.coerce.string().max(CONSTANTS.INT16_UNSIGNED_MAX), + }); + expectSchemaShape(t, expected).from(result); + Expect>(); +}); + +test('type coercion - mixed', (t) => { + const table = singlestoreTable('test', ({ + timestamp, + int, + }) => ({ + timestamp: timestamp().notNull(), + int: int().notNull(), + })); + + const { createSelectSchema } = createSchemaFactory({ + coerce: { + date: true, + }, + }); + const result = createSelectSchema(table); + const expected = z.object({ + timestamp: z.coerce.date(), + int: z.number().min(CONSTANTS.INT32_MIN).max(CONSTANTS.INT32_MAX).int(), + }); + expectSchemaShape(t, expected).from(result); + Expect>(); +}); + /* Disallow unknown keys in table refinement - select */ { const table = singlestoreTable('test', { id: int() }); // @ts-expect-error diff --git a/drizzle-zod/tests/sqlite.test.ts b/drizzle-zod/tests/sqlite.test.ts index bb0f254b5..a5a80fb5f 100644 --- a/drizzle-zod/tests/sqlite.test.ts +++ b/drizzle-zod/tests/sqlite.test.ts @@ -4,7 +4,7 @@ import { test } from 'vitest'; import { z } from 'zod'; import { bufferSchema, jsonSchema } from '~/column.ts'; import { CONSTANTS } from '~/constants.ts'; -import { createInsertSchema, createSelectSchema, createUpdateSchema } from '../src'; +import { createInsertSchema, createSchemaFactory, createSelectSchema, createUpdateSchema } from '../src'; import { Expect, expectSchemaShape } from './utils.ts'; const intSchema = z.number().min(Number.MIN_SAFE_INTEGER).max(Number.MAX_SAFE_INTEGER).int(); @@ -350,6 +350,62 @@ test('all data types', (t) => { Expect>(); }); +test('type coercion - all', (t) => { + const table = sqliteTable('test', ({ + blob, + integer, + text, + }) => ({ + blob: blob({ mode: 'bigint' }).notNull(), + integer1: integer({ mode: 'boolean' }).notNull(), + integer2: integer({ mode: 'timestamp' }).notNull(), + integer3: integer().notNull(), + text: text().notNull(), + })); + + const { createSelectSchema } = createSchemaFactory({ + coerce: { + bigint: true, + boolean: true, + date: true, + number: true, + string: true, + }, + }); + const result = createSelectSchema(table); + const expected = z.object({ + blob: z.coerce.bigint().min(CONSTANTS.INT64_MIN).max(CONSTANTS.INT64_MAX), + integer1: z.coerce.boolean(), + integer2: z.coerce.date(), + integer3: z.coerce.number().min(Number.MIN_SAFE_INTEGER).max(Number.MAX_SAFE_INTEGER).int(), + text: z.coerce.string(), + }); + expectSchemaShape(t, expected).from(result); + Expect>(); +}); + +test('type coercion - mixed', (t) => { + const table = sqliteTable('test', ({ + integer, + }) => ({ + integer1: integer({ mode: 'timestamp' }).notNull(), + integer2: integer().notNull(), + })); + + const { createSelectSchema } = createSchemaFactory({ + coerce: { + date: true, + }, + }); + const result = createSelectSchema(table); + const expected = z.object({ + integer1: z.coerce.date(), + integer2: z.number().min(Number.MIN_SAFE_INTEGER).max(Number.MAX_SAFE_INTEGER).int(), + }); + expectSchemaShape(t, expected).from(result); + Expect>(); +}); + /* Disallow unknown keys in table refinement - select */ { const table = sqliteTable('test', { id: int() }); // @ts-expect-error diff --git a/drizzle-zod/tests/utils.ts b/drizzle-zod/tests/utils.ts index 6a36f66c5..da473b116 100644 --- a/drizzle-zod/tests/utils.ts +++ b/drizzle-zod/tests/utils.ts @@ -9,6 +9,7 @@ export function expectSchemaShape>(t: TaskC for (const key of Object.keys(actual.shape)) { expect(actual.shape[key]!._def.typeName).toStrictEqual(expected.shape[key]?._def.typeName); expect(actual.shape[key]!._def?.checks).toEqual(expected.shape[key]?._def?.checks); + expect(actual.shape[key]!._def?.coerce).toEqual(expected.shape[key]?._def?.coerce); if (actual.shape[key]?._def.typeName === 'ZodOptional') { expect(actual.shape[key]!._def.innerType._def.typeName).toStrictEqual( actual.shape[key]!._def.innerType._def.typeName, From 62016292c82078b22a8284305f92c956f759b2ff Mon Sep 17 00:00:00 2001 From: Mario564 Date: Thu, 26 Dec 2024 10:41:30 -0800 Subject: [PATCH 2/2] Add shorthand syntax to coerce option --- drizzle-zod/src/column.ts | 10 +++++----- drizzle-zod/src/schema.types.ts | 2 +- drizzle-zod/tests/mysql.test.ts | 8 +------- drizzle-zod/tests/pg.test.ts | 8 +------- drizzle-zod/tests/singlestore.test.ts | 8 +------- drizzle-zod/tests/sqlite.test.ts | 8 +------- 6 files changed, 10 insertions(+), 34 deletions(-) diff --git a/drizzle-zod/src/column.ts b/drizzle-zod/src/column.ts index b6c50b85a..e300ff6d1 100644 --- a/drizzle-zod/src/column.ts +++ b/drizzle-zod/src/column.ts @@ -105,9 +105,9 @@ export function columnToSchema(column: Column, factory: CreateSchemaFactoryOptio } else if (column.dataType === 'bigint') { schema = bigintColumnToSchema(column, z, coerce); } else if (column.dataType === 'boolean') { - schema = coerce.boolean ? z.coerce.boolean() : z.boolean(); + schema = coerce === true || coerce.boolean ? z.coerce.boolean() : z.boolean(); } else if (column.dataType === 'date') { - schema = coerce.date ? z.coerce.date() : z.date(); + schema = coerce === true || coerce.date ? z.coerce.date() : z.date(); } else if (column.dataType === 'string') { schema = stringColumnToSchema(column, z, coerce); } else if (column.dataType === 'json') { @@ -230,7 +230,7 @@ function numberColumnToSchema( max = Number.MAX_SAFE_INTEGER; } - let schema = coerce?.number ? z.coerce.number() : z.number(); + let schema = coerce === true || coerce?.number ? z.coerce.number() : z.number(); schema = schema.min(min).max(max); return integer ? schema.int() : schema; } @@ -244,7 +244,7 @@ function bigintColumnToSchema( const min = unsigned ? 0n : CONSTANTS.INT64_MIN; const max = unsigned ? CONSTANTS.INT64_UNSIGNED_MAX : CONSTANTS.INT64_MAX; - const schema = coerce?.bigint ? z.coerce.bigint() : z.bigint(); + const schema = coerce === true || coerce?.bigint ? z.coerce.bigint() : z.bigint(); return schema.min(min).max(max); } @@ -295,7 +295,7 @@ function stringColumnToSchema( max = column.dimensions; } - let schema = coerce?.string ? z.coerce.string() : z.string(); + let schema = coerce === true || coerce?.string ? z.coerce.string() : z.string(); schema = regex ? schema.regex(regex) : schema; return max && fixed ? schema.length(max) : max ? schema.max(max) : schema; } diff --git a/drizzle-zod/src/schema.types.ts b/drizzle-zod/src/schema.types.ts index 4d67d1d6f..9ec093593 100644 --- a/drizzle-zod/src/schema.types.ts +++ b/drizzle-zod/src/schema.types.ts @@ -49,5 +49,5 @@ export interface CreateUpdateSchema { export interface CreateSchemaFactoryOptions { zodInstance?: any; - coerce?: Partial>; + coerce?: Partial> | true; } diff --git a/drizzle-zod/tests/mysql.test.ts b/drizzle-zod/tests/mysql.test.ts index 18f8b3d50..314631b6e 100644 --- a/drizzle-zod/tests/mysql.test.ts +++ b/drizzle-zod/tests/mysql.test.ts @@ -470,13 +470,7 @@ test('type coercion - all', (t) => { })); const { createSelectSchema } = createSchemaFactory({ - coerce: { - bigint: true, - boolean: true, - date: true, - number: true, - string: true, - }, + coerce: true, }); const result = createSelectSchema(table); const expected = z.object({ diff --git a/drizzle-zod/tests/pg.test.ts b/drizzle-zod/tests/pg.test.ts index 68642f073..4f82afc2d 100644 --- a/drizzle-zod/tests/pg.test.ts +++ b/drizzle-zod/tests/pg.test.ts @@ -516,13 +516,7 @@ test('type coercion - all', (t) => { })); const { createSelectSchema } = createSchemaFactory({ - coerce: { - bigint: true, - boolean: true, - date: true, - number: true, - string: true, - }, + coerce: true, }); const result = createSelectSchema(table); const expected = z.object({ diff --git a/drizzle-zod/tests/singlestore.test.ts b/drizzle-zod/tests/singlestore.test.ts index f4824e67e..cf54ea0c0 100644 --- a/drizzle-zod/tests/singlestore.test.ts +++ b/drizzle-zod/tests/singlestore.test.ts @@ -472,13 +472,7 @@ test('type coercion - all', (t) => { })); const { createSelectSchema } = createSchemaFactory({ - coerce: { - bigint: true, - boolean: true, - date: true, - number: true, - string: true, - }, + coerce: true, }); const result = createSelectSchema(table); const expected = z.object({ diff --git a/drizzle-zod/tests/sqlite.test.ts b/drizzle-zod/tests/sqlite.test.ts index a5a80fb5f..5950f6efe 100644 --- a/drizzle-zod/tests/sqlite.test.ts +++ b/drizzle-zod/tests/sqlite.test.ts @@ -364,13 +364,7 @@ test('type coercion - all', (t) => { })); const { createSelectSchema } = createSchemaFactory({ - coerce: { - bigint: true, - boolean: true, - date: true, - number: true, - string: true, - }, + coerce: true, }); const result = createSelectSchema(table); const expected = z.object({