import { isEmpty } from "lodash";
import { Prisma } from "prisma/generated";

/**
 * Middleware featured below
 */

// This naming convention specifically matches the
// safe word to the exported constant to support the
// middleware error handling
export const safeWordDelete = "FORCE_DELETE_ALL";
export const FORCE_DELETE_ALL = { id: safeWordDelete };

export const safeWordFind = "FIND_ALL";
export const FIND_ALL = { id: safeWordFind };

export const safeWordUpdate = "UPDATE_ALL";
export const UPDATE_ALL = { id: safeWordUpdate };

/**
 * If you have found yourself here after a hard lesson:
 *
 * These middleware specifically do NOT guard in the case of a nested find, update, delete.
 * These middleware also don't protect in the case of an undefined variable being passed in alongside a specific condition.
 *
 * Ex: {
 *    id, // undefined
 *    otherAttr: null
 * }
 *
 * The above where clause will delete/modify/find all regardless of undefined status of variable appId.
 *
 * The !! best !! guard here is proper type adherence - do not bang, do not coerce.
 */
const _guardEnMasse = ({
  params,
  actions,
  safeWord,
}: {
  params: Prisma.MiddlewareParams;
  actions: Prisma.MiddlewareParams["action"][];
  safeWord: string;
}) => {
  // Check if empty, if not and safeWord is provided, then set where = {} and proceed to execute
  if (actions.includes(params.action)) {
    // Don't allow delete all queries
    if (isEmpty(params.args.where)) {
      throw new Error(
        `It looks like you just tried to perform a ${params.action} on all of the ${params.model}s. If this was intentional, pass 'where: ${safeWord}'`,
      );
    }
  }

  if (params.args?.where?.id === safeWord) {
    params.args.where = {};
  }
};

export const guardEnMasse: PrismaMiddleware = async (params, next) => {
  /* DELETION PROTECTION MIDDLEWARE */
  _guardEnMasse({
    actions: ["delete", "deleteMany"],
    params,
    safeWord: safeWordDelete,
  });

  /* MASS FIND PROTECTION MIDDLEWARE */
  _guardEnMasse({
    actions: ["findMany"],
    params,
    safeWord: safeWordFind,
  });

  /* MASS UPDATE PROTECTION MIDDLEWARE */
  _guardEnMasse({
    actions: ["updateMany"],
    params,
    safeWord: safeWordUpdate,
  });

  return next(params);
};

const dbClient = new PrismaClient();

dbClient.$use(guardEnMasse);

const whoops = undefined;

// Throws an error
dbClient.sensitiveData.deleteMany({ where: { id: whoops } });

// Deletes everything
dbClient.insensitiveData.deleteMany({ where: FORCE_DELETE_ALL });

#!/usr/bin/env node

import { generatorHandler } from "@prisma/generator-helper";
import { logger } from "@prisma/internals";

import { convert } from "./convertor";

generatorHandler({
  onManifest: () => ({
    prettyName: "where-required",
    requiresGenerators: ["prisma-client-js"],
    defaultOutput: ".",
  }),
  onGenerate: async (options) => {
    const nodeModulePath = options.generator.config.nodeModulePath;
    if (!nodeModulePath) {
      logger.error("No nodeModulePath config.");
      return;
    }

    const debug = options.generator.config.debug === "true";

    const models = options.dmmf.datamodel.models;

    convert({
      path: `${nodeModulePath}/.prisma/client/index.d.ts`,
      models: models.map((model) => model.name),
      debug,
    });
  },
});

import { Project, SourceFile, ts } from "ts-morph";
import isPropertySignature = ts.isPropertySignature;

function removeQuestionTokenFromWhereFieldProperty(
  sourceFile: SourceFile,
  models: string[]
) {
  const prismaNamespace = sourceFile.getModuleOrThrow("Prisma");
  let isInTypeAlias = false;
  prismaNamespace.transform((tr) => {
    const currentNode = tr.currentNode;
    if (
      ts.isTypeAliasDeclaration(currentNode) &&
      models.some(
        (targetModel) =>
          currentNode.name.getText().startsWith(targetModel) &&
          !currentNode.name.getText().startsWith(`${targetModel}$`) &&
          (currentNode.name.getText().endsWith("FindManyArgs") ||
            currentNode.name.getText().endsWith("DeleteManyArgs") ||
            currentNode.name.getText().endsWith("DeleteArgs") ||
            currentNode.name.getText().endsWith("UpdateManyArgs"))
      )
    ) {
      isInTypeAlias = true;
      const newNode = tr.visitChildren();
      isInTypeAlias = false;
      return newNode;
    }
    if (
      isInTypeAlias &&
      isPropertySignature(currentNode) &&
      currentNode.name.getText() === "where"
    ) {
      isInTypeAlias = false;
      return tr.factory.updatePropertySignature(
        currentNode,
        currentNode.modifiers,
        currentNode.name,
        undefined,
        currentNode.type
      );
    }
    return tr.visitChildren();
  });
}

function removeQuestionTokenFromActionFunctionProperty(
  sourceFile: SourceFile,
  models: string[]
) {
  const prismaNamespace = sourceFile.getModuleOrThrow("Prisma");

  let isInDelegate = false;
  let isInFindMany = false;
  prismaNamespace.transform((tr) => {
    const currentNode = tr.currentNode;
    if (
      ts.isInterfaceDeclaration(currentNode) &&
      models.some(
        (targetModel) => currentNode.name.getText() === `${targetModel}Delegate`
      )
    ) {
      isInDelegate = true;
      const newNode = tr.visitChildren();
      isInDelegate = false;
      return newNode;
    }
    if (
      isInDelegate &&
      ts.isMethodSignature(currentNode) &&
      (currentNode.name.getText() === "findMany" ||
        currentNode.name.getText() === "deleteMany" ||
        currentNode.name.getText() === "delete" ||
        currentNode.name.getText() === "updateMany")
    ) {
      isInFindMany = true;
      const newNode = tr.visitChildren();
      isInFindMany = false;
      return newNode;
    }
    if (isInFindMany && ts.isParameter(currentNode)) {
      isInDelegate = false;
      isInFindMany = false;
      return tr.factory.updateParameterDeclaration(
        currentNode,
        currentNode.modifiers,
        currentNode.dotDotDotToken,
        currentNode.name,
        undefined,
        currentNode.type,
        currentNode.initializer
      );
    }
    return tr.visitChildren();
  });
}

export function convert(arguments_: {
  path: string;
  models: string[];
  debug: boolean;
}) {
  const project = new Project();
  const sourceFile = project.addSourceFileAtPath(arguments_.path);

  removeQuestionTokenFromWhereFieldProperty(sourceFile, arguments_.models);
  removeQuestionTokenFromActionFunctionProperty(sourceFile, arguments_.models);

  sourceFile.saveSync();
}

// Compile error (missing an arg)
dbClient.sensitiveData.deleteMany()

// Compile error (missing a where)
dbClient.sensitiveData.findMany({select: {...}})

// Runtime error (as above)
dbClient.sensitiveData.findMany({ where: {id: undefined}})

// Success via runtime override with maximal clarity on what's happening
dbClient.sensitiveData.findMany({ where: FIND_ALL })