From 60dd62993241df56bb4f27afe79e1e7e1956d6ff Mon Sep 17 00:00:00 2001
From: Christopher Sundersingh <83315412+sundersc@users.noreply.github.com>
Date: Fri, 20 Sep 2024 14:50:18 -0700
Subject: [PATCH] fix(amplify-codegen): support multiple indexes on the same
 field in introspection schema (#879)

---
 .../src/__tests__/utils/process-index.test.ts | 72 +++++++++++++++++++
 .../src/utils/fieldUtils.ts                   |  3 +
 .../src/utils/process-index.ts                | 32 +++++----
 3 files changed, 94 insertions(+), 13 deletions(-)

diff --git a/packages/appsync-modelgen-plugin/src/__tests__/utils/process-index.test.ts b/packages/appsync-modelgen-plugin/src/__tests__/utils/process-index.test.ts
index a2c354ff..b81aa85e 100644
--- a/packages/appsync-modelgen-plugin/src/__tests__/utils/process-index.test.ts
+++ b/packages/appsync-modelgen-plugin/src/__tests__/utils/process-index.test.ts
@@ -68,6 +68,78 @@ describe('processIndex', () => {
     ]);
   });
 
+  it('support multiple @index directives on a field', () => {
+    const model: CodeGenModel = {
+      directives: [
+        {
+          name: 'model',
+          arguments: {},
+        },
+      ],
+      name: 'testModel',
+      type: 'model',
+      fields: [
+        {
+          type: 'field',
+          isList: false,
+          isNullable: true,
+          name: 'connectionField',
+          directives: [
+            {
+              name: 'index',
+              arguments: {
+                name: 'byItemAndSortField',
+                sortKeyFields: ['sortField'],
+              },
+            },
+            {
+              name: 'index',
+              arguments: {
+                name: 'byItemAndAnotherSortField',
+                sortKeyFields: ['anotherSortField'],
+              },
+            },
+            {
+              name: 'index',
+              arguments: {
+                name: 'byItemAndSomeOtherSortField',
+                sortKeyFields: ['someOtherSortField'],
+              },
+            },
+          ],
+        },
+      ],
+    };
+    processIndex(model);
+    expect(model.directives).toEqual([
+      {
+        name: 'model',
+        arguments: {},
+      },
+      {
+        name: 'key',
+        arguments: {
+          name: 'byItemAndSortField',
+          fields: ['connectionField', 'sortField'],
+        },
+      },
+      {
+        name: 'key',
+        arguments: {
+          name: 'byItemAndAnotherSortField',
+          fields: ['connectionField', 'anotherSortField'],
+        },
+      },
+      {
+        name: 'key',
+        arguments: {
+          name: 'byItemAndSomeOtherSortField',
+          fields: ['connectionField', 'someOtherSortField'],
+        },
+      },
+    ]);
+  });
+
   it('adds simple @index directives as model key attributes', () => {
     const model: CodeGenModel = {
       directives: [
diff --git a/packages/appsync-modelgen-plugin/src/utils/fieldUtils.ts b/packages/appsync-modelgen-plugin/src/utils/fieldUtils.ts
index 17232fe4..948d7be5 100644
--- a/packages/appsync-modelgen-plugin/src/utils/fieldUtils.ts
+++ b/packages/appsync-modelgen-plugin/src/utils/fieldUtils.ts
@@ -15,6 +15,9 @@ export function removeFieldFromModel(model: CodeGenModel, fieldName: string): vo
 export const getDirective = (fieldOrModel: CodeGenField | CodeGenModel) => (directiveName: string): CodeGenDirective | undefined =>
   fieldOrModel.directives.find(d => d.name === directiveName);
 
+export const getDirectives = (fieldOrModel: CodeGenField | CodeGenModel) => (directiveName: string): CodeGenDirective[] | undefined =>
+  fieldOrModel.directives.filter(d => d.name === directiveName);
+
 // Function matching to GraphQL transformer so that the auto-generated field
 export function toCamelCase(words: string[]): string {
   const formatted = words.map((w, i) => (i === 0 ? w.charAt(0).toLowerCase() + w.slice(1) : w.charAt(0).toUpperCase() + w.slice(1)));
diff --git a/packages/appsync-modelgen-plugin/src/utils/process-index.ts b/packages/appsync-modelgen-plugin/src/utils/process-index.ts
index dc00d085..c0492afe 100644
--- a/packages/appsync-modelgen-plugin/src/utils/process-index.ts
+++ b/packages/appsync-modelgen-plugin/src/utils/process-index.ts
@@ -1,5 +1,5 @@
 import { CodeGenDirective, CodeGenModel } from '../visitors/appsync-visitor';
-import { getDirective } from './fieldUtils';
+import { getDirectives } from './fieldUtils';
 import pluralize from 'pluralize';
 import { toLower, toUpper } from './stringUtils';
 
@@ -9,21 +9,27 @@ import { toLower, toUpper } from './stringUtils';
  */
 export const processIndex = (model: CodeGenModel) => {
   const indexMap = model.fields.reduce((acc, field) => {
-    const indexDirective = getDirective(field)('index');
-    if (!indexDirective) {
+    const indexDirectives = getDirectives(field)('index');
+    if (!indexDirectives) {
       return acc;
     }
-    return { ...acc, [field.name]: indexDirective };
-  }, {} as Record<string, CodeGenDirective>);
+    return { ...acc, [field.name]: indexDirectives };
+  }, {} as Record<string, CodeGenDirective[]>);
+
+  const keyList: CodeGenDirective[] = [];
+  Object.entries(indexMap).forEach(([fieldName, directives]) => {
+    directives.forEach(directive => {
+      keyList.push({
+        name: 'key',
+        arguments: {
+          name: directive.arguments.name ?? generateDefaultIndexName(model.name, [fieldName].concat((directive.arguments.sortKeyFields as string[]) ?? [])),
+          queryField: directive.arguments.queryField,
+          fields: [fieldName].concat((directive.arguments.sortKeyFields as string[]) ?? []),
+        },
+      });
+    });
+  });
 
-  const keyList: CodeGenDirective[] = Object.entries(indexMap).map(([fieldName, directive]) => ({
-    name: 'key',
-    arguments: {
-      name: directive.arguments.name ?? generateDefaultIndexName(model.name, [fieldName].concat((directive.arguments.sortKeyFields as string[]) ?? [])),
-      queryField: directive.arguments.queryField,
-      fields: [fieldName].concat((directive.arguments.sortKeyFields as string[]) ?? []),
-    },
-  }));
   const existingIndexNames = model.directives
     .filter(directive => directive.name === 'key' && !!directive.arguments.name)
     .map(directive => directive.arguments.name);