fix(policy): Prisma extension computed fields are not returned when fields level polices are used

fixes #2104

Author: ymc9

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

@@ -1204,47 +1204,42 @@ export class PolicyUtil extends QueryUtils {
             return;
         }
 
-        let target: any; // injection target
-        let isInclude = false; // if the target is include or select
-
-        if (args.select) {
-            target = args.select;
-            isInclude = false;
-        } else if (args.include) {
-            target = args.include;
-            isInclude = true;
-        } else {
-            target = args.select = this.makeAllScalarFieldSelect(model);
-            isInclude = false;
-        }
-
-        if (!isInclude) {
-            // merge selects
-            for (const [k, v] of Object.entries(input.select)) {
-                if (v === true) {
-                    if (!target[k]) {
-                        target[k] = true;
-                    }
+        // process scalar field selections first
+        for (const [k, v] of Object.entries<any>(input.select)) {
+            const field = resolveField(this.modelMeta, model, k);
+            if (!field || field.isDataModel) {
+                continue;
+            }
+            if (v === true) {
+                if (!args.select) {
+                    // do nothing since all scalar fields are selected by default
+                } else if (args.include) {
+                    // do nothing since include implies selecting all scalar fields
+                } else {
+                    args.select[k] = true;
                 }
             }
         }
 
-        // recurse into nested selects (relation fields)
+        // process relation selections
         for (const [k, v] of Object.entries<any>(input.select)) {
-            if (typeof v === 'object' && v?.select) {
-                const field = resolveField(this.modelMeta, model, k);
-                if (field?.isDataModel) {
-                    // recurse into relation
-                    if (isInclude && target[k] === true) {
-                        // select all fields for the relation
-                        target[k] = { select: this.makeAllScalarFieldSelect(field.type) };
-                    } else if (!target[k]) {
-                        // ensure an empty select clause
-                        target[k] = { select: {} };
-                    }
-                    // recurse
-                    this.doInjectReadCheckSelect(field.type, target[k], v);
-                }
+            const field = resolveField(this.modelMeta, model, k);
+            if (!field || !field.isDataModel) {
+                continue;
+            }
+
+            // prepare the next level of args
+            let nextArgs = args.select ?? args.include;
+            if (!nextArgs) {
+                nextArgs = args.include = {};
+            }
+            if (!nextArgs[k] || typeof nextArgs[k] !== 'object') {
+                nextArgs[k] = {};
+            }
+
+            if (v && typeof v === 'object') {
+                // recurse into relation
+                this.doInjectReadCheckSelect(field.type, nextArgs[k], v);
             }
         }
     }

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

@@ -0,0 +1,53 @@
+import { loadSchema } from '@zenstackhq/testtools';
+
+describe('issue 2104', () => {
+    it('regression', async () => {
+        const { enhance } = await loadSchema(
+            `
+        model Post {
+            id Int @id @default(autoincrement())
+            tickets TicketInPost[]
+            @@allow('all', true)
+        }
+
+        model View {
+            id Int @id @default(autoincrement())
+            ability String
+            tickets TicketInPost[]
+            @@allow('all', true)
+        }
+
+        model TicketInPost {
+            id Int @id @default(autoincrement())
+            postId Int
+            post Post @relation(fields: [postId], references: [id])
+            viewId Int
+            view View @relation(fields: [viewId], references: [id])
+            @@allow('all', true)
+        }
+            `,
+            { logPrismaQuery: true, enhancements: ['policy', 'delegate'] }
+        );
+
+        const db = enhance();
+        const view = await db.view.create({ data: { ability: 'bla' } });
+        const post = await db.post.create({ data: {} });
+        const updatedPost = await db.post.update({
+            where: {
+                id: post.id,
+            },
+            data: {
+                tickets: {
+                    deleteMany: {},
+                    create: {
+                        viewId: view.id,
+                    },
+                },
+            },
+            include: {
+                tickets: true,
+            },
+        });
+        console.log('created!', updatedPost);
+    });
+});