From 3a8fee86c0b45597158af3d8dcb8024299ca6373 Mon Sep 17 00:00:00 2001
From: JG <jiashengguo@outlook.com>
Date: Mon, 20 Mar 2023 18:26:25 +0000
Subject: [PATCH] fix:Support code action of generating relation fields for
 both sides

---
 .../src/language-server/zmodel-code-action.ts | 58 ++++++++++++-------
 1 file changed, 37 insertions(+), 21 deletions(-)

diff --git a/packages/schema/src/language-server/zmodel-code-action.ts b/packages/schema/src/language-server/zmodel-code-action.ts
index 23a64ca62..d8d5d69bb 100644
--- a/packages/schema/src/language-server/zmodel-code-action.ts
+++ b/packages/schema/src/language-server/zmodel-code-action.ts
@@ -66,11 +66,7 @@ export class ZModelCodeActionProvider implements CodeActionProvider {
 
             const container = getContainerOfType(cstNode?.element, isDataModel) as DataModel;
 
-            const idField = container.fields.find((f) =>
-                f.attributes.find((attr) => attr.decl.ref?.name === '@id')
-            ) as DataModelField;
-
-            if (container && container.$cstNode && idField) {
+            if (container && container.$cstNode) {
                 // indent
                 let indent = '\t';
                 const formatOptions = this.formatter.getFormatOptions();
@@ -79,17 +75,41 @@ export class ZModelCodeActionProvider implements CodeActionProvider {
                 }
                 indent = indent.repeat(this.formatter.getIndent());
 
-                const typeName = container.name;
-                const fieldName = this.lowerCaseFirstLetter(typeName);
-
-                // might already exist
-                let referenceField = '';
-
-                const idFieldName = idField.name;
-                const referenceIdFieldName = fieldName + this.upperCaseFirstLetter(idFieldName);
-
-                if (!oppositeModel.fields.find((f) => f.name === referenceIdFieldName)) {
-                    referenceField = '\n' + indent + `${referenceIdFieldName} ${idField.type.type}`;
+                let newText = '';
+                if (astNode.type.array) {
+                    //post Post[]
+                    const idField = container.fields.find((f) =>
+                        f.attributes.find((attr) => attr.decl.ref?.name === '@id')
+                    ) as DataModelField;
+
+                    // if no id field, we can't generate reference
+                    if (!idField) {
+                        return undefined;
+                    }
+
+                    const typeName = container.name;
+                    const fieldName = this.lowerCaseFirstLetter(typeName);
+
+                    // might already exist
+                    let referenceField = '';
+
+                    const idFieldName = idField.name;
+                    const referenceIdFieldName = fieldName + this.upperCaseFirstLetter(idFieldName);
+
+                    if (!oppositeModel.fields.find((f) => f.name === referenceIdFieldName)) {
+                        referenceField = '\n' + indent + `${referenceIdFieldName} ${idField.type.type}`;
+                    }
+
+                    newText =
+                        '\n' +
+                        indent +
+                        `${fieldName} ${typeName} @relation(fields: [${referenceIdFieldName}], references: [${idFieldName}])` +
+                        referenceField;
+                } else {
+                    // user User @relation(fields: [userAbc], references: [id])
+                    const typeName = container.name;
+                    const fieldName = this.lowerCaseFirstLetter(typeName);
+                    newText = '\n' + indent + `${fieldName} ${typeName}[]`;
                 }
 
                 return {
@@ -107,11 +127,7 @@ export class ZModelCodeActionProvider implements CodeActionProvider {
                                         // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
                                         end: lastField.$cstNode!.range.end,
                                     },
-                                    newText:
-                                        '\n' +
-                                        indent +
-                                        `${fieldName} ${typeName} @relation(fields: [${referenceIdFieldName}], references: [${idFieldName}])` +
-                                        referenceField,
+                                    newText,
                                 },
                             ],
                         },