web/packages/trpc/lib/utils/discriminatedUnion.ts

import { z } from "zod"

import type { ZodDiscriminatedUnionOption, ZodError } from "zod"

export interface DiscriminatedUnionError {
  error: ZodError<unknown>
}

export interface Option extends ZodDiscriminatedUnionOption<"__typename"> {}

/**
 * This file is created to handle our discriminated unions
 * validations primarily for union returns from Contentstacks
 * GraphQL server.
 *
 * In the case of a new block being added to the union in Contenstack,
 * Zod will throw because that typename is not expected ("invalid_union_discriminator").
 * Therefore we add a safety that we return everytime we stumble upon
 * the issue and then we filter it out in the transform.
 *
 * This replaces the `cleanEmptyObjects` function that would require
 * everyone to never make a mistake in adding __typename to root
 * anywhere (or any other potentially global fields in case the return type
 * is an Interface e.g).
 */

export function discriminatedUnion<R>(options: Option[]) {
  return z
    .discriminatedUnion("__typename", [
      z.object({ __typename: z.literal(undefined) }),
      ...options,
    ])
    .catch(({ error }: DiscriminatedUnionError) => {
      if (
        error.issues.find(
          (issue) => issue.code === "invalid_union_discriminator"
        )
      ) {
        return { __typename: undefined }
      }
      throw new Error(error.message)
    })
    .transform((data) => {
      if (data.__typename === "undefined" || data.__typename === undefined) {
        return null
      }
      return data as R
    })
}

export function discriminatedUnionArray<T extends Option>(options: T[]) {
  return z
    .array(discriminatedUnion(options))
    .transform((blocks) =>
      blocks.filter((block) => !!block)
    ) as unknown as z.ZodEffects<z.ZodArray<T>>
}