drizzle-effect-schema

drizzle-effect-schema.ts

import * as S from "@effect/schema/Schema";
import { Assume, Column, Equal, Table, getTableColumns, is } from "drizzle-orm";
import {
  MySqlChar,
  MySqlVarBinary,
  MySqlVarChar,
} from "drizzle-orm/mysql-core";
import { PgArray, PgChar, PgUUID, PgVarchar } from "drizzle-orm/pg-core";
import { SQLiteText } from "drizzle-orm/sqlite-core";
import { Simplify } from "effect/Types";

const uuidRegex =
  /^[0-9a-fA-F]{8}\b-[0-9a-fA-F]{4}\b-[0-9a-fA-F]{4}\b-[0-9a-fA-F]{4}\b-[0-9a-fA-F]{12}$/i;
const UUID = S.string.pipe(S.pattern(uuidRegex));

type Columns<TTable extends Table> =
  TTable["_"]["columns"] extends infer TColumns extends Record<
    string,
    Column<any>
  >
    ? TColumns
    : never;

type PropertySignatureFrom<T> =
  T extends S.PropertySignature<infer From, any, any, any> ? From : never;

type PropertySignatureTo<T> =
  T extends S.PropertySignature<any, any, infer To, any> ? To : never;

type RefineArg<TTable extends Table, Col extends keyof Columns<TTable>> =
  | S.Schema<any, any>
  | ((s: {
      [S in keyof InsertColumnPropertySignatures<TTable>]: InsertColumnPropertySignatures<TTable>[S] extends S.PropertySignature<
        any,
        any,
        any,
        any
      >
        ? S.Schema<
            Exclude<
              PropertySignatureFrom<InsertColumnPropertySignatures<TTable>[S]>,
              undefined
            >,
            Exclude<
              PropertySignatureTo<InsertColumnPropertySignatures<TTable>[S]>,
              undefined
            >
          >
        : InsertColumnPropertySignatures<TTable>[S];
    }) => InsertColumnPropertySignatures<TTable>[Col] extends S.PropertySignature<
      any,
      any,
      any,
      any
    >
      ? S.Schema<
          Exclude<
            PropertySignatureFrom<InsertColumnPropertySignatures<TTable>[Col]>,
            undefined
          >,
          any
        >
      : S.Schema<
          Exclude<
            S.Schema.From<InsertColumnPropertySignatures<TTable>[Col]>,
            undefined
          >,
          any
        >);

export type InsertRefine<TTable extends Table> = {
  [K in keyof Columns<TTable>]?: RefineArg<TTable, K>;
};

const literalSchema = S.union(S.string, S.number, S.boolean, S.null);
export const jsonSchema = S.suspend(() =>
  S.union(literalSchema, S.array(jsonSchema), S.record(S.string, jsonSchema)),
);
type Literal = S.Schema.To<typeof literalSchema>;
type Json = Literal | { [key: string]: Json } | Json[];

type GetSchemaForType<TColumn extends Column> =
  TColumn["_"]["dataType"] extends infer TDataType
    ? TDataType extends "custom"
      ? S.Schema<any>
      : TDataType extends "json"
        ? S.Schema<Json>
        : TColumn extends { enumValues: [string, ...string[]] }
          ? Equal<TColumn["enumValues"], [string, ...string[]]> extends true
            ? S.Schema<string>
            : S.Schema<TColumn["enumValues"][number]>
          : TDataType extends "array"
            ? S.Schema<
                Array<
                  GetSchemaForType<
                    Assume<TColumn["_"], { baseColumn: Column }>["baseColumn"]
                  >
                >
              >
            : TDataType extends "bigint"
              ? S.Schema<bigint>
              : TDataType extends "number"
                ? S.Schema<number>
                : TDataType extends "string"
                  ? S.Schema<string>
                  : TDataType extends "boolean"
                    ? S.Schema<boolean>
                    : TDataType extends "date"
                      ? S.Schema<Date>
                      : S.Schema<any>
    : never;

type MapInsertColumnToPropertySignature<TColumn extends Column> =
  TColumn["_"]["notNull"] extends false
    ? S.PropertySignature<
        S.Schema.From<GetSchemaForType<TColumn>> | undefined,
        true,
        S.Schema.To<GetSchemaForType<TColumn>> | undefined,
        true
      >
    : TColumn["_"]["hasDefault"] extends true
      ? S.PropertySignature<
          S.Schema.From<GetSchemaForType<TColumn>> | undefined,
          true,
          S.Schema.To<GetSchemaForType<TColumn>>,
          true
        >
      : GetSchemaForType<TColumn>;

type MapSelectColumnToPropertySignature<TColumn extends Column> =
  TColumn["_"]["notNull"] extends false
    ? S.PropertySignature<
        S.Schema.From<GetSchemaForType<TColumn>> | undefined,
        true,
        S.Schema.To<GetSchemaForType<TColumn>> | undefined,
        true
      >
    : TColumn["_"]["hasDefault"] extends true
      ? S.PropertySignature<
          S.Schema.From<GetSchemaForType<TColumn>>,
          false,
          S.Schema.To<GetSchemaForType<TColumn>>,
          false
        >
      : GetSchemaForType<TColumn>;

export type InsertColumnPropertySignatures<TTable extends Table> = {
  [K in keyof Columns<TTable>]: MapInsertColumnToPropertySignature<
    Columns<TTable>[K]
  >;
};

export type SelectColumnPropertySignatures<TTable extends Table> = {
  [K in keyof Columns<TTable>]: MapSelectColumnToPropertySignature<
    Columns<TTable>[K]
  >;
};

type PropertySignatureReplaceTo<S, With> =
  S extends S.PropertySignature<infer From, infer IO, any, infer OO>
    ? S.PropertySignature<From, IO, With, OO>
    : never;

