From 1686686d5111648a301680bae7f26d605aabf3de Mon Sep 17 00:00:00 2001
From: Mario564 <ka.mario564@gmail.com>
Date: Tue, 10 Dec 2024 09:43:17 -0800
Subject: [PATCH] Fix type refinement type mapping in drizzle-zod

---
 drizzle-zod/src/schema.types.internal.ts | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/drizzle-zod/src/schema.types.internal.ts b/drizzle-zod/src/schema.types.internal.ts
index 5732e2e0f..8b89187f2 100644
--- a/drizzle-zod/src/schema.types.internal.ts
+++ b/drizzle-zod/src/schema.types.internal.ts
@@ -39,13 +39,20 @@ export type BuildRefine<
 	: never;
 
 type HandleRefinement<
+	TType extends 'select' | 'insert' | 'update',
 	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 extends (schema: any) => z.ZodTypeAny ? (TColumn['_']['notNull'] extends true ? ReturnType<TRefinement>
+		: z.ZodNullable<ReturnType<TRefinement>>) extends infer TSchema
+		? TType extends 'update' ? z.ZodOptional<Assume<TSchema, z.ZodTypeAny>> : TSchema
+	: z.ZodTypeAny
 	: TRefinement;
 
+type IsRefinementDefined<TRefinements, TKey extends string> = TKey extends keyof TRefinements
+	? TRefinements[TKey] extends z.ZodTypeAny | ((schema: any) => any) ? true
+	: false
+	: false;
+
 export type BuildSchema<
 	TType extends 'select' | 'insert' | 'update',
 	TColumns extends Record<string, any>,
@@ -56,9 +63,8 @@ export type BuildSchema<
 			{
 				[K in keyof TColumns]: TColumns[K] extends infer TColumn extends Column
 					? TRefinements extends object
-						? TRefinements[Assume<K, keyof TRefinements>] extends
-							infer TRefinement extends z.ZodTypeAny | ((schema: z.ZodTypeAny) => z.ZodTypeAny)
-							? HandleRefinement<TRefinement, TColumn>
+						? IsRefinementDefined<TRefinements, Assume<K, string>> extends true
+							? HandleRefinement<TType, TRefinements[Assume<K, keyof TRefinements>], TColumn>
 						: HandleColumn<TType, TColumn>
 					: HandleColumn<TType, TColumn>
 					: TColumns[K] extends infer TObject extends SelectedFieldsFlat<Column> | Table | View ? BuildSchema<