more fixes

Author: ymc9

packages/runtime/src/enhancements/node/delegate.ts

@@ -587,6 +587,7 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler {
         let curr = args;
         let base = this.getBaseModel(model);
         let sub = this.getModelInfo(model);
+        const hasDelegateBase = !!base;
 
         while (base) {
             const baseRelationName = this.makeAuxRelationName(base);
@@ -615,6 +616,50 @@ export class DelegateProxyHandler extends DefaultPrismaProxyHandler {
             sub = base;
             base = this.getBaseModel(base.name);
         }
+
+        if (hasDelegateBase) {
+            // A delegate base model creation is added, this can be incompatible if
+            // the user-provided payload assigns foreign keys directly, because Prisma
+            // doesn't permit mixed "checked" and "unchecked" fields in a payload.
+            //
+            // {
+            //   delegate_aux_base: { ... },
+            //   [fkField]: value  // <- this is not compatible
+            // }
+            //
+            // We need to convert foreign key assignments to `connect`.
+            this.fkAssignmentToConnect(model, args);
+        }
+    }
+
+    // convert foreign key assignments to `connect` payload
+    // e.g.: { authorId: value } -> { author: { connect: { id: value } } }
+    private fkAssignmentToConnect(model: string, args: any) {
+        for (const key of Object.keys(args)) {
+            const value = args[key];
+            if (value === undefined) {
+                continue;
+            }
+
+            const fieldInfo = this.queryUtils.getModelField(model, key);
+            if (
+                !fieldInfo?.inheritedFrom && // fields from delegate base are handled outside
+                fieldInfo?.isForeignKey
+            ) {
+                const relationInfo = this.queryUtils.getRelationForForeignKey(model, key);
+                if (relationInfo) {
+                    // turn { [fk]: value } into { [relation]: { connect: { [id]: value } } }
+                    if (!args[relationInfo.relation.name]) {
+                        args[relationInfo.relation.name] = {};
+                    }
+                    if (!args[relationInfo.relation.name].connect) {
+                        args[relationInfo.relation.name].connect = {};
+                    }
+                    args[relationInfo.relation.name].connect[relationInfo.idField] = value;
+                    delete args[key];
+                }
+            }
+        }
     }
 
     // inject field data that belongs to base type into proper nesting structure

packages/runtime/src/enhancements/node/query-utils.ts

@@ -232,4 +232,25 @@ export class QueryUtils {
 
         return model;
     }
+
+    /**
+     * Gets relation info for a foreign key field.
+     */
+    getRelationForForeignKey(model: string, fkField: string) {
+        const modelInfo = getModelInfo(this.options.modelMeta, model);
+        if (!modelInfo) {
+            return undefined;
+        }
+
+        for (const field of Object.values(modelInfo.fields)) {
+            if (field.foreignKeyMapping) {
+                const entry = Object.entries(field.foreignKeyMapping).find(([, v]) => v === fkField);
+                if (entry) {
+                    return { relation: field, idField: entry[0], fkField: entry[1] };
+                }
+            }
+        }
+
+        return undefined;
+    }
 }

packages/schema/src/plugins/zod/transformer.ts

@@ -1,6 +1,12 @@
 /* eslint-disable @typescript-eslint/ban-ts-comment */
-import { indentString, isDiscriminatorField, type PluginOptions } from '@zenstackhq/sdk';
-import { DataModel, Enum, isDataModel, isEnum, isTypeDef, type Model } from '@zenstackhq/sdk/ast';
+import {
+    getForeignKeyFields,
+    hasAttribute,
+    indentString,
+    isDiscriminatorField,
+    type PluginOptions,
+} from '@zenstackhq/sdk';
+import { DataModel, DataModelField, Enum, isDataModel, isEnum, isTypeDef, type Model } from '@zenstackhq/sdk/ast';
 import { checkModelHasModelRelation, findModelByName, isAggregateInputType } from '@zenstackhq/sdk/dmmf-helpers';
 import { supportCreateMany, type DMMF as PrismaDMMF } from '@zenstackhq/sdk/prisma';
 import path from 'path';
@@ -241,7 +247,8 @@ export default class Transformer {
                             this.addSchemaImport(inputType.type);
                         }
 
-                        result.push(this.generatePrismaStringLine(field, inputType, lines.length));
+                        const contextField = contextDataModel?.fields.find((f) => f.name === field.name);
+                        result.push(this.generatePrismaStringLine(field, inputType, lines.length, contextField));
                     }
                 }
 
