fix(codegen): include partition keys in hooks update/delete mutation variables

graphql/codegen/src/core/codegen/index.ts

@@ -249,6 +249,7 @@ export function generate(options: GenerateOptions): GenerateResult {
   const mutationHooks = generateAllMutationHooks(tables, {
     reactQueryEnabled,
     useCentralizedKeys,
+    typeRegistry: customOperations?.typeRegistry,
   });
   for (const hook of mutationHooks) {
     files.push({

graphql/codegen/src/core/codegen/mutations.ts

@@ -9,7 +9,7 @@
  */
 import * as t from '@babel/types';
 
-import type { Table } from '../../types/schema';
+import type { Table, TypeRegistry } from '../../types/schema';
 import {
   addJSDocComment,
   buildSelectionArgsCall,
@@ -50,11 +50,14 @@ import {
   getCreateMutationFileName,
   getCreateMutationHookName,
   getCreateMutationName,
+  getDeleteInputTypeName,
   getDeleteMutationFileName,
   getDeleteMutationHookName,
   getDeleteMutationName,
+  getExtraInputKeys,
   getPrimaryKeyInfo,
   getTableNames,
+  getUpdateInputTypeName,
   getUpdateMutationFileName,
   getUpdateMutationHookName,
   getUpdateMutationName,
@@ -70,6 +73,7 @@ export interface GeneratedMutationFile {
 export interface MutationGeneratorOptions {
   reactQueryEnabled?: boolean;
   useCentralizedKeys?: boolean;
+  typeRegistry?: TypeRegistry;
 }
 
 function buildMutationResultType(
@@ -332,7 +336,7 @@ export function generateUpdateMutationHook(
   table: Table,
   options: MutationGeneratorOptions = {},
 ): GeneratedMutationFile | null {
-  const { reactQueryEnabled = true, useCentralizedKeys = true } = options;
+  const { reactQueryEnabled = true, useCentralizedKeys = true, typeRegistry } = options;
 
   if (!reactQueryEnabled) return null;
   if (table.query?.update === null) return null;
@@ -355,6 +359,14 @@ export function generateUpdateMutationHook(
   const pkTsType =
     pkField.tsType === 'string' ? t.tsStringKeyword() : t.tsNumberKeyword();
 
+  const updateInputTypeName = getUpdateInputTypeName(table);
+  const updateExtraKeys = getExtraInputKeys(
+    updateInputTypeName,
+    new Set(pkFields.map((pk) => pk.name)),
+    patchFieldName,
+    typeRegistry,
+  );
+
   const statements: t.Statement[] = [];
 
   // Imports
@@ -409,12 +421,22 @@ export function generateUpdateMutationHook(
     ),
   );
 
-  // Variable type: { pkField: type; patchFieldName: PatchType }
+  // Variable type: { pkField: type; extraKeys...; patchFieldName: PatchType }
   const updateVarType = t.tsTypeLiteral([
     t.tsPropertySignature(
       t.identifier(pkField.name),
       t.tsTypeAnnotation(pkTsType),
     ),
+    ...updateExtraKeys.map((ek) =>
+      t.tsPropertySignature(
+        t.identifier(ek.name),
+        t.tsTypeAnnotation(
+          ek.tsType === 'number' ? t.tsNumberKeyword() :
+          ek.tsType === 'boolean' ? t.tsBooleanKeyword() :
+          t.tsStringKeyword()
+        ),
+      ),
+    ),
     t.tsPropertySignature(
       t.identifier(patchFieldName),
       t.tsTypeAnnotation(typeRef(patchTypeName)),
@@ -485,6 +507,7 @@ export function generateUpdateMutationHook(
   //   getClient().singular.update({ where: { pkField }, data: patchFieldName, select: ... }).unwrap()
   const destructParam = t.objectPattern([
     shorthandProp(pkField.name),
+    ...updateExtraKeys.map((ek) => shorthandProp(ek.name)),
     shorthandProp(patchFieldName),
   ]);
   destructParam.typeAnnotation = t.tsTypeAnnotation(updateVarType);
@@ -494,7 +517,10 @@ export function generateUpdateMutationHook(
       singularName,
       'update',
       t.objectExpression([
-        objectProp('where', t.objectExpression([shorthandProp(pkField.name)])),
+        objectProp('where', t.objectExpression([
+          shorthandProp(pkField.name),
+          ...updateExtraKeys.map((ek) => shorthandProp(ek.name)),
+        ])),
         objectProp('data', t.identifier(patchFieldName)),
         objectProp(
           'select',
@@ -592,7 +618,7 @@ export function generateDeleteMutationHook(
   table: Table,
   options: MutationGeneratorOptions = {},
 ): GeneratedMutationFile | null {
-  const { reactQueryEnabled = true, useCentralizedKeys = true } = options;
+  const { reactQueryEnabled = true, useCentralizedKeys = true, typeRegistry } = options;
 
   if (!reactQueryEnabled) return null;
   if (table.query?.delete === null) return null;
@@ -612,6 +638,14 @@ export function generateDeleteMutationHook(
   const pkTsType =
     pkField.tsType === 'string' ? t.tsStringKeyword() : t.tsNumberKeyword();
 
+  const deleteInputTypeName = getDeleteInputTypeName(table);
+  const deleteExtraKeys = getExtraInputKeys(
+    deleteInputTypeName,
+    new Set(pkFields.map((pk) => pk.name)),
+    null,
+    typeRegistry,
+  );
+
   const statements: t.Statement[] = [];
 
   // Imports
@@ -666,12 +700,22 @@ export function generateDeleteMutationHook(
     ),
   );
 
-  // Variable type: { pkField: type }
+  // Variable type: { pkField: type; extraKeys... }
   const deleteVarType = t.tsTypeLiteral([
     t.tsPropertySignature(
       t.identifier(pkField.name),
       t.tsTypeAnnotation(pkTsType),
     ),
+    ...deleteExtraKeys.map((ek) =>
+      t.tsPropertySignature(
+        t.identifier(ek.name),
+        t.tsTypeAnnotation(
+          ek.tsType === 'number' ? t.tsNumberKeyword() :
+          ek.tsType === 'boolean' ? t.tsBooleanKeyword() :
+          t.tsStringKeyword()
+        ),
+      ),
+    ),
   ]);
 
   const resultType = (sel: t.TSType) =>
@@ -736,15 +780,21 @@ export function generateDeleteMutationHook(
 
   // mutationFn: ({ pkField }: { pkField: type }) =>
   //   getClient().singular.delete({ where: { pkField }, select: ... }).unwrap()
-  const destructParam = t.objectPattern([shorthandProp(pkField.name)]);
+  const destructParam = t.objectPattern([
+    shorthandProp(pkField.name),
+    ...deleteExtraKeys.map((ek) => shorthandProp(ek.name)),
+  ]);
   destructParam.typeAnnotation = t.tsTypeAnnotation(deleteVarType);
   const mutationFnExpr = t.arrowFunctionExpression(
     [destructParam],
     getClientCallUnwrap(
       singularName,
       'delete',
       t.objectExpression([
-        objectProp('where', t.objectExpression([shorthandProp(pkField.name)])),
+        objectProp('where', t.objectExpression([
+          shorthandProp(pkField.name),
+          ...deleteExtraKeys.map((ek) => shorthandProp(ek.name)),
+        ])),
         objectProp(
           'select',
           t.memberExpression(t.identifier('args'), t.identifier('select')),

graphql/codegen/src/core/codegen/orm/model-generator.ts

@@ -15,6 +15,7 @@ import {
   getCreateMutationName,
   getDeleteInputTypeName,
   getDeleteMutationName,
+  getExtraInputKeys,
   getFilterTypeName,
   getGeneratedFileHeader,
   getOrderByTypeName,
@@ -25,6 +26,7 @@ import {
   lcFirst,
   ucFirst,
 } from '../utils';
+import type { ExtraInputKey } from '../utils';
 
 export interface GeneratedModelFile {
   fileName: string;
@@ -168,43 +170,7 @@ function strictSelectGuard(selectTypeName: string): t.TSType {
   );
 }
 
-interface ExtraInputKey {
-  name: string;
-  gqlType: string;
-  tsType: string;
-}
-
-/**
- * Discover extra required fields on a mutation input type beyond the standard
- * ones (clientMutationId, PK fields, patch field). PostGraphile adds these for
- * partitioned tables (e.g. databaseId as the partition key).
- */
-function getExtraInputKeys(
-  inputTypeName: string,
-  pkFieldNames: Set<string>,
-  patchFieldName: string | null,
-  typeRegistry?: TypeRegistry,
-): ExtraInputKey[] {
-  if (!typeRegistry) return [];
-  const inputType = typeRegistry.get(inputTypeName);
-  if (!inputType || inputType.kind !== 'INPUT_OBJECT' || !inputType.inputFields) return [];
-
-  const skip = new Set<string>(['clientMutationId', ...(patchFieldName ? [patchFieldName] : [])]);
-  for (const pk of pkFieldNames) skip.add(pk);
-
-  const extras: ExtraInputKey[] = [];
-  for (const field of inputType.inputFields) {
-    if (skip.has(field.name)) continue;
-    if (field.type.kind !== 'NON_NULL') continue;
-    const innerName = field.type.ofType?.name;
-    if (!innerName) continue;
-    let tsType = 'string';
-    if (innerName === 'Int' || innerName === 'Float' || innerName === 'BigFloat') tsType = 'number';
-    else if (innerName === 'Boolean') tsType = 'boolean';
-    extras.push({ name: field.name, gqlType: innerName, tsType });
-  }
-  return extras;
-}
+// ExtraInputKey type and getExtraInputKeys are imported from ../utils
 
 export function generateModelFile(
   table: Table,

graphql/codegen/src/core/codegen/utils.ts

@@ -314,6 +314,48 @@ export function getDeleteInputTypeName(table: Table): string {
   return mutationName ? ucFirst(mutationName) + 'Input' : `Delete${table.name}Input`;
 }
 
+// ============================================================================
+// Extra input keys (partition keys for update/delete mutations)
+// ============================================================================
+
+export interface ExtraInputKey {
+  name: string;
+  gqlType: string;
+  tsType: string;
+}
+
+/**
+ * Discover extra required fields on a mutation input type beyond the standard
+ * ones (clientMutationId, PK fields, patch field). PostGraphile adds these for
+ * partitioned tables (e.g. databaseId as the partition key).
+ */
+export function getExtraInputKeys(
+  inputTypeName: string,
+  pkFieldNames: Set<string>,
+  patchFieldName: string | null,
+  typeRegistry?: TypeRegistry,
+): ExtraInputKey[] {
+  if (!typeRegistry) return [];
+  const inputType = typeRegistry.get(inputTypeName);
+  if (!inputType || inputType.kind !== 'INPUT_OBJECT' || !inputType.inputFields) return [];
+
+  const skip = new Set<string>(['clientMutationId', ...(patchFieldName ? [patchFieldName] : [])]);
+  for (const pk of pkFieldNames) skip.add(pk);
+
+  const extras: ExtraInputKey[] = [];
+  for (const field of inputType.inputFields) {
+    if (skip.has(field.name)) continue;
+    if (field.type.kind !== 'NON_NULL') continue;
+    const innerName = field.type.ofType?.name;
+    if (!innerName) continue;
+    let tsType = 'string';
+    if (innerName === 'Int' || innerName === 'Float' || innerName === 'BigFloat') tsType = 'number';
+    else if (innerName === 'Boolean') tsType = 'boolean';
+    extras.push({ name: field.name, gqlType: innerName, tsType });
+  }
+  return extras;
+}
+
 // ============================================================================
 // Type mapping: GraphQL → TypeScript
 // ============================================================================