Disallow unknown keys in drizzle-zod refinement

L-Mario564 · Nov 15, 2024 · 2382a09 · 2382a09
1 parent 9acbabc
commit 2382a09
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 15 deletions.
diff --git a/drizzle-zod/src/create-schema.ts b/drizzle-zod/src/create-schema.ts
@@ -1,4 +1,4 @@
-import { optional, z } from 'zod';
+import { z } from 'zod';
 import { Column, getTableColumns, getViewSelectedFields, is, isTable, isView, SQL } from 'drizzle-orm';
 import { columnToSchema } from './column';
 import { isPgEnum, PgEnum } from 'drizzle-orm/pg-core';

diff --git a/drizzle-zod/src/types.ts b/drizzle-zod/src/types.ts
@@ -1,7 +1,7 @@
 import type { z } from 'zod';
-import type { Assume, Column, Equal, SelectedFieldsFlat, Simplify, Table, View } from 'drizzle-orm';
+import type { Assume, Column, DrizzleTypeError, Equal, SelectedFieldsFlat, Simplify, Table, View } from 'drizzle-orm';
 import type { literalSchema } from './column';
-import { PgEnum } from 'drizzle-orm/pg-core';
+import type { PgEnum } from 'drizzle-orm/pg-core';
 
 type Literal = z.infer<typeof literalSchema>;
 export type Json = Literal | { [key: string]: Json } | Json[];
@@ -175,25 +175,36 @@ export type BuildSchema<
   'strip'
 >;
 
-export interface CreateSelectSchema {
-  <TView extends View>(view: TView): BuildSchema<'select', TView['_']['selectedFields'], undefined>;
-  <
-    TView extends View,
-    TRefine extends BuildRefine<TView['_']['selectedFields']>
-  >(
-    view: TView,
-    refine: TRefine
-  ): BuildSchema<'select', TView['_']['selectedFields'], TRefine>;
+type NoUnknownKeys<
+  TRefinement extends Record<string, any>,
+  TCompare extends Record<string, any>
+> = {
+  [K in keyof TRefinement]: K extends keyof TCompare
+    ? TRefinement[K] extends Record<string, z.ZodTypeAny>
+      ? NoUnknownKeys<TRefinement[K], TCompare[K]>
+      : TRefinement[K]
+    : DrizzleTypeError<`Found unknown key in refinement: "${K & string}"`>;
+};
 
+export interface CreateSelectSchema {
   <TTable extends Table>(table: TTable): BuildSchema<'select', TTable['_']['columns'], undefined>;
   <
     TTable extends Table,
     TRefine extends BuildRefine<TTable['_']['columns']>
   >(
     table: TTable,
-    refine?: TRefine
+    refine?: NoUnknownKeys<TRefine, TTable['$inferSelect']>
   ): BuildSchema<'select', TTable['_']['columns'], TRefine>;
 
+  <TView extends View>(view: TView): BuildSchema<'select', TView['_']['selectedFields'], undefined>;
+  <
+    TView extends View,
+    TRefine extends BuildRefine<TView['_']['selectedFields']>
+  >(
+    view: TView,
+    refine: NoUnknownKeys<TRefine, TView['$inferSelect']>
+  ): BuildSchema<'select', TView['_']['selectedFields'], TRefine>;
+
   <TEnum extends PgEnum<any>>(enum_: TEnum): z.ZodEnum<TEnum['enumValues']>;
 }
 
@@ -204,7 +215,7 @@ export interface CreateInsertSchema {
     TRefine extends BuildRefine<Pick<TTable['_']['columns'], keyof TTable['$inferInsert']>>
   >(
     table: TTable,
-    refine?: TRefine
+    refine?: NoUnknownKeys<TRefine, TTable['$inferInsert']>
   ): BuildSchema<'insert', TTable['_']['columns'], TRefine>;
 }
 

diff --git a/drizzle-zod/tests/pg.test.ts b/drizzle-zod/tests/pg.test.ts
@@ -1,4 +1,4 @@
-import { char, date, getViewConfig, integer, pgEnum, pgMaterializedView, pgTable, pgView, serial, text, timestamp, varchar } from 'drizzle-orm/pg-core';
+import { integer, pgEnum, pgMaterializedView, pgTable, pgView, serial, text, timestamp, varchar } from 'drizzle-orm/pg-core';
 import { test } from 'vitest';
 import { z } from 'zod';
 import { createInsertSchema, createSelectSchema, createUpdateSchema } from '../src';
@@ -294,3 +294,43 @@ test('refine view - select', (t) => {
 	});
 	expectSchemaShape(t, expected).from(result);
 });
+
+/* Disallow unknown keys in table refinement - select */ {
+	const table = pgTable('test', { id: integer() });
+	// @ts-expect-error
+	createSelectSchema(table, { unknown: z.string() });
+}
+
+/* Disallow unknown keys in table refinement - insert */ {
+	const table = pgTable('test', { id: integer() });
+	// @ts-expect-error
+	createInsertSchema(table, { unknown: z.string() });
+}
+
+/* Disallow unknown keys in table refinement - update */ {
+	const table = pgTable('test', { id: integer() });
+	// @ts-expect-error
+	createUpdateSchema(table, { unknown: z.string() });
+}
+
+/* Disallow unknown keys in view qb - select */ {
+	const table = pgTable('test', { id: integer() });
+	const view = pgView('test').as((qb) => qb.select().from(table));
+	const mView = pgMaterializedView('test').as((qb) => qb.select().from(table));
+	const nestedSelect = pgView('test').as((qb) => qb.select({ table }).from(table));
+	// @ts-expect-error
+	createSelectSchema(view, { unknown: z.string() });
+	// @ts-expect-error
+	createSelectSchema(mView, { unknown: z.string() });
+	// @ts-expect-error
+	createSelectSchema(nestedSelect, { table: { unknown: z.string() } });
+}
+
+/* Disallow unknown keys in view columns - select */ {
+	const view = pgView('test', { id: integer() }).as(sql``);
+	const mView = pgView('test', { id: integer() }).as(sql``);
+	// @ts-expect-error
+	createSelectSchema(view, { unknown: z.string() });
+	// @ts-expect-error
+	createSelectSchema(mView, { unknown: z.string() });
+}