type PropertySignatureReplaceFrom<S, With> =
  S extends S.PropertySignature<any, infer IO, infer To, infer OO>
    ? S.PropertySignature<With, IO, To, OO>
    : never;

type BuildInsertSchema<
  TTable extends Table,
  TRefine extends InsertRefine<TTable> | {} = {},
> = S.Schema<
  Simplify<
    S.FromStruct<
      {
        [K in Exclude<
          keyof S.FromStruct<InsertColumnPropertySignatures<TTable>>,
          keyof TRefine
        >]: InsertColumnPropertySignatures<TTable>[K];
      } & {
        [K in keyof TRefine]: K extends string
          ? TRefine[K] extends S.Schema<any, any>
            ? TRefine[K]
            : TRefine[K] extends (...a: any[]) => any
              ? InsertColumnPropertySignatures<TTable>[K] extends S.PropertySignature<
                  any,
                  any,
                  any,
                  any
                >
                ? PropertySignatureReplaceFrom<
                    InsertColumnPropertySignatures<TTable>[K],
                    S.Schema.From<ReturnType<TRefine[K]>>
                  >
                : ReturnType<TRefine[K]>
              : never
          : never;
      }
    >
  >,
  Simplify<
    S.ToStruct<
      {
        [K in Exclude<
          keyof InsertColumnPropertySignatures<TTable>,
          keyof TRefine
        >]: InsertColumnPropertySignatures<TTable>[K];
      } & {
        [K in keyof TRefine]: K extends string
          ? TRefine[K] extends S.Schema<any, any>
            ? TRefine[K]
            : TRefine[K] extends (...a: any[]) => any
              ? InsertColumnPropertySignatures<TTable>[K] extends S.PropertySignature<
                  any,
                  any,
                  any,
                  any
                >
                ? PropertySignatureReplaceTo<
                    InsertColumnPropertySignatures<TTable>[K],
                    S.Schema.To<ReturnType<TRefine[K]>>
                  >
                : ReturnType<TRefine[K]>
              : never
          : never;
      }
    >
  >
>;

export function createInsertSchema<
  TTable extends Table,
  TRefine extends InsertRefine<TTable>,
>(table: TTable, refine?: TRefine): BuildInsertSchema<TTable, TRefine> {
  const columns = getTableColumns(table);
  const columnEntries = Object.entries(columns);

  let schemaEntries = Object.fromEntries(
    columnEntries.map(([name, column]) => {
      return [name, mapColumnToSchema(column)];
    }),
  );

  if (refine) {
    schemaEntries = Object.assign(
      schemaEntries,
      Object.fromEntries(
        Object.entries(refine).map(([name, refineColumn]) => {
          return [
            name,
            typeof refineColumn === "function"
              ? refineColumn(schemaEntries as any)
              : refineColumn,
          ];
        }),
      ),
    );
  }

  for (const [name, column] of columnEntries) {
    if (!column.notNull) {
      schemaEntries[name] = S.optional(
        schemaEntries[name]!.pipe(S.nullable),
      ) as any;
    } else if (column.hasDefault) {
      schemaEntries[name] = S.optional(schemaEntries[name]!) as any;
    }
  }

  return S.struct(schemaEntries) as any;
}

export function createSelectSchema<TTable extends Table>(
  table: TTable,
): S.Schema<
  Simplify<S.FromStruct<SelectColumnPropertySignatures<TTable>>>,
  Simplify<S.ToStruct<SelectColumnPropertySignatures<TTable>>>
> {
  const columns = getTableColumns(table);
  const columnEntries = Object.entries(columns);

  const schemaEntries = Object.fromEntries(
    columnEntries.map(([name, column]) => {
      return [name, mapColumnToSchema(column)];
    }),
  );

  for (const [name, column] of columnEntries) {
    if (!column.notNull) {
      schemaEntries[name] = schemaEntries[name]!.pipe(S.nullable);
    }
  }

  return S.struct(schemaEntries) as any;
}

function mapColumnToSchema(column: Column): S.Schema<any, any> {
  let type: S.Schema<any, any> | undefined;

  if (isWithEnum(column)) {
    type = column.enumValues.length
      ? S.literal(...column.enumValues)
      : S.string;
  }

  if (!type) {
    if (is(column, PgUUID)) {
      type = UUID;
    } else if (column.dataType === "custom") {
      type = S.any;
    } else if (column.dataType === "json") {
      type = jsonSchema;
    } else if (column.dataType === "array") {
      type = S.array(
        mapColumnToSchema((column as PgArray<any, any>).baseColumn),
      );
    } else if (column.dataType === "number") {
      type = S.number;
    } else if (column.dataType === "bigint") {
      type = S.bigint;
    } else if (column.dataType === "boolean") {
      type = S.boolean;
    } else if (column.dataType === "date") {
      type = S.Date;
    } else if (column.dataType === "string") {
      let sType = S.string;

      if (
        (is(column, PgChar) ||
          is(column, PgVarchar) ||
          is(column, MySqlVarChar) ||
          is(column, MySqlVarBinary) ||
          is(column, MySqlChar) ||
          is(column, SQLiteText)) &&
        typeof column.length === "number"
      ) {
        sType = sType.pipe(S.maxLength(column.length));
      }

      type = sType;
    }
  }

  if (!type) {
    type = S.any;
  }

  return type;
}

function isWithEnum(
  column: Column,
): column is typeof column & { enumValues: [string, ...string[]] } {
  return (
    "enumValues" in column &&
    Array.isArray(column.enumValues) &&
    column.enumValues.length > 0
  );
}

I would love to see an updated version that works with the latest version of @effect/schema