Support views in drizzle-zod

drizzle-zod/src/create-schema.ts

@@ -1,51 +1,41 @@
 import { z } from 'zod';
-import { getTableColumns } from 'drizzle-orm';
+import { Column, getTableColumns, getViewSelectedFields, is, isTable, SQL } from 'drizzle-orm';
 import { columnToSchema } from './column';
 import type { Table, View } from 'drizzle-orm';
-import type { BuildTableSelectSchema, BuildViewSelectSchema } from './types';
+import type { BuildSelectSchema } from './types';
 
-// export function createSelectSchema<
-// 	TView extends View
-// >(
-// 	view: TView
-// ): z.ZodObject<BuildViewSelectSchema<TView>>
-// export function createSelectSchema<
-// 	TTable extends Table
-// >(
-// 	table: TTable
-// ): z.ZodObject<BuildTableSelectSchema<TTable>>;
-// export function createSelectSchema<
-// 	TTableLike extends Table | View
-// >(
-// 	tableLike: TTableLike
-// ) {
-//   const columns = isTable(tableLike) ? getTableColumns(tableLike) : getViewSelectedFields(tableLike);
-//   let columnSchemas: Record<string, z.ZodTypeAny> = {};
+function createSelectColumns(columns: Record<string, any>): z.ZodTypeAny {
+  const columnSchemas: Record<string, z.ZodTypeAny> = {};
 
-//   for (const [key, column] of Object.entries(columns)) {
-//     if (is(column, Column)) {
-//       const schema = columnToSchema(column);
-//       columnSchemas[key] = column.notNull ? schema : schema.nullable();
-//     } else if (is(column, SQL)) {
-//       column.
-//     }
-//   }
+  for (const [key, selected] of Object.entries(columns)) {
+    if (!is(selected, Column) && !is(selected, SQL) && !is(selected, SQL.Aliased) && typeof selected === 'object') {
+      columnSchemas[key] = createSelectColumns(selected);
+      continue;
+    }
 
-//   return z.object(columnSchemas) as any;
-// }
+    const column = is(selected, Column) ? selected : undefined;
+    const schema = !!column ? columnToSchema(column) : z.any();
+    columnSchemas[key] = (!!column && !column.notNull) ? schema.nullable() : schema;
+  }
+
+  return z.object(columnSchemas) as any;
+}
 
+export function createSelectSchema<
+	TView extends View
+>(
+	view: TView
+): BuildSelectSchema<TView['_']['selectedFields']>;
 export function createSelectSchema<
 	TTable extends Table
 >(
 	table: TTable
-): z.ZodObject<BuildTableSelectSchema<TTable>> {
-  const columns = getTableColumns(table);
-  let columnSchemas: Record<string, z.ZodTypeAny> = {};
-
-  for (const [key, column] of Object.entries(columns)) {
-    const schema = columnToSchema(column);
-    columnSchemas[key] = column.notNull ? schema : schema.nullable();
-  }
-
-  return z.object(columnSchemas) as any;
+): BuildSelectSchema<TTable['_']['columns']>;
+export function createSelectSchema<
+	TTableLike extends Table | View
+>(
+	tableLike: TTableLike
+) {
+  const columns = isTable(tableLike) ? getTableColumns(tableLike) : getViewSelectedFields(tableLike);
+  return createSelectColumns(columns);
 }

drizzle-zod/src/types.ts

@@ -1,16 +1,25 @@
 import { z } from 'zod';
-import type { Column, Equal, Simplify, Table, View } from 'drizzle-orm';
+import type { Assume, Column, Equal, SelectedFieldsFlat, Simplify } from 'drizzle-orm';
 import type { literalSchema } from './column';
 
 type Literal = z.infer<typeof literalSchema>;
 export type Json = Literal | { [key: string]: Json } | Json[];
 
+type EnumHasAtLeastOneValue<TEnum extends [string, ...string[]] | undefined> =
+  TEnum extends [infer TString, ...string[]]
+    ? TString extends `${infer TLiteral}`
+      ? TLiteral extends string
+        ? true
+        : false
+      : false
+  : false;
+
 export type GetZodType<
   TData,
   TDataType extends string,
   TEnumValues extends [string, ...string[]] | undefined,
-> = TEnumValues extends [string, ...string[]]
-  ? z.ZodEnum<TEnumValues>
+> = EnumHasAtLeastOneValue<TEnumValues> extends true
+  ? z.ZodEnum<Assume<TEnumValues, [string, ...string[]]>>
   : TData extends infer TTuple extends [any, ...any[]]
   ? z.ZodTuple<{ [K in keyof TTuple]: GetZodType<TTuple[K], never, never> }>
   : TData extends infer TDict extends Record<string, any>
@@ -31,35 +40,24 @@ export type GetZodType<
   ? z.ZodType<Buffer>
   : z.ZodTypeAny;
 
-
-export type BuildTableSelectSchema<
-  TTable extends Table
-> = Simplify<{
-  [K in keyof TTable['_']['columns']]: TTable['_']['columns'][K] extends infer TColumn extends Column
-    ? 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>
-      : never
-    : never
-}>;
-
-export type BuildViewSelectSchema<
-  TTable extends View
-> = {
-  [K in keyof TTable['_']['selectedFields']]: TTable['_']['selectedFields'][K] extends infer TColumn extends Column
-    ? GetZodType<
-        TColumn['_']['data'],
-        TColumn['_']['dataType'],
-        TColumn['_'] extends { enumValues: [string, ...string[]] } ? TColumn['_']['enumValues'] : never
-      > extends infer TSchema extends z.ZodTypeAny
-      ? TColumn['_']['notNull'] extends true
-        ? TSchema
-        : z.ZodNullable<TSchema>
-      : never
-    : never
-};
+export type BuildSelectSchema<
+  TColumns extends Record<string, any>
+> = z.ZodObject<
+  Simplify<{
+    [K in keyof TColumns]:
+      TColumns[K] extends infer TColumn extends Column
+        ? 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
+      : TColumns[K] extends infer TObject extends SelectedFieldsFlat<Column>
+        ? BuildSelectSchema<TObject>
+        : z.ZodAny
+  }>,
+  'strip'
+>;