From 0926a020f4d3019bd3d391ab073aa42ae1a1325f Mon Sep 17 00:00:00 2001
From: ymc9 <104139426+ymc9@users.noreply.github.com>
Date: Sun, 3 Sep 2023 22:10:41 +0800
Subject: [PATCH] fix: add the missing "count" schema/router for zod/trpc

---
 packages/plugins/trpc/src/generator.ts         |  4 ++++
 packages/plugins/trpc/src/helpers.ts           |  4 ++++
 packages/schema/src/plugins/zod/transformer.ts | 14 ++++++++++++++
 3 files changed, 22 insertions(+)

diff --git a/packages/plugins/trpc/src/generator.ts b/packages/plugins/trpc/src/generator.ts
index 32530b836..32cc64582 100644
--- a/packages/plugins/trpc/src/generator.ts
+++ b/packages/plugins/trpc/src/generator.ts
@@ -193,6 +193,10 @@ function createAppRouter(
             writer.block(() => {
                 for (const modelOperation of filteredModelOperations) {
                     const { model, ...operations } = modelOperation;
+
+                    // "count" operation is missing from Prisma DMMF, add it here
+                    operations.count = `count${model}`;
+
                     generateModelCreateRouter(
                         project,
                         model,
diff --git a/packages/plugins/trpc/src/helpers.ts b/packages/plugins/trpc/src/helpers.ts
index 4bbe18460..fadd752b2 100644
--- a/packages/plugins/trpc/src/helpers.ts
+++ b/packages/plugins/trpc/src/helpers.ts
@@ -300,6 +300,9 @@ export const getInputSchemaByOpName = (opName: string, modelName: string) => {
         case 'groupBy':
             inputType = `${modelName}InputSchema.groupBy`;
             break;
+        case 'count':
+            inputType = `${modelName}InputSchema.count`;
+            break;
         default:
             console.log('getInputTypeByOpName: ', { opName, modelName });
     }
@@ -316,6 +319,7 @@ export const getProcedureTypeByOpName = (opName: string) => {
         case 'aggregate':
         case 'aggregateRaw':
         case 'groupBy':
+        case 'count':
             procType = 'query';
             break;
         case 'createOne':
diff --git a/packages/schema/src/plugins/zod/transformer.ts b/packages/schema/src/plugins/zod/transformer.ts
index 2ad22d2a8..9d5bf9e20 100644
--- a/packages/schema/src/plugins/zod/transformer.ts
+++ b/packages/schema/src/plugins/zod/transformer.ts
@@ -594,6 +594,20 @@ export const ${this.name}ObjectSchema: SchemaType = ${schema} as SchemaType;`;
                 }
             }
 
+            // count
+            {
+                imports.push(
+                    `import { ${modelName}WhereInputObjectSchema } from '../objects/${modelName}WhereInput.schema'`,
+                    `import { ${orderByWithRelationInput}ObjectSchema } from '../objects/${orderByWithRelationInput}.schema'`,
+                    `import { ${modelName}WhereUniqueInputObjectSchema } from '../objects/${modelName}WhereUniqueInput.schema'`,
+                    `import { ${modelName}ScalarFieldEnumSchema } from '../enums/${modelName}ScalarFieldEnum.schema'`,
+                    `import { ${modelName}CountAggregateInputObjectSchema } from '../objects/${modelName}CountAggregateInput.schema'`
+                );
+
+                codeBody += `count: z.object({ where: ${modelName}WhereInputObjectSchema.optional(), orderBy: z.union([${orderByWithRelationInput}ObjectSchema, ${orderByWithRelationInput}ObjectSchema.array()]).optional(), cursor: ${modelName}WhereUniqueInputObjectSchema.optional(), take: z.number().optional(), skip: z.number().optional(), distinct: z.array(${modelName}ScalarFieldEnumSchema).optional(), select: z.union([ z.literal(true), ${modelName}CountAggregateInputObjectSchema ]).optional() })`;
+                operations.push(['count', origModelName]);
+            }
+
             imports = [...new Set(imports)];
 
             const filePath = path.join(Transformer.outputPath, `input/${modelName}Input.schema.ts`);