Prisma extensions

findInBatches.ts

import { Prisma } from '@prisma/client'

interface Omit {
  <T extends object, K extends [...(keyof T)[]]>(obj: T, ...keys: K): {
    [K2 in Exclude<keyof T, K[number]>]: T[K2]
  }
}

const omit: Omit = (obj, ...keys) => {
  const ret = {} as {
    [K in keyof typeof obj]: (typeof obj)[K]
  }
  let key: keyof typeof obj
  for (key in obj) {
    if (!keys.includes(key)) {
      ret[key] = obj[key]
    }
  }
  return ret
}

async function findInBatches<T, A, TResult extends Prisma.Result<T, A, 'findMany'>>(
  this: T,
  args: Prisma.Exact<A, Prisma.Args<T, 'findMany'>> & {
    size: number
  },
  callback: (results: TResult, page: number) => Promise<void>
) {
  const { size = 500 } = args

  let count = 0
  let pages = 0

  while (true) {
    const results = await (this as any).findMany({
      ...omit(args, 'size'),
      skip: size * pages,
    })

    if (results.length === 0) {
      break
    }

    await callback(results, pages)

    count += args.size
    pages++
  }

  return { count }
}

const extension = Prisma.defineExtension({
  name: 'findInBatches',
  model: {
    $allModels: {
      findInBatches,
    },
  },
})

export default extension