I’m creating an extension that allows me to modify the output of a field based on a user’s permissions.

The extension looks like this:

import { isArray } from 'lodash-es';

import { PrismaModels } from '@dapi/utils/prisma';
import { Prisma } from '@prisma/client';

type PrismaMapping = {
  [ModelName in keyof PrismaModels]?: (
    userId: string,
    model: PrismaModels[ModelName],
  ) => PrismaModels[ModelName] | Promise<PrismaModels[ModelName]>;
};

type UnionToIntersection<Union> = (
  Union extends unknown ? (union: Union) => void : never
) extends (intersection: infer Intersection) => void
  ? Intersection
  : never;

type Query = UnionToIntersection<PrismaModels[keyof PrismaModels]> | Array<UnionToIntersection<PrismaModels[keyof PrismaModels]>>

/**
 * Handle a Prisma query
 *
 * @param {any} query The Prisma query
 * @param {any} args The Prisma query arguments
 * @param {keyof PrismaModels} model The Prisma model
 * @param {string} userId The user ID
 * @param {PrismaMapping} models The Prisma modeli mapping
 * @returns {Promise<Query>} The Prisma query result
 */
async function handleQuery(
  query: any,
  args: any,
  model: keyof PrismaModels,
  userId: string,
  models: PrismaMapping,
) {
  // TODO?: Do we need to type assert the query?
  const res = await query(args) as Query;

  if (res) {
    if (isArray(res)) {
      return Promise.all(
        res.map(async (r) => {
          return (await models[model]?.(userId, r)) ?? r;
        }),
      );
    } else {
      return (await models[model]?.(userId, res)) ?? res;
    }
  } else {
    return res;
  }
}

/**
 * Limit fields on a Prisma query
 *
 * @param {string} userId The user ID
 * @param {PrismaMapping} models The Prisma model mapping
 * @returns {Prisma.ExtensionDeclaration} The Prisma extension declaration
 */
export function limitFields(userId: string, models: PrismaMapping) {
  return Prisma.defineExtension((prisma) => {
    return prisma.$extends({
      query: {
        $allModels: {
          async findFirst({ query, args, model }) {
            return await handleQuery(query, args, model, userId, models);
          },
          async findFirstOrThrow({ query, args, model }) {
            return await handleQuery(query, args, model, userId, models);
          },
          async findMany({ query, args, model }) {
            return await handleQuery(query, args, model, userId, models);
          },
          async findUnique({ query, args, model }) {
            return await handleQuery(query, args, model, userId, models);
          },
          async findUniqueOrThrow({ query, args, model }) {
            return await handleQuery(query, args, model, userId, models);
          },
          async create({ query, args, model }) {
            return await handleQuery(query, args, model, userId, models);
          },
          async update({ query, args, model }) {
            return await handleQuery(query, args, model, userId, models);
          },
          async upsert({ query, args, model }) {
            return await handleQuery(query, args, model, userId, models);
          },
          async delete({ query, args, model }) {
            return await handleQuery(query, args, model, userId, models);
          },
        },
      },
    });
  });
}

The extension works, but I’m struggling to find the right types for the query and args parameters for the handleQuery function.

I’ve tried using a few generics since query and args will vary depending on the type of Prisma query being used. I’m just trying to get better type safety.