Merge pull request #68 from lottamus/feat/optional-model-actions

Add config option to enable specific model actions
omar-dulaimi · Feb 24, 2023 · 89710fb · 89710fb
2 parents 54c0d0d + cd64efa
commit 89710fb
Show file tree

Hide file tree

Showing 5 changed files with 77 additions and 68 deletions.
diff --git a/README.md b/README.md
@@ -171,29 +171,32 @@ 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`                       |
-| `showModelNameInProcedure`  | When disabled, the generated procedure no longer includes the name of the Prisma model | `boolean` | `true`                        |
+| 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`                                                                                                                                                                      |
+| `showModelNameInProcedure` | When disabled, the generated procedure no longer includes the name of the Prisma model | `boolean` | `true`                                                                                                                                                                       |
+| `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`
 
 ```prisma
 generator trpc {
-  provider                  = "prisma-trpc-generator"
-  output                    = "./trpc"
-  withMiddleware            = false
-  withShield                = false
-  contextPath               = "../context"
-  trpcOptionsPath           = "../trpcOptions"
-  isGenerateSelect          = true
-  isGenerateInclude         = true
+  provider           = "prisma-trpc-generator"
+  output             = "./trpc"
+  withMiddleware     = false
+  withShield         = false
+  contextPath        = "../context"
+  trpcOptionsPath    = "../trpcOptions"
+  isGenerateSelect   = true
+  isGenerateInclude  = true
   showModelNameInProcedure  = false
+  generateModelActions = "aggregate,aggregateRaw,count,create,createMany,delete,deleteMany,findFirst,findFirstOrThrow,findMany,findRaw,findUnique,findUniqueOrThrow,groupBy,update,updateMany,upsert"
 }
 ```
diff --git a/prisma/schema.prisma b/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 {

diff --git a/src/config.ts b/src/config.ts
@@ -1,15 +1,26 @@
+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(),
-  showModelNameInProcedure: configBoolean.default('true')
+  showModelNameInProcedure: configBoolean.default('true'),
+  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>;
diff --git a/src/helpers.ts b/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;

diff --git a/src/prisma-generator.ts b/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(