Add config option to enable specific model actions

README.md

@@ -171,15 +171,16 @@ model User {
 
 # Additional Options
 
-| Option              | Description                                                                | Type      | Default                       |
-| ------------------- | -------------------------------------------------------------------------- | --------- | ----------------------------- |
-| `output`            | Output directory for the generated routers and zod schemas                 | `string`  | `./generated`                 |
-| `withMiddleware`    | Attaches a global middleware that runs before all procedures               | `boolean` | `true`                        |
-| `withShield`        | Generates a tRPC Shield to use as a permissions layer                      | `boolean` | `true`                        |
-| `contextPath`       | Sets the context path used in your routers                                 | `string`  | `../../../../src/context`     |
-| `trpcOptionsPath`   | Sets the tRPC instance options                                             | `string`  | `../../../../src/trpcOptions` |
-| `isGenerateSelect`  | Enables the generation of Select related schemas and the select property   | `boolean` | `false`                       |
-| `isGenerateInclude` | Enables the generation of Include related schemas and the include property | `boolean` | `false`                       |
+| Option                 | Description                                                                | Type      | Default                                                                                                                                                                      |
+| ---------------------- | -------------------------------------------------------------------------- | --------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `output`               | Output directory for the generated routers and zod schemas                 | `string`  | `./generated`                                                                                                                                                                |
+| `withMiddleware`       | Attaches a global middleware that runs before all procedures               | `boolean` | `true`                                                                                                                                                                       |
+| `withShield`           | Generates a tRPC Shield to use as a permissions layer                      | `boolean` | `true`                                                                                                                                                                       |
+| `contextPath`          | Sets the context path used in your routers                                 | `string`  | `../../../../src/context`                                                                                                                                                    |
+| `trpcOptionsPath`      | Sets the tRPC instance options                                             | `string`  | `../../../../src/trpcOptions`                                                                                                                                                |
+| `isGenerateSelect`     | Enables the generation of Select related schemas and the select property   | `boolean` | `false`                                                                                                                                                                      |
+| `isGenerateInclude`    | Enables the generation of Include related schemas and the include property | `boolean` | `false`                                                                                                                                                                      |
+| `generateModelActions` | Enables the generation of specific model actions                           | `string`  | `aggregate,aggregateRaw,count,create,createMany,delete,deleteMany,findFirst,findFirstOrThrow,findMany,findRaw,findUnique,findUniqueOrThrow,groupBy,update,updateMany,upsert` |
 
 Use additional options in the `schema.prisma`
 
@@ -193,5 +194,6 @@ generator trpc {
   trpcOptionsPath        = "../trpcOptions"
   isGenerateSelect   = true
   isGenerateInclude  = true
+  generateModelActions = "aggregate,aggregateRaw,count,create,createMany,delete,deleteMany,findFirst,findFirstOrThrow,findMany,findRaw,findUnique,findUniqueOrThrow,groupBy,update,updateMany,upsert"
 }
 ```

prisma/schema.prisma

@@ -6,10 +6,10 @@ generator trpc {
   provider          = "node ./lib/generator.js"
   isGenerateSelect  = true
   isGenerateInclude = true
-  withMiddleware = true
-  withShield     = true
-  contextPath    = "./context"
-  trpcOptionsPath = "./trpcOptions"
+  withMiddleware    = true
+  withShield        = true
+  contextPath       = "./context"
+  trpcOptionsPath   = "./trpcOptions"
 }
 
 datasource db {

src/config.ts

@@ -1,14 +1,25 @@
+import { DMMF } from '@prisma/generator-helper';
 import { z } from 'zod';
 
 const configBoolean = z
   .enum(['true', 'false'])
   .transform((arg) => JSON.parse(arg));
 
+const modelActionEnum = z.nativeEnum(DMMF.ModelAction);
+
 export const configSchema = z.object({
   withMiddleware: configBoolean.default('true'),
   withShield: configBoolean.default('true'),
   contextPath: z.string().default('../../../../src/context'),
-  trpcOptionsPath: z.string().optional()
+  trpcOptionsPath: z.string().optional(),
+  generateModelActions: z
+    .string()
+    .default(Object.values(DMMF.ModelAction).join(','))
+    .transform((arg) => {
+      return arg
+        .split(',')
+        .map((action) => modelActionEnum.parse(action.trim()));
+    }),
 });
 
 export type Config = z.infer<typeof configSchema>;

src/helpers.ts

@@ -168,42 +168,33 @@ export function generateProcedure(
 
 export function generateRouterSchemaImports(
   sourceFile: SourceFile,
-  name: string,
-  hasCreateMany: boolean,
-  provider: string,
+  modelName: string,
+  modelActions: string[],
 ) {
-  let statements = [
-    `import { ${name}FindUniqueSchema } from "../schemas/findUnique${name}.schema";`,
-    `import { ${name}FindFirstSchema } from "../schemas/findFirst${name}.schema";`,
-    `import { ${name}FindManySchema } from "../schemas/findMany${name}.schema";`,
-    `import { ${name}CreateOneSchema } from "../schemas/createOne${name}.schema";`,
-  ];
-
-  if (hasCreateMany) {
-    statements.push(
-      `import { ${name}CreateManySchema } from "../schemas/createMany${name}.schema";`,
-    );
-  }
-
-  statements = statements.concat([
-    `import { ${name}DeleteOneSchema } from "../schemas/deleteOne${name}.schema";`,
-    `import { ${name}UpdateOneSchema } from "../schemas/updateOne${name}.schema";`,
-    `import { ${name}DeleteManySchema } from "../schemas/deleteMany${name}.schema";`,
-    `import { ${name}UpdateManySchema } from "../schemas/updateMany${name}.schema";`,
-    `import { ${name}UpsertSchema } from "../schemas/upsertOne${name}.schema";`,
-    `import { ${name}AggregateSchema } from "../schemas/aggregate${name}.schema";`,
-    `import { ${name}GroupBySchema } from "../schemas/groupBy${name}.schema";`,
-  ]);
+  sourceFile.addStatements(
+    /* ts */
+    [
+      // remove any duplicate import statements
+      ...new Set(
+        modelActions.map((opName) =>
+          getRouterSchemaImportByOpName(opName, modelName),
+        ),
+      ),
+    ].join('\n'),
+  );
+}
 
-  if (provider === 'mongodb') {
-    statements = statements.concat([
-      `import { ${name}FindRawObjectSchema } from "../schemas/objects/${name}FindRaw.schema";`,
-      `import { ${name}AggregateRawObjectSchema } from "../schemas/objects/${name}AggregateRaw.schema";`,
-    ]);
-  }
+export const getRouterSchemaImportByOpName = (
+  opName: string,
+  modelName: string,
+) => {
+  const opType = opName.replace('OrThrow', '');
+  const inputType = getInputTypeByOpName(opType, modelName);
 
-  sourceFile.addStatements(/* ts */ statements.join('\n'));
-}
+  return inputType
+    ? `import { ${inputType} } from "../schemas/${opType}${modelName}.schema";`
+    : '';
+};
 
 export const getInputTypeByOpName = (opName: string, modelName: string) => {
   let inputType;

src/prisma-generator.ts

@@ -1,4 +1,4 @@
-import { EnvValue, GeneratorOptions } from '@prisma/generator-helper';
+import { DMMF, EnvValue, GeneratorOptions } from '@prisma/generator-helper';
 import { getDMMF, parseEnvValue } from '@prisma/internals';
 import { promises as fs } from 'fs';
 import path from 'path';
@@ -15,7 +15,7 @@ import {
   generateShieldImport,
   generatetRPCImport,
   getInputTypeByOpName,
-  resolveModelsComments
+  resolveModelsComments,
 } from './helpers';
 import { project } from './project';
 import removeDir from './utils/removeDir';
@@ -53,8 +53,6 @@ export async function generate(options: GeneratorOptions) {
     (it) => parseEnvValue(it.provider) === 'prisma-client-js',
   );
 
-  const dataSource = options.datasources?.[0];
-
   const prismaClientDmmf = await getDMMF({
     datamodel: options.datamodel,
     previewFeatures: prismaClientProvider.previewFeatures,
@@ -96,8 +94,17 @@ export async function generate(options: GeneratorOptions) {
   for (const modelOperation of modelOperations) {
     const { model, ...operations } = modelOperation;
     if (hiddenModels.includes(model)) continue;
+
+    const modelActions = Object.keys(operations).filter<DMMF.ModelAction>(
+      (opType): opType is DMMF.ModelAction =>
+        config.generateModelActions.includes(
+          opType.replace('One', '') as DMMF.ModelAction,
+        ),
+    );
+    if (!modelActions.length) continue;
+
     const plural = pluralize(model.toLowerCase());
-    const hasCreateMany = Boolean(operations.createMany);
+
     generateRouterImport(appRouter, plural, model);
     const modelRouter = project.createSourceFile(
       path.resolve(outputDir, 'routers', `${model}.router.ts`),
@@ -110,16 +117,13 @@ export async function generate(options: GeneratorOptions) {
       config,
     });
 
-    generateRouterSchemaImports(
-      modelRouter,
-      model,
-      hasCreateMany,
-      dataSource.provider,
-    );
+    generateRouterSchemaImports(modelRouter, model, modelActions);
 
     modelRouter.addStatements(/* ts */ `
       export const ${plural}Router = t.router({`);
-    for (const [opType, opNameWithModel] of Object.entries(operations)) {
+
+    for (const opType of modelActions) {
+      const opNameWithModel = operations[opType];
       const baseOpType = opType.replace('OrThrow', '');
 
       generateProcedure(