fix(api-graphql): incorrect list sk arg type (aws-amplify#13249)

* fix: incorrect sk arg type * bump bundle size * pr feedback
ashwinkumar6 · Apr 15, 2024 · f37faeb · f37faeb
1 parent 390c159
commit f37faeb
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 31 deletions.
diff --git a/packages/api-graphql/__tests__/internals/APIClient.test.ts b/packages/api-graphql/__tests__/internals/APIClient.test.ts
@@ -563,7 +563,7 @@ describe('generateGraphQLDocument()', () => {
 					modelOperation
 				);
 
-				expect(document.includes(expectedArgs)).toBe(true);
+				expect(document).toEqual(expect.stringContaining(expectedArgs))
 			}
 		);
 	});

diff --git a/packages/api-graphql/__tests__/internals/__snapshots__/generateClient.test.ts.snap b/packages/api-graphql/__tests__/internals/__snapshots__/generateClient.test.ts.snap
@@ -4092,12 +4092,12 @@ exports[`generateClient basic model operations can list() with sortDirection (AS
       "abortController": AbortController {},
       "options": {
         "body": {
-          "query": "query ($filter: ModelThingWithCustomPkFilterInput, $sortDirection: ModelSortDirection, $cpk_cluster_key: String, $cpk_sort_key: String, $limit: Int, $nextToken: String) {
+          "query": "query ($cpk_cluster_key: String, $sortDirection: ModelSortDirection, $cpk_sort_key: ModelStringKeyConditionInput, $filter: ModelThingWithCustomPkFilterInput, $limit: Int, $nextToken: String) {
   listThingWithCustomPks(
-    filter: $filter
-    sortDirection: $sortDirection
     cpk_cluster_key: $cpk_cluster_key
+    sortDirection: $sortDirection
     cpk_sort_key: $cpk_sort_key
+    filter: $filter
     limit: $limit
     nextToken: $nextToken
   ) {
@@ -4139,12 +4139,12 @@ exports[`generateClient basic model operations can list() with sortDirection (DE
       "abortController": AbortController {},
       "options": {
         "body": {
-          "query": "query ($filter: ModelThingWithCustomPkFilterInput, $sortDirection: ModelSortDirection, $cpk_cluster_key: String, $cpk_sort_key: String, $limit: Int, $nextToken: String) {
+          "query": "query ($cpk_cluster_key: String, $sortDirection: ModelSortDirection, $cpk_sort_key: ModelStringKeyConditionInput, $filter: ModelThingWithCustomPkFilterInput, $limit: Int, $nextToken: String) {
   listThingWithCustomPks(
-    filter: $filter
-    sortDirection: $sortDirection
     cpk_cluster_key: $cpk_cluster_key
+    sortDirection: $sortDirection
     cpk_sort_key: $cpk_sort_key
+    filter: $filter
     limit: $limit
     nextToken: $nextToken
   ) {

diff --git a/packages/api-graphql/src/internals/APIClient.ts b/packages/api-graphql/src/internals/APIClient.ts
@@ -721,6 +721,45 @@ export function generateGraphQLDocument(
 		selectionSet as ListArgs['selectionSet'],
 	);
 
+	// default PK args for get and list operations
+	// modified below for CPK
+	const getPkArgs = {
+		[primaryKeyFieldName]: `${fields[primaryKeyFieldName].type}!`,
+	};
+	const listPkArgs = {};
+
+	const generateSkArgs = (op: 'get' | 'list') => {
+		return sortKeyFieldNames.reduce(
+			(acc: Record<string, any>, fieldName: string) => {
+				const fieldType = fields[fieldName].type;
+
+				if (op === 'get') {
+					acc[fieldName] = `${fieldType}!`;
+				} else if (op === 'list') {
+					acc[fieldName] = `Model${fieldType}KeyConditionInput`;
+				}
+
+				return acc;
+			},
+			{},
+		);
+	};
+
+	if (isCustomPrimaryKey) {
+		Object.assign(getPkArgs, generateSkArgs('get'));
+
+		Object.assign(
+			listPkArgs,
+			{
+				// PK is only included in list query field args in the generated GQL
+				// when explicitly specifying PK with .identifier(['fieldName']) or @primaryKey in the schema definition
+				[primaryKeyFieldName]: `${fields[primaryKeyFieldName].type}`, // PK is always a nullable arg for list (no `!` after the type)
+				sortDirection: 'ModelSortDirection',
+			},
+			generateSkArgs('list'),
+		);
+	}
+
 	switch (modelOperation) {
 		case 'CREATE':
 		case 'UPDATE':
@@ -736,36 +775,15 @@ export function generateGraphQLDocument(
 		// TODO(Eslint): this this case clause correct without the break statement?
 		// eslint-disable-next-line no-fallthrough
 		case 'READ':
-			graphQLArguments ??
-				(graphQLArguments = isCustomPrimaryKey
-					? [primaryKeyFieldName, ...sortKeyFieldNames].reduce(
-							(acc: Record<string, any>, fieldName) => {
-								acc[fieldName] = `${fields[fieldName].type}!`;
-
-								return acc;
-							},
-							{},
-						)
-					: {
-							[primaryKeyFieldName]: `${fields[primaryKeyFieldName].type}!`,
-						});
+			graphQLArguments ?? (graphQLArguments = getPkArgs);
 			graphQLSelectionSet ?? (graphQLSelectionSet = selectionSetFields);
 		// TODO(Eslint): this this case clause correct without the break statement?
 		// eslint-disable-next-line no-fallthrough
 		case 'LIST':
 			graphQLArguments ??
 				(graphQLArguments = {
+					...listPkArgs,
 					filter: `Model${name}FilterInput`,
-					...(sortKeyFieldNames.length > 0
-						? [primaryKeyFieldName, ...sortKeyFieldNames].reduce(
-								(acc: Record<string, any>, fieldName) => {
-									acc[fieldName] = `${fields[fieldName].type}`;
-
-									return acc;
-								},
-								{ sortDirection: 'ModelSortDirection' },
-							)
-						: []),
 					limit: 'Int',
 					nextToken: 'String',
 				});

diff --git a/packages/aws-amplify/package.json b/packages/aws-amplify/package.json
@@ -335,7 +335,7 @@
 			"name": "[API] generateClient (AppSync)",
 			"path": "./dist/esm/api/index.mjs",
 			"import": "{ generateClient }",
-			"limit": "38.5 kB"
+			"limit": "39.0 kB"
 		},
 		{
 			"name": "[API] REST API handlers",