Merge pull request #573 from aws-amplify/add-e2e-test-codegen

Refactor the `getPrimaryKeyComponentFields` method

Author: phani-srikar

packages/amplify-codegen-e2e-tests/schemas/modelgen/schema_with_connected_pk.graphql

@@ -0,0 +1,10 @@
+type Post @model {
+  postId: ID! @primaryKey
+  node: PostNode! @belongsTo(fields: ["postId"])
+  title: String!
+}
+
+type PostNode @model {
+  id: ID!
+  post: Post! @hasOne
+}

packages/amplify-codegen-e2e-tests/src/__tests__/model-introspection-codegen.test.ts

@@ -28,6 +28,7 @@ describe('Model Introspection Codegen test', () => {
       // Model introspection is generated at correct location
       expect(isNotEmptyDir(join(projectRoot, outputDir))).toBe(true);
     });
+
     it('should throw error when the output directory is not defined in the command', async () => {
       // init project and add API category
       await initJSProjectWithProfile(projectRoot);
@@ -36,6 +37,7 @@ describe('Model Introspection Codegen test', () => {
       //generate introspection schema
       await expect(generateModelIntrospection(projectRoot)).rejects.toThrowError();
     });
+
     it('should throw error if the GraphQL schema is invalid', async () => {
       const invalidSchema = 'modelgen/model_gen_schema_with_errors.graphql';
       // init project and add API category
@@ -46,5 +48,20 @@ describe('Model Introspection Codegen test', () => {
       //generate introspection schema
       await expect(generateModelIntrospection(projectRoot, { outputDir })).rejects.toThrowError();
     });
+
+    it(`should handle a schema with connected PK`, async () => {
+      const schemaName = 'modelgen/schema_with_connected_pk.graphql';
+
+      // init project and add API category
+      await initJSProjectWithProfile(projectRoot);
+      await addApiWithoutSchema(projectRoot, { apiName });
+      await updateApiSchema(projectRoot, apiName, schemaName);
+
+      const outputDir = 'output';
+      //generate introspection schema
+      await expect(generateModelIntrospection(projectRoot, { outputDir })).resolves.not.toThrow();
+      // Model introspection is generated at correct location
+      expect(isNotEmptyDir(join(projectRoot, outputDir))).toBe(true);
+    });
 });
 

packages/appsync-modelgen-plugin/src/utils/fieldUtils.ts

@@ -25,8 +25,19 @@ export function getOtherSideBelongsToField(type: string, otherSideModel: CodeGen
   return otherSideModel.fields.filter(f => f.type === type).find(f => f.directives.find(d => d.name === TransformerV2DirectiveName.BELONGS_TO));
 }
 
+/**
+ * Given a model, it returns the primary and sort key fields if present, an empty list otherwise.
+ * @param model Codegen Model object
+ * @returns Array of primary and sort key codegen fields if present or an empty list
+ */
 export function getModelPrimaryKeyComponentFields(model: CodeGenModel): CodeGenField[] {
-  const primaryKeyField = model.fields.find(field => field.primaryKeyInfo)!;
-  const { sortKeyFields } = primaryKeyField.primaryKeyInfo!;
-  return [ primaryKeyField, ...sortKeyFields ];
+  const primaryKeyField = model.fields.find(field => field.primaryKeyInfo);
+  const keyFields: CodeGenField[] = [];
+  if (primaryKeyField) {
+    keyFields.push(primaryKeyField);
+    if ( primaryKeyField?.primaryKeyInfo?.sortKeyFields ) {
+      keyFields.push(...primaryKeyField.primaryKeyInfo.sortKeyFields);
+    };
+  }
+  return keyFields;
 }

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

@@ -941,15 +941,17 @@ export class AppSyncModelVisitor<
           } else if (connectionInfo.kind === CodeGenConnectionType.HAS_ONE) {
             if (isCustomPKEnabled) {
               const connectedModelFields = getModelPrimaryKeyComponentFields(connectionInfo.connectedModel);
-              connectionInfo.targetNames.forEach((target, index) => {
-                addFieldToModel(model, {
-                  name: target,
-                  directives: [],
-                  type: connectedModelFields[index].type,
-                  isList: false,
-                  isNullable: field.isNullable,
+              if (connectedModelFields?.length > 0) {
+                connectionInfo.targetNames.forEach((target, index) => {
+                  addFieldToModel(model, {
+                    name: target,
+                    directives: [],
+                    type: connectedModelFields[index].type,
+                    isList: false,
+                    isNullable: field.isNullable,
+                  });
                 });
-              });
+              }
             } else {
               addFieldToModel(model, {
                 name: connectionInfo.targetName,
@@ -962,15 +964,17 @@ export class AppSyncModelVisitor<
           } else if (connectionInfo.kind === CodeGenConnectionType.BELONGS_TO) {
             if (isCustomPKEnabled) {
               const connectedModelFields = getModelPrimaryKeyComponentFields(connectionInfo.connectedModel);
-              connectionInfo.targetNames.forEach((target, index) => {
-                addFieldToModel(model, {
-                  name: target,
-                  directives: [],
-                  type: connectedModelFields[index].type,
-                  isList: false,
-                  isNullable: field.isNullable,
+              if (connectedModelFields?.length > 0) {
+                connectionInfo.targetNames.forEach((target, index) => {
+                  addFieldToModel(model, {
+                    name: target,
+                    directives: [],
+                    type: connectedModelFields[index].type,
+                    isList: false,
+                    isNullable: field.isNullable,
+                  });
                 });
-              });
+              }
             } else {
               addFieldToModel(model, {
                 name: connectionInfo.targetName,
@@ -1022,6 +1026,8 @@ export class AppSyncModelVisitor<
       }
     } else {
       Object.values(this.modelMap).forEach(model => {
+        const primaryKeyFields = getModelPrimaryKeyComponentFields(model);
+        const primaryKeyName = (primaryKeyFields?.length > 0) ? this.getFieldName(primaryKeyFields[0]) : undefined;
         model.fields.forEach(field => {
           const connectionInfo = field.connectionInfo;
           if (
@@ -1030,7 +1036,7 @@ export class AppSyncModelVisitor<
             connectionInfo.kind !== CodeGenConnectionType.HAS_ONE &&
             connectionInfo.targetName !== 'id' &&
             !(this.config.target === 'introspection' &&
-              this.getFieldName(getModelPrimaryKeyComponentFields(model)[0]) === connectionInfo.targetName)
+              primaryKeyName && primaryKeyName === connectionInfo.targetName)
           ) {
             // Need to remove the field that is targetName
             removeFieldFromModel(model, connectionInfo.targetName);