chore: refactor getprimarykeyfields method; add e2e test

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,17 @@ 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, throws meaningful error otherwise.
+ * @param model Codegen Model object
+ * @returns Array of primary and sort key codegen fields
+ */
 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);
+  if (primaryKeyField && primaryKeyField?.primaryKeyInfo) {
+    const { sortKeyFields } = primaryKeyField.primaryKeyInfo;
+    return [ primaryKeyField, ...sortKeyFields ];
+  }
+
+  throw new Error(`Unable to get the primary key component fields for model ${model?.name}`);
 }

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

@@ -1022,6 +1022,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 +1032,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);