From e73cf7acbc5f4509e94c0ae6ba2b2a8548cf7ece Mon Sep 17 00:00:00 2001
From: Andy Yen <38731840+onlyjackfrost@users.noreply.github.com>
Date: Mon, 25 Mar 2024 19:38:15 +0800
Subject: [PATCH] Fix(wren-ai): relation should generate column in related
 models (#34)

* fix(wren-ai): relation should generate column in related models

* rename function
---
 wren-ui/src/apollo/server/mdl/mdlBuilder.ts   | 58 ++++++++++++++++++-
 .../apollo/server/mdl/test/mdlBuilder.test.ts | 14 +++++
 wren-ui/src/apollo/server/mdl/type.ts         |  1 +
 3 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts
index beff8152f..1368a2f73 100644
--- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts
+++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts
@@ -133,11 +133,29 @@ export class MDLBuilder implements IMDLBuilder {
   public addRelation(): void {
     this.manifest.relationships = this.relations.map(
       (relation: RelationInfo) => {
+        const {
+          name,
+          joinType,
+          fromModelName,
+          fromColumnName,
+          toModelName,
+          toColumnName,
+        } = relation;
         const condition = this.getRelationCondition(relation);
+        this.addRelationColumn(fromModelName, {
+          modelReferenceName: toModelName,
+          columnReferenceName: toColumnName,
+          relation: name,
+        });
+        this.addRelationColumn(toModelName, {
+          modelReferenceName: fromModelName,
+          columnReferenceName: fromColumnName,
+          relation: name,
+        });
         return {
-          name: relation.name,
-          models: [relation.fromModelName, relation.toModelName],
-          joinType: relation.joinType,
+          name: name,
+          models: [fromModelName, toModelName],
+          joinType: joinType,
           condition,
         };
       },
@@ -149,6 +167,40 @@ export class MDLBuilder implements IMDLBuilder {
     this.manifest.catalog = this.project.catalog;
   }
 
+  protected addRelationColumn(
+    modelName: string,
+    columnData: {
+      modelReferenceName: string;
+      columnReferenceName: string;
+      relation: string;
+    },
+  ) {
+    const model = this.manifest.models.find(
+      (model: any) => model.name === modelName,
+    );
+    if (!model) {
+      logger.debug(`Can not find model "${modelName}" to add relation column`);
+      return;
+    }
+    if (!model.columns) {
+      model.columns = [];
+    }
+    // check if the modelReferenceName is already in the model column
+    const modelNameDuplicated = model.columns.find(
+      (column: any) => column.name === columnData.modelReferenceName,
+    );
+
+    model.columns.push({
+      name: modelNameDuplicated
+        ? `${columnData.modelReferenceName}_${columnData.columnReferenceName}`
+        : columnData.modelReferenceName,
+      type: columnData.modelReferenceName,
+      relationship: columnData.relation,
+      isCalculated: false,
+      notNull: false,
+    });
+  }
+
   protected getColumnExpression(column: ModelColumn): string {
     if (column.isCalculated) {
       // calculated field
diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts
index 3139b1493..d6afccfa3 100644
--- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts
+++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts
@@ -137,6 +137,13 @@ describe('MDLBuilder', () => {
               notNull: true,
               properties: { description: 'bar' },
             },
+            {
+              name: 'customer',
+              type: 'customer',
+              isCalculated: false,
+              relationship: 'OrderCustomer',
+              notNull: false,
+            },
           ],
           cached: false,
           refreshTime: null,
@@ -155,6 +162,13 @@ describe('MDLBuilder', () => {
               notNull: true,
               properties: null,
             },
+            {
+              name: 'order',
+              type: 'order',
+              isCalculated: false,
+              relationship: 'OrderCustomer',
+              notNull: false,
+            },
           ],
           primaryKey: '',
           cached: false,
diff --git a/wren-ui/src/apollo/server/mdl/type.ts b/wren-ui/src/apollo/server/mdl/type.ts
index 44069ca70..1a4c0e500 100644
--- a/wren-ui/src/apollo/server/mdl/type.ts
+++ b/wren-ui/src/apollo/server/mdl/type.ts
@@ -3,6 +3,7 @@ export interface ColumnMDL {
   type?: string; // eg: "integer", "string", "relationName"
   isCalculated: boolean; // eg: true, false
   notNull?: boolean; // eg: true, false
+  relationship?: string; //eg: OrdersCustomer
   properties?: {
     description?: string; // eg: "the key of each order"
   };