Mongo DataSource + Zod + DataLoader

user.ts

import { AsyncLocalStorage } from 'async_hooks'

import DataLoader from 'dataloader'
import { ObjectId } from 'mongodb'
import { z } from 'zod'

import {
  Connection,
  ConnectionArgs,
  connectionArgsSchema,
} from './common-place-to-put-connection-args'
import { db } from './where-db-is-initialized'

const loadersStorage = new AsyncLocalStorage<
  Record<string, DataLoader<ObjectId | string, unknown>>
>()

const getLoaders = () => {
  const store = loadersStorage.getStore()
  if (!store) throw new Error('Loaders are not initialized')
  return store
}

export const createDataLoaders = () => {
  const userByIdLoader = new DataLoader<ObjectId, User | null>(async (ids) => {
    const users = await collection
      .find({ _id: { $in: ids } })
      .toArray()
      .then((users) => users.map((user) => userSchema.parse(user)))

    return ids.map((id) => {
      const user = users.find((user) => user._id.equals(id))
      if (!user) return null

      userByEmailLoader.prime(user.email, user)
      return user
    })
  })

  const userByEmailLoader = new DataLoader<string, User | null>(
    async (emails) => {
      const users = await collection
        .find({ email: { $in: emails } })
        .toArray()
        .then((users) => users.map((user) => userSchema.parse(user)))

      return emails.map((email) => {
        const user = users.find((user) => user.email === email)
        if (!user) return null

        userByIdLoader.prime(user._id, user)
        return user
      })
    },
  )
  return { userByIdLoader, userByEmailLoader }
}

const userSchema = z.object({
  _id: z.instanceof(ObjectId),
  name: z.string(),
  email: z.string().email(),
})
type UserInput = z.input<typeof userSchema>
type User = z.infer<typeof userSchema>

const collection = db.collection<User>('users')

const create = async (user: UserInput | unknown): Promise<User> => {
  const data = userSchema.parse(user)
  const result = await collection.insertOne(data)

  const { userByIdLoader, userByEmailLoader } = getLoaders()
  userByIdLoader
    .clear(result.ops[0]._id)
    .prime(result.ops[0]._id, result.ops[0])
  userByEmailLoader
    .clear(result.ops[0].email)
    .prime(result.ops[0].email, result.ops[0])

  return result.ops[0]
}

const update = async (
  _id: ObjectId,
  user: Partial<UserInput> | unknown,
): Promise<User | null> => {
  const partialUserSchema = userSchema.partial()
  const data = partialUserSchema.parse(user)

  const result = await collection.findOneAndUpdate(
    { _id },
    { $set: data },
    { returnDocument: 'after' },
  )
  if (!result.value) return null

  const updatedUser = userSchema.parse(result.value)

  const { userByIdLoader, userByEmailLoader } = getLoaders()
  userByIdLoader.clear(updatedUser._id).prime(updatedUser._id, updatedUser)
  userByEmailLoader
    .clear(updatedUser.email)
    .prime(updatedUser.email, updatedUser)

  return updatedUser
}

const remove = async (_id: ObjectId): Promise<boolean> => {
  const result = await collection.findOneAndUpdate(
    { _id },
    { $set: { removedAt: new Date() } },
  )

  if (result.value) {
    const { userByIdLoader, userByEmailLoader } = getLoaders()
    userByIdLoader.clear(result.value._id)
    userByEmailLoader.clear(result.value.email)
  }

  return !!result.value
}

const hardDelete = async (_id: ObjectId): Promise<boolean> => {
  const result = await collection.deleteOne({ _id })

  if (result.deletedCount) {
    const { userByIdLoader, userByEmailLoader } = getLoaders()
    userByIdLoader.clear(_id)
    userByEmailLoader.clear(_id)
  }

  return !!result.deletedCount
}

const findById = async (_id: ObjectId): Promise<User | null> => {
  const { userByIdLoader } = getLoaders()
  return userByIdLoader.load(_id)
}

const findByEmail = async (email: string): Promise<User | null> => {
  const { userByEmailLoader } = getLoaders()
  return userByEmailLoader.load(email)
}

const list = async (
  args: ConnectionArgs | unknown,
): Promise<Connection<User>> => {
  const { before, after, first, last } = connectionArgsSchema.parse(args)
  const query = {}
  const sort = { _id: 1 }

  if (before) {
    query['_id'] = { $lt: new ObjectId(before) }
  }

  if (after) {
    query['_id'] = { $gt: new ObjectId(after) }
    sort._id = -1
  }

  const limit = first || last || 10
  const [users, totalCount] = await Promise.all([
    collection
      .find(query)
      .sort(sort)
      .limit(limit + 1)
      .toArray(),
    collection.countDocuments(query),
  ])

  const hasNextPage = users.length > limit
  const hasPreviousPage = !!before

  const { userByIdLoader, userByEmailLoader } = getLoaders()

  const edges = users.slice(0, limit).map((user) => {
    const node = userSchema.parse(user)
    userByIdLoader.prime(node._id, node)
    userByEmailLoader.prime(node.email, node)

    return {
      cursor: user._id.toHexString(),
      node,
    }
  })

  return {
    edges,
    pageInfo: {
      hasNextPage,
      hasPreviousPage,
      startCursor: edges[0]?.cursor,
      endCursor: edges.at(-1)?.cursor,
    },
    totalCount,
  }
}

export const UserDataSource = {
  create,
  update,
  remove,
  hardDelete,
  findById,
  findByEmail,
  list,
}