constructive-io · pyramation · Jun 14, 2026 · Jun 14, 2026
diff --git a/graphql/codegen/src/__tests__/codegen/__snapshots__/model-generator.test.ts.snap b/graphql/codegen/src/__tests__/codegen/__snapshots__/model-generator.test.ts.snap
@@ -118,7 +118,7 @@ export class UserModel {
     const {
       document,
       variables
-    } = buildUpdateByPkDocument("User", "updateUser", "user", args.select, args.where.id, args.data, "UpdateUserInput", "id", "userPatch", connectionFieldsMap);
+    } = buildUpdateByPkDocument("User", "updateUser", "user", args.select, args.where.id, args.data, "UpdateUserInput", "id", "userPatch", connectionFieldsMap, undefined);
     return new QueryBuilder({
       client: this.client,
       operation: "mutation",
@@ -273,7 +273,7 @@ exports[`model-generator handles custom query/mutation names 1`] = `
 import { OrmClient } from "../client";
 import { QueryBuilder, buildFindManyDocument, buildFindFirstDocument, buildFindOneDocument, buildCreateDocument, buildUpdateByPkDocument, buildDeleteByPkDocument } from "../query-builder";
 import type { ConnectionResult, FindManyArgs, FindFirstArgs, CreateArgs, UpdateArgs, DeleteArgs, InferSelectResult, StrictSelect } from "../select-types";
-import type { Organization, OrganizationWithRelations, OrganizationSelect, OrganizationFilter, OrganizationsOrderBy, CreateOrganizationInput, UpdateOrganizationInput, OrganizationPatch } from "../input-types";
+import type { Organization, OrganizationWithRelations, OrganizationSelect, OrganizationFilter, OrganizationsOrderBy, CreateOrganizationInput, ModifyOrganizationInput, OrganizationPatch } from "../input-types";
 import { connectionFieldsMap } from "../input-types";
 export class OrganizationModel {
   constructor(private client: OrmClient) {}
@@ -382,7 +382,7 @@ export class OrganizationModel {
     const {
       document,
       variables
-    } = buildUpdateByPkDocument("Organization", "modifyOrganization", "organization", args.select, args.where.id, args.data, "UpdateOrganizationInput", "id", "organizationPatch", connectionFieldsMap);
+    } = buildUpdateByPkDocument("Organization", "modifyOrganization", "organization", args.select, args.where.id, args.data, "ModifyOrganizationInput", "id", "organizationPatch", connectionFieldsMap, undefined);
     return new QueryBuilder({
       client: this.client,
       operation: "mutation",

diff --git a/graphql/codegen/src/core/codegen/orm/index.ts b/graphql/codegen/src/core/codegen/orm/index.ts
@@ -95,7 +95,7 @@ export function generateOrm(options: GenerateOrmOptions): GenerateOrmResult {
   });
 
   // 2. Generate model files
-  const modelFiles = generateAllModelFiles(tables, useSharedTypes);
+  const modelFiles = generateAllModelFiles(tables, useSharedTypes, typeRegistry);
   for (const modelFile of modelFiles) {
     files.push({
       path: `models/${modelFile.fileName}`,

diff --git a/graphql/codegen/src/core/codegen/orm/model-generator.ts b/graphql/codegen/src/core/codegen/orm/model-generator.ts
@@ -8,7 +8,7 @@ import * as t from '@babel/types';
 
 import { singularize } from 'inflekt';
 
-import type { Table } from '../../../types/schema';
+import type { Table, TypeRegistry } from '../../../types/schema';
 import { asConst, generateCode } from '../babel-ast';
 import {
   getCreateInputTypeName,
@@ -20,6 +20,7 @@ import {
   getOrderByTypeName,
   getPrimaryKeyInfo,
   getTableNames,
+  getUpdateInputTypeName,
   hasValidPrimaryKey,
   lcFirst,
   ucFirst,
@@ -167,11 +168,50 @@ 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;
+}
+
 export function generateModelFile(
   table: Table,
   _useSharedTypes: boolean,
   options?: Record<string, never>,
   allTables?: Table[],
+  typeRegistry?: TypeRegistry,
 ): GeneratedModelFile {
   const { typeName, singularName, pluralName } = getTableNames(table);
   const modelName = `${typeName}Model`;
@@ -185,7 +225,7 @@ export function generateModelFile(
   const whereTypeName = getFilterTypeName(table);
   const orderByTypeName = getOrderByTypeName(table);
   const createInputTypeName = `Create${typeName}Input`;
-  const updateInputTypeName = `Update${typeName}Input`;
+  const updateInputTypeName = getUpdateInputTypeName(table);
   const patchTypeName = `${typeName}Patch`;
 
   const pkFields = getPrimaryKeyInfo(table);
@@ -845,6 +885,14 @@ export function generateModelFile(
 
   // ── update ─────────────────────────────────────────────────────────────
   if (updateMutationName) {
+    const patchFieldName =
+      table.query?.patchFieldName ?? lcFirst(typeName) + 'Patch';
+    const updateExtraKeys = getExtraInputKeys(
+      updateInputTypeName,
+      new Set(pkFields.map((pk) => pk.name)),
+      patchFieldName,
+      typeRegistry,
+    );
     const whereLiteral = () =>
       t.tsTypeLiteral([
         (() => {
@@ -855,6 +903,14 @@ export function generateModelFile(
           prop.optional = false;
           return prop;
         })(),
+        ...updateExtraKeys.map((ek) => {
+          const prop = t.tsPropertySignature(
+            t.identifier(ek.name),
+            t.tsTypeAnnotation(tsTypeFromPrimitive(ek.tsType)),
+          );
+          prop.optional = false;
+          return prop;
+        }),
       ]);
     const argsType = (sel: t.TSType) =>
       t.tsTypeReference(
@@ -907,8 +963,20 @@ export function generateModelFile(
       t.identifier('args'),
       t.identifier('select'),
     );
-    const patchFieldName =
-      table.query?.patchFieldName ?? lcFirst(typeName) + 'Patch';
+    // Build extraKeys object for partitioned table fields
+    const extraKeysArg = updateExtraKeys.length > 0
+      ? t.objectExpression(
+          updateExtraKeys.map((ek) =>
+            t.objectProperty(
+              t.identifier(ek.name),
+              t.memberExpression(
+                t.memberExpression(t.identifier('args'), t.identifier('where')),
+                t.identifier(ek.name),
+              ),
+            ),
+          ),
+        )
+      : t.identifier('undefined');
     const bodyArgs = [
       t.stringLiteral(typeName),
       t.stringLiteral(updateMutationName),
@@ -923,6 +991,7 @@ export function generateModelFile(
       t.stringLiteral(pkField.name),
       t.stringLiteral(patchFieldName),
       t.identifier('connectionFieldsMap'),
+      extraKeysArg,
     ];
     classBody.push(
       createClassMethod(
@@ -943,18 +1012,32 @@ export function generateModelFile(
 
   // ── delete ─────────────────────────────────────────────────────────────
   if (deleteMutationName) {
-    // Build where type with ALL PK fields (supports composite PKs)
+    const deleteExtraKeys = getExtraInputKeys(
+      deleteInputTypeName,
+      new Set(pkFields.map((pk) => pk.name)),
+      null,
+      typeRegistry,
+    );
+    // Build where type with ALL PK fields + extra keys (supports composite PKs + partition keys)
     const whereLiteral = () =>
-      t.tsTypeLiteral(
-        pkFields.map((pk) => {
+      t.tsTypeLiteral([
+        ...pkFields.map((pk) => {
           const prop = t.tsPropertySignature(
             t.identifier(pk.name),
             t.tsTypeAnnotation(tsTypeFromPrimitive(pk.tsType ?? 'string')),
           );
           prop.optional = false;
           return prop;
         }),
-      );
+        ...deleteExtraKeys.map((ek) => {
+          const prop = t.tsPropertySignature(
+            t.identifier(ek.name),
+            t.tsTypeAnnotation(tsTypeFromPrimitive(ek.tsType)),
+          );
+          prop.optional = false;
+          return prop;
+        }),
+      ]);
     const argsType = (sel: t.TSType) =>
       t.tsTypeReference(
         t.identifier('DeleteArgs'),
@@ -1002,9 +1085,9 @@ export function generateModelFile(
       t.identifier('args'),
       t.identifier('select'),
     );
-    // Build keys object: { field1: args.where.field1, field2: args.where.field2, ... }
-    const keysObj = t.objectExpression(
-      pkFields.map((pk) =>
+    // Build keys object: { field1: args.where.field1, ... } including extra keys
+    const keysObj = t.objectExpression([
+      ...pkFields.map((pk) =>
         t.objectProperty(
           t.identifier(pk.name),
           t.memberExpression(
@@ -1013,7 +1096,16 @@ export function generateModelFile(
           ),
         ),
       ),
-    );
+      ...deleteExtraKeys.map((ek) =>
+        t.objectProperty(
+          t.identifier(ek.name),
+          t.memberExpression(
+            t.memberExpression(t.identifier('args'), t.identifier('where')),
+            t.identifier(ek.name),
+          ),
+        ),
+      ),
+    ]);
     const bodyArgs = [
       t.stringLiteral(typeName),
       t.stringLiteral(deleteMutationName),
@@ -1516,6 +1608,7 @@ export function generateModelFile(
 export function generateAllModelFiles(
   tables: Table[],
   useSharedTypes: boolean,
+  typeRegistry?: TypeRegistry,
 ): GeneratedModelFile[] {
-  return tables.map((table) => generateModelFile(table, useSharedTypes, undefined, tables));
+  return tables.map((table) => generateModelFile(table, useSharedTypes, undefined, tables, typeRegistry));
 }
diff --git a/graphql/codegen/src/core/codegen/templates/query-builder.ts b/graphql/codegen/src/core/codegen/templates/query-builder.ts
@@ -509,6 +509,7 @@ export function buildUpdateByPkDocument<TSelect, TData>(
   idFieldName: string,
   patchFieldName: string,
   connectionFieldsMap?: Record<string, Record<string, string>>,
+  extraKeys?: Record<string, unknown>,
 ): { document: string; variables: Record<string, unknown> } {
   const selections = select
     ? buildSelections(
@@ -533,6 +534,7 @@ export function buildUpdateByPkDocument<TSelect, TData>(
     variables: {
       input: {
         [idFieldName]: id,
+        ...extraKeys,
         [patchFieldName]: data,
       },
     },