feat(appsync-modelgen-plugin): add sort key field in manyToMany models (#401)

* test: disable flow modern tests

* feat(appsync-modelgen-plugin): add sort key field in manyToMany models

* fix: rm only test

* fix: remove old sort key impl

Author: AaronZyLee

packages/appsync-modelgen-plugin/src/__tests__/visitors/appsync-visitor.test.ts

@@ -913,6 +913,41 @@ describe('AppSyncModelVisitor', () => {
     });
   });
 
+  describe('manyToMany with sort key testing', () => {
+    const schema = /* GraphQL */ `
+      type ModelA @model {
+        id: ID! @primaryKey(sortKeyFields: ["sortId"])
+        sortId: ID!
+        models: [ModelB] @manyToMany(relationName: "ModelAModelB")
+      }
+      type ModelB @model {
+        id: ID! @primaryKey(sortKeyFields: ["sortId"])
+        sortId: ID!
+        models: [ModelA] @manyToMany(relationName: "ModelAModelB")
+      }
+    `;
+    const { models } = createAndGeneratePipelinedTransformerVisitor(schema);
+    const { ModelAModelB } = models;
+    it('should generate correct fields and secondary keys for intermediate type', () => {
+      expect(ModelAModelB).toBeDefined();
+      expect(ModelAModelB.fields.length).toEqual(7);
+      const modelASortKeyField = ModelAModelB.fields.find(f => f.name === 'modelAsortId');
+      expect(modelASortKeyField).toBeDefined();
+      expect(modelASortKeyField.type).toEqual('ID');
+      expect(modelASortKeyField.isNullable).toBe(false);
+      const modelBSortKeyField = ModelAModelB.fields.find(f => f.name === 'modelBsortId');
+      expect(modelBSortKeyField).toBeDefined();
+      expect(modelBSortKeyField.type).toEqual('ID');
+      expect(modelBSortKeyField.isNullable).toBe(false);
+      const modelAIndexDirective = ModelAModelB.directives.find(d => d.name === 'key' && d.arguments.name === 'byModelA');
+      expect(modelAIndexDirective).toBeDefined();
+      expect(modelAIndexDirective.arguments.fields).toEqual(['modelAID', 'modelAsortId']);
+      const modelBIndexDirective = ModelAModelB.directives.find(d => d.name === 'key' && d.arguments.name === 'byModelB');
+      expect(modelBIndexDirective).toBeDefined();
+      expect(modelBIndexDirective.arguments.fields).toEqual(['modelBID', 'modelBsortId']);
+    });
+  });
+
   describe('Graphql V2 fix tests for multiple has many relations of only one model type', () => {
     const schema = /* GraphQL*/ `
       type Registration @model {
@@ -936,5 +971,5 @@ describe('AppSyncModelVisitor', () => {
     it(`should not throw error when processing models`, () => {
       expect(() => createAndGeneratePipelinedTransformerVisitor(schema)).not.toThrow();
     });
-  })
+  });
 });

packages/appsync-modelgen-plugin/src/visitors/appsync-visitor.ts

@@ -296,7 +296,7 @@ export class AppSyncModelVisitor<
   }
   processDirectives(
     // TODO: Remove us when we have a fix to roll-forward.
-    shouldUseModelNameFieldInHasManyAndBelongsTo: boolean
+    shouldUseModelNameFieldInHasManyAndBelongsTo: boolean,
   ) {
     if (this.config.usePipelinedTransformer || this.config.transformerVersion === 2) {
       this.processV2KeyDirectives();
@@ -551,15 +551,12 @@ export class AppSyncModelVisitor<
     });
   }
 
-  protected generateIntermediateModel(
-    firstModel: CodeGenModel,
-    secondModel: CodeGenModel,
-    firstField: CodeGenField,
-    secondField: CodeGenField,
-    relationName: string,
-  ) {
+  protected generateIntermediateModel(firstModel: CodeGenModel, secondModel: CodeGenModel, relationName: string) {
     const firstModelKeyFieldName = `${camelCase(firstModel.name)}ID`;
+    const firstModelSortKeyFields: CodeGenField[] = this.getSortKeyFields(firstModel);
     const secondModelKeyFieldName = `${camelCase(secondModel.name)}ID`;
+    const secondModelSortKeyFields: CodeGenField[] = this.getSortKeyFields(secondModel);
+
     let intermediateModel: CodeGenModel = {
       name: relationName,
       type: 'model',
@@ -577,15 +574,49 @@ export class AppSyncModelVisitor<
           isNullable: false,
           isList: false,
           name: firstModelKeyFieldName,
-          directives: [{ name: 'index', arguments: { name: 'by' + firstModel.name, sortKeyFields: [secondModelKeyFieldName] } }],
+          directives: [
+            {
+              name: 'index',
+              arguments: {
+                name: 'by' + firstModel.name,
+                sortKeyFields: firstModelSortKeyFields.map(f => this.generateIntermediateModelSortKeyFieldName(firstModel, f)),
+              },
+            },
+          ],
         },
+        ...firstModelSortKeyFields.map(field => {
+          return {
+            type: field.type,
+            isNullable: false,
+            isList: field.isList,
+            name: this.generateIntermediateModelSortKeyFieldName(firstModel, field),
+            directives: [],
+          };
+        }),
         {
           type: 'ID',
           isNullable: false,
           isList: false,
           name: secondModelKeyFieldName,
-          directives: [{ name: 'index', arguments: { name: 'by' + secondModel.name, sortKeyFields: [firstModelKeyFieldName] } }],
+          directives: [
+            {
+              name: 'index',
+              arguments: {
+                name: 'by' + secondModel.name,
+                sortKeyFields: secondModelSortKeyFields.map(f => this.generateIntermediateModelSortKeyFieldName(secondModel, f)),
+              },
+            },
+          ],
         },
+        ...secondModelSortKeyFields.map(field => {
+          return {
+            type: field.type,
+            isNullable: false,
+            isList: field.isList,
+            name: this.generateIntermediateModelSortKeyFieldName(secondModel, field),
+            directives: [],
+          };
+        }),
         {
           type: firstModel.name,
           isNullable: false,
@@ -606,6 +637,18 @@ export class AppSyncModelVisitor<
     return intermediateModel;
   }
 
+  protected generateIntermediateModelSortKeyFieldName(model: CodeGenModel, sortKeyField: CodeGenField): string {
+    const modelName = model.name.charAt(0).toLocaleLowerCase() + model.name.slice(1);
+    return `${modelName}${sortKeyField.name}`;
+  }
+
+  protected getSortKeyFields(model: CodeGenModel): CodeGenField[] {
+    const keyDirective = model.directives.find(d => d.name === 'key' && !d.arguments.name);
+    return keyDirective
+      ? (keyDirective.arguments.fields as string[]).slice(1).map(fieldName => model.fields.find(f => f.name === fieldName)!)
+      : [];
+  }
+
   protected determinePrimaryKeyFieldname(model: CodeGenModel): string {
     let primaryKeyFieldName = 'id';
     model.fields.forEach(field => {
@@ -664,12 +707,13 @@ export class AppSyncModelVisitor<
       let intermediateModel = this.generateIntermediateModel(
         value[0].model,
         value[1].model,
-        value[0].field,
-        value[1].field,
         graphqlName(toUpper(value[0].directive.arguments.relationName)),
       );
       const modelDirective = intermediateModel.directives.find(directive => directive.name === 'model');
       if (modelDirective) {
+        // Maps @primaryKey and @index of intermediate model to old @key
+        processPrimaryKey(intermediateModel);
+        processIndex(intermediateModel);
         this.ensureIdField(intermediateModel);
         this.addTimestampFields(intermediateModel, modelDirective);
         this.sortFields(intermediateModel);
@@ -682,7 +726,7 @@ export class AppSyncModelVisitor<
 
   protected processConnectionDirectivesV2(
     // TODO: Remove us when we have a fix to roll-forward.
-    shouldUseModelNameFieldInHasManyAndBelongsTo: boolean
+    shouldUseModelNameFieldInHasManyAndBelongsTo: boolean,
   ): void {
     this.processManyToManyDirectives();
 
@@ -726,15 +770,17 @@ export class AppSyncModelVisitor<
     Object.values(this.modelMap).forEach(model => {
       model.fields.forEach(field => {
         const connectionInfo = field.connectionInfo;
-        if (connectionInfo
-          && connectionInfo.kind !== CodeGenConnectionType.HAS_MANY
-          && connectionInfo.kind !== CodeGenConnectionType.HAS_ONE
-          && connectionInfo.targetName !== 'id') {
+        if (
+          connectionInfo &&
+          connectionInfo.kind !== CodeGenConnectionType.HAS_MANY &&
+          connectionInfo.kind !== CodeGenConnectionType.HAS_ONE &&
+          connectionInfo.targetName !== 'id'
+        ) {
           // Need to remove the field that is targetName
           removeFieldFromModel(model, connectionInfo.targetName);
         }
       });
-    })
+    });
   }
 
   protected processV2KeyDirectives(): void {