From 74736ac85f7a0c3463424ae7c8771bac06b38d2a Mon Sep 17 00:00:00 2001
From: ymc9 <104139426+ymc9@users.noreply.github.com>
Date: Fri, 5 May 2023 15:24:28 -0700
Subject: [PATCH] fix: wrong validation error when relation field is marked @id
 instead of @unique

---
 packages/schema/src/language-server/utils.ts  |  4 +-
 .../validator/datamodel-validator.ts          |  2 +-
 .../tests/regression/issues.test.ts           | 61 +++++++++++++++++++
 3 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/packages/schema/src/language-server/utils.ts b/packages/schema/src/language-server/utils.ts
index dcf6a9f37..775858fa7 100644
--- a/packages/schema/src/language-server/utils.ts
+++ b/packages/schema/src/language-server/utils.ts
@@ -51,7 +51,9 @@ export function getIdFields(model: DataModel) {
  * Gets lists of unique fields declared at the data model level
  */
 export function getUniqueFields(model: DataModel) {
-    const uniqueAttrs = model.attributes.filter((attr) => attr.decl.ref?.name === '@@unique');
+    const uniqueAttrs = model.attributes.filter(
+        (attr) => attr.decl.ref?.name === '@@unique' || attr.decl.ref?.name === '@@id'
+    );
     return uniqueAttrs.map((uniqueAttr) => {
         const fieldsArg = uniqueAttr.args.find((a) => a.$resolvedParam?.name === 'fields');
         if (!fieldsArg || !isArrayExpr(fieldsArg.value)) {
diff --git a/packages/schema/src/language-server/validator/datamodel-validator.ts b/packages/schema/src/language-server/validator/datamodel-validator.ts
index ec305330f..e8da86d22 100644
--- a/packages/schema/src/language-server/validator/datamodel-validator.ts
+++ b/packages/schema/src/language-server/validator/datamodel-validator.ts
@@ -342,7 +342,7 @@ export default class DataModelValidator implements AstValidator<DataModel> {
             thisRelation.fields?.forEach((ref) => {
                 const refField = ref.target.ref as DataModelField;
                 if (refField) {
-                    if (refField.attributes.find((a) => a.decl.ref?.name === '@unique')) {
+                    if (refField.attributes.find((a) => a.decl.ref?.name === '@id' || a.decl.ref?.name === '@unique')) {
                         return;
                     }
                     if (uniqueFieldList.some((list) => list.includes(refField))) {
diff --git a/tests/integration/tests/regression/issues.test.ts b/tests/integration/tests/regression/issues.test.ts
index 6683a9de2..f62d2ebe7 100644
--- a/tests/integration/tests/regression/issues.test.ts
+++ b/tests/integration/tests/regression/issues.test.ts
@@ -54,4 +54,65 @@ describe('GitHub issues regression', () => {
         expect(queried.posts[0].zenstack_guard).toBeUndefined();
         expect(queried.posts[0].zenstack_transaction).toBeUndefined();
     });
+
+    it('issue 392', async () => {
+        await loadSchema(
+            `
+            model M1 {
+                m2_id String @id
+                m2 M2 @relation(fields: [m2_id], references: [id])
+            }
+              
+            model M2 {
+                id String @id  
+                m1 M1?
+            }
+              `
+        );
+
+        await loadSchema(
+            `
+            model M1 {
+                id String @id
+                m2_id String @unique
+                m2 M2 @relation(fields: [m2_id], references: [id])
+            }
+              
+            model M2 {
+                id String @id  
+                m1 M1?
+            }
+              `
+        );
+
+        await loadSchema(
+            `
+            model M1 {
+                m2_id String
+                m2 M2 @relation(fields: [m2_id], references: [id])
+                @@id([m2_id])
+            }
+              
+            model M2 {
+                id String @id  
+                m1 M1?
+            }
+              `
+        );
+
+        await loadSchema(
+            `
+            model M1 {
+                m2_id String
+                m2 M2 @relation(fields: [m2_id], references: [id])
+                @@unique([m2_id])
+            }
+              
+            model M2 {
+                id String @id  
+                m1 M1?
+            }
+              `
+        );
+    });
 });