Fix upsert validation to merge $create and $update schemas (#604)

* Initial plan

* fix: merge $create and $update schemas for upsert validation

- Handle upsert operation specially to match TypeScript type behavior
- When both $create and $update schemas exist, merge them for upsert
- Add test case to verify the fix works correctly

Co-authored-by: ymc9 <104139426+ymc9@users.noreply.github.com>

* fix: improve comment accuracy about Zod merge behavior

Co-authored-by: ymc9 <104139426+ymc9@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: ymc9 <104139426+ymc9@users.noreply.github.com>

Author: Copilot

packages/orm/src/client/crud/validator/index.ts

@@ -392,6 +392,34 @@ export class InputValidator<Schema extends SchemaDef> {
         if (operation in plugin.queryArgs && plugin.queryArgs[operation]) {
             // most specific operation takes highest precedence
             result = plugin.queryArgs[operation];
+        } else if (operation === 'upsert') {
+            // upsert is special: it's in both CoreCreateOperations and CoreUpdateOperations
+            // so we need to merge both $create and $update schemas to match the type system
+            const createSchema =
+                '$create' in plugin.queryArgs && plugin.queryArgs['$create']
+                    ? plugin.queryArgs['$create']
+                    : undefined;
+            const updateSchema =
+                '$update' in plugin.queryArgs && plugin.queryArgs['$update']
+                    ? plugin.queryArgs['$update']
+                    : undefined;
+
+            if (createSchema && updateSchema) {
+                invariant(
+                    createSchema instanceof z.ZodObject,
+                    'Plugin extended query args schema must be a Zod object',
+                );
+                invariant(
+                    updateSchema instanceof z.ZodObject,
+                    'Plugin extended query args schema must be a Zod object',
+                );
+                // merge both schemas (combines their properties)
+                result = createSchema.merge(updateSchema);
+            } else if (createSchema) {
+                result = createSchema;
+            } else if (updateSchema) {
+                result = updateSchema;
+            }
         } else if (
             // then comes grouped operations: $create, $read, $update, $delete
             CoreCreateOperations.includes(operation as CoreCreateOperations) &&

tests/e2e/orm/plugin-infra/ext-query-args.test.ts

@@ -230,4 +230,92 @@ describe('Plugin extended query args', () => {
         await expect(db.user.findMany({ cache: { ttl: 2000 } })).rejects.toThrow('Unrecognized key');
         await expect(extDb.user.findMany({ cache: { ttl: 2000 } })).toResolveWithLength(0);
     });
+
+    it('should merge $create and $update schemas for upsert operation', async () => {
+        // Define different schemas for $create and $update
+        const createOnlySchema = z.object({
+            tracking: z
+                .strictObject({
+                    source: z.string().optional(),
+                })
+                .optional(),
+        });
+
+        const updateOnlySchema = z.object({
+            audit: z
+                .strictObject({
+                    reason: z.string().optional(),
+                })
+                .optional(),
+        });
+
+        const extDb = db.$use(
+            definePlugin({
+                id: 'test',
+                queryArgs: {
+                    $create: createOnlySchema,
+                    $update: updateOnlySchema,
+                },
+            }),
+        );
+
+        // upsert should accept both tracking (from $create) and audit (from $update)
+        await expect(
+            extDb.user.upsert({
+                where: { id: 999 },
+                create: { name: 'Alice' },
+                update: { name: 'Alice Updated' },
+                tracking: { source: 'test' },
+                audit: { reason: 'testing merge' },
+            }),
+        ).resolves.toMatchObject({ name: 'Alice' });
+
+        // upsert should reject tracking-only in update operations
+        await expect(
+            extDb.user.update({
+                where: { id: 1 },
+                data: { name: 'Test' },
+                // @ts-expect-error - tracking is only for $create
+                tracking: { source: 'test' },
+            }),
+        ).rejects.toThrow('Unrecognized key');
+
+        // upsert should reject audit-only in create operations
+        await expect(
+            extDb.user.create({
+                data: { name: 'Bob' },
+                // @ts-expect-error - audit is only for $update
+                audit: { reason: 'test' },
+            }),
+        ).rejects.toThrow('Unrecognized key');
+
+        // verify that upsert without both is fine
+        await expect(
+            extDb.user.upsert({
+                where: { id: 888 },
+                create: { name: 'Charlie' },
+                update: { name: 'Charlie Updated' },
+            }),
+        ).resolves.toMatchObject({ name: 'Charlie' });
+
+        // verify that upsert with only tracking is fine
+        await expect(
+            extDb.user.upsert({
+                where: { id: 777 },
+                create: { name: 'David' },
+                update: { name: 'David Updated' },
+                tracking: { source: 'test' },
+            }),
+        ).resolves.toMatchObject({ name: 'David' });
+
+        // verify that upsert with only audit is fine
+        await expect(
+            extDb.user.upsert({
+                where: { id: 666 },
+                create: { name: 'Eve' },
+                update: { name: 'Eve Updated' },
+                audit: { reason: 'testing' },
+            }),
+        ).resolves.toMatchObject({ name: 'Eve' });
+    });
 });