@@ -315,7 +322,12 @@ export default class Transformer {
         this.schemaImports.add(upperCaseFirst(name));
     }
 
-    generatePrismaStringLine(field: PrismaDMMF.SchemaArg, inputType: PrismaDMMF.InputTypeRef, inputsLength: number) {
+    generatePrismaStringLine(
+        field: PrismaDMMF.SchemaArg,
+        inputType: PrismaDMMF.InputTypeRef,
+        inputsLength: number,
+        contextField: DataModelField | undefined
+    ) {
         const isEnum = inputType.location === 'enumTypes';
 
         const { isModelQueryType, modelName, queryName } = this.checkIsModelQueryType(inputType.type as string);
@@ -330,11 +342,36 @@ export default class Transformer {
 
         const arr = inputType.isList ? '.array()' : '';
 
-        const opt = !field.isRequired ? '.optional()' : '';
+        const optional =
+            !field.isRequired ||
+            // also check if the zmodel field infers the field as optional
+            (contextField && this.isFieldOptional(contextField));
 
         return inputsLength === 1
-            ? `  ${field.name}: z.lazy(() => ${schema})${arr}${opt}`
-            : `z.lazy(() => ${schema})${arr}${opt}`;
+            ? `  ${field.name}: z.lazy(() => ${schema})${arr}${optional ? '.optional()' : ''}`
+            : `z.lazy(() => ${schema})${arr}${optional ? '.optional()' : ''}`;
+    }
+
+    private isFieldOptional(dmField: DataModelField) {
+        if (hasAttribute(dmField, '@default')) {
+            // it's possible that ZModel field has a default but it's transformed away
+            // when generating Prisma schema, e.g.: `@default(auth().id)`
+            return true;
+        }
+
+        if (isDataModel(dmField.type.reference?.ref)) {
+            // if field is a relation, we need to check if the corresponding fk field has a default
+            // {
+            //   authorId Int @default(auth().id)
+            //   author User @relation(...)  // <- author should be optional
+            // }
+            const fkFields = getForeignKeyFields(dmField);
+            if (fkFields.every((fkField) => hasAttribute(fkField, '@default'))) {
+                return true;
+            }
+        }
+
+        return false;
     }
 
     generateFieldValidators(zodStringWithMainType: string, field: PrismaDMMF.SchemaArg) {

packages/sdk/src/utils.ts

@@ -381,6 +381,13 @@ export function getRelationField(fkField: DataModelField) {
     });
 }
 
+/**
+ * Gets the foreign key fields of the given relation field.
+ */
+export function getForeignKeyFields(relationField: DataModelField) {
+    return getRelationKeyPairs(relationField).map((pair) => pair.foreignKey);
+}
+
 export function resolvePath(_path: string, options: Pick<PluginOptions, 'schemaPath'>) {
     if (path.isAbsolute(_path)) {
         return _path;

tests/regression/tests/issue-1843.test.ts

@@ -2,7 +2,7 @@ import { loadSchema } from '@zenstackhq/testtools';
 
 describe('issue 1843', () => {
     it('regression', async () => {
-        await loadSchema(
+        const { zodSchemas, enhance, prisma } = await loadSchema(
             `
             model User {
                 id              String    @id @default(cuid())
@@ -43,9 +43,9 @@ describe('issue 1843', () => {
                 coauthorId String
 
                 @@allow('all', true)
-                }
+            }
 
-                model Post extends Content {
+            model Post extends Content {
                 title      String
 
                 @@allow('all', true)
@@ -85,5 +85,24 @@ describe('issue 1843', () => {
                 ],
             }
         );
+
+        const user = await prisma.user.create({ data: { email: 'abc', password: '123' } });
+        const db = enhance({ id: user.id });
+
+        // connect
+        await expect(
+            db.postWithCoauthor.create({ data: { title: 'new post', coauthor: { connect: { id: user.id } } } })
+        ).toResolveTruthy();
+
+        // fk setting
+        await expect(
+            db.postWithCoauthor.create({ data: { title: 'new post', coauthorId: user.id } })
+        ).toResolveTruthy();
+
+        // zod validation
+        zodSchemas.models.PostWithCoauthorCreateSchema.parse({
+            title: 'new post',
+            coauthorId: '1',
+        });
     });
 });