From 2017e620dc62f26c47e582b30a5ff175575665ac Mon Sep 17 00:00:00 2001
From: ymc9 <104139426+ymc9@users.noreply.github.com>
Date: Mon, 6 Jan 2025 21:45:07 +0800
Subject: [PATCH] fix(delegate): don't generate both "@@schema" attributes from
 base and sub models

---
 packages/schema/src/utils/ast-utils.ts    | 17 +++++-
 tests/regression/tests/issue-1647.test.ts | 69 +++++++++++++++++++++++
 2 files changed, 85 insertions(+), 1 deletion(-)
 create mode 100644 tests/regression/tests/issue-1647.test.ts

diff --git a/packages/schema/src/utils/ast-utils.ts b/packages/schema/src/utils/ast-utils.ts
index 0e462547f..f59ee7faa 100644
--- a/packages/schema/src/utils/ast-utils.ts
+++ b/packages/schema/src/utils/ast-utils.ts
@@ -16,7 +16,14 @@ import {
     ModelImport,
     TypeDef,
 } from '@zenstackhq/language/ast';
-import { getAttribute, getInheritanceChain, getRecursiveBases, isDelegateModel, isFromStdlib } from '@zenstackhq/sdk';
+import {
+    getAttribute,
+    getInheritanceChain,
+    getRecursiveBases,
+    hasAttribute,
+    isDelegateModel,
+    isFromStdlib,
+} from '@zenstackhq/sdk';
 import {
     AstNode,
     copyAstNode,
@@ -96,6 +103,9 @@ function filterBaseAttribute(forModel: DataModel, base: DataModel, attr: DataMod
     // uninheritable attributes for delegate inheritance (they reference fields from the base)
     const uninheritableFromDelegateAttributes = ['@@unique', '@@index', '@@fulltext'];
 
+    // attributes that are inherited but can be overridden
+    const overrideAttributes = ['@@schema'];
+
     if (uninheritableAttributes.includes(attr.decl.$refText)) {
         return false;
     }
@@ -109,6 +119,11 @@ function filterBaseAttribute(forModel: DataModel, base: DataModel, attr: DataMod
         return false;
     }
 
+    if (hasAttribute(forModel, attr.decl.$refText) && overrideAttributes.includes(attr.decl.$refText)) {
+        // don't inherit an attribute if it's overridden in the sub model
+        return false;
+    }
+
     return true;
 }
 
diff --git a/tests/regression/tests/issue-1647.test.ts b/tests/regression/tests/issue-1647.test.ts
new file mode 100644
index 000000000..e93f63cfb
--- /dev/null
+++ b/tests/regression/tests/issue-1647.test.ts
@@ -0,0 +1,69 @@
+import { loadSchema } from '@zenstackhq/testtools';
+import fs from 'fs';
+
+describe('issue 1647', () => {
+    it('inherits @@schema by default', async () => {
+        const { projectDir } = await loadSchema(
+            `
+            datasource db {
+                provider = 'postgresql'
+                url = env('DATABASE_URL')
+                schemas = ['public', 'post']
+            }
+
+            generator client {
+                provider = 'prisma-client-js'
+                previewFeatures = ['multiSchema']
+            }
+
+            model Asset {
+                id Int @id
+                type String
+                @@delegate(type)
+                @@schema('public')
+            }
+            
+            model Post extends Asset {
+                title String
+            }
+            `,
+            { addPrelude: false, pushDb: false, getPrismaOnly: true }
+        );
+
+        const prismaSchema = fs.readFileSync(`${projectDir}/prisma/schema.prisma`, 'utf-8');
+        expect(prismaSchema.split('\n').filter((l) => l.includes('@@schema("public")'))).toHaveLength(2);
+    });
+    it('respects sub model @@schema overrides', async () => {
+        const { projectDir } = await loadSchema(
+            `
+            datasource db {
+                provider = 'postgresql'
+                url = env('DATABASE_URL')
+                schemas = ['public', 'post']
+            }
+
+            generator client {
+                provider = 'prisma-client-js'
+                previewFeatures = ['multiSchema']
+            }
+
+            model Asset {
+                id Int @id
+                type String
+                @@delegate(type)
+                @@schema('public')
+            }
+            
+            model Post extends Asset {
+                title String
+                @@schema('post')
+            }
+            `,
+            { addPrelude: false, pushDb: false, getPrismaOnly: true }
+        );
+
+        const prismaSchema = fs.readFileSync(`${projectDir}/prisma/schema.prisma`, 'utf-8');
+        expect(prismaSchema.split('\n').filter((l) => l.includes('@@schema("public")'))).toHaveLength(1);
+        expect(prismaSchema.split('\n').filter((l) => l.includes('@@schema("post")'))).toHaveLength(1);
+    });
+});