fix(api-graphql): custom operation selection set on nested custom types

packages/api-graphql/__tests__/fixtures/modeled/amplifyconfiguration.ts

@@ -1367,6 +1367,10 @@ const amplifyConfig = {
 				name: 'Status',
 				values: ['NOT_STARTED', 'STARTED', 'DONE', 'CANCELED'],
 			},
+			ProductMetaStatus: {
+				name: 'ProductMetaStatus',
+				values: ['discontinued', 'in_production'],
+			},
 		},
 		nonModels: {
 			CommunityPostMetadata: {
@@ -1508,6 +1512,38 @@ const amplifyConfig = {
 					},
 				},
 			},
+			echoNestedCustomTypes: {
+				name: 'echoNestedCustomTypes',
+				isArray: false,
+				type: {
+					nonModel: 'ProductTrackingMeta',
+				},
+				isRequired: false,
+				arguments: {
+					argumentContent: {
+						name: 'input',
+						isArray: false,
+						type: 'String',
+						isRequired: true,
+					},
+				},
+			},
+			echoModelHasNestedTypes: {
+				name: 'echoModelHasNestedTypes',
+				isArray: false,
+				type: {
+					model: 'Product',
+				},
+				isRequired: false,
+				arguments: {
+					argumentContent: {
+						name: 'input',
+						isArray: false,
+						type: 'String',
+						isRequired: true,
+					},
+				},
+			},
 		},
 		mutations: {
 			likePost: {

packages/api-graphql/__tests__/fixtures/modeled/schema.ts

@@ -87,6 +87,26 @@ const schema = a.schema({
 			a.index('title'),
 			a.index('description').sortKeys(['viewCount']),
 		]),
+	Product: a
+		.model({
+			sku: a.string().required(),
+			factoryId: a.string().required(),
+			warehouseId: a.string().required(),
+			description: a.string(),
+			trackingMeta: a.customType({
+				productMeta: a.ref('ProductMeta'),
+				note: a.string(),
+			}),
+		})
+		.identifier(['sku', 'factoryId', 'warehouseId'])
+		.authorization([a.allow.public()]),
+	ProductMeta: a.customType({
+		releaseDate: a.date(),
+		status: a.enum(['in_production', 'discontinued']),
+		deepMeta: a.customType({
+			content: a.string(),
+		}),
+	}),
 
 	// #region Custom queries and mutations
 	EchoResult: a.customType({
@@ -112,7 +132,22 @@ const schema = a.schema({
 		.returns(a.string())
 		.function('echoFunction')
 		.authorization([a.allow.public()]),
-
+	echoNestedCustomTypes: a
+		.query()
+		.arguments({
+			input: a.string().required(),
+		})
+		.returns(a.ref('ProductTrackingMeta'))
+		.function('echoFunction')
+		.authorization([a.allow.public()]),
+	echoModelHasNestedTypes: a
+		.query()
+		.arguments({
+			input: a.string().required(),
+		})
+		.returns(a.ref('Product'))
+		.function('echoFunction')
+		.authorization([a.allow.public()]),
 	// custom mutation returning a non-model type
 	PostLikeResult: a.customType({
 		likes: a.integer().required(),

packages/api-graphql/__tests__/internals/__snapshots__/generateClient.test.ts.snap

@@ -4680,6 +4680,54 @@ exports[`generateClient custom operations can mutate with returnType of model (P
 ]
 `;
 
+exports[`generateClient custom operations can query with returnType of a model that has nested custom types 1`] = `
+[
+  [
+    "AmplifyClassV6",
+    {
+      "abortController": AbortController {},
+      "options": {
+        "body": {
+          "query": "query ($argumentContent: String!) {
+  echoModelHasNestedTypes(argumentContent: $argumentContent) {
+    sku
+    factoryId
+    warehouseId
+    description
+    trackingMeta {
+      productMeta {
+        releaseDate
+        status
+        deepMeta {
+          content
+        }
+      }
+      note
+    }
+    owner
+    createdAt
+    updatedAt
+  }
+}
+",
+          "variables": {
+            "input": "test input",
+          },
+        },
+        "headers": {
+          "Authorization": "amplify-config-auth-token",
+          "X-Api-Key": "FAKE-KEY",
+          "x-amz-user-agent": "aws-amplify/latest api/latest framework/latest",
+        },
+        "signingServiceInfo": undefined,
+        "withCredentials": undefined,
+      },
+      "url": "https://localhost/graphql",
+    },
+  ],
+]
+`;
+
 exports[`generateClient custom operations can query with returnType of customType 1`] = `
 [
   [
@@ -4712,6 +4760,45 @@ exports[`generateClient custom operations can query with returnType of customTyp
 ]
 `;
 
+exports[`generateClient custom operations can query with returnType of nested custom types 1`] = `
+[
+  [
+    "AmplifyClassV6",
+    {
+      "abortController": AbortController {},
+      "options": {
+        "body": {
+          "query": "query ($argumentContent: String!) {
+  echoNestedCustomTypes(argumentContent: $argumentContent) {
+    productMeta {
+      releaseDate
+      status
+      deepMeta {
+        content
+      }
+    }
+    note
+  }
+}
+",
+          "variables": {
+            "input": "test input",
+          },
+        },
+        "headers": {
+          "Authorization": "amplify-config-auth-token",
+          "X-Api-Key": "FAKE-KEY",
+          "x-amz-user-agent": "aws-amplify/latest api/latest framework/latest",
+        },
+        "signingServiceInfo": undefined,
+        "withCredentials": undefined,
+      },
+      "url": "https://localhost/graphql",
+    },
+  ],
+]
+`;
+
 exports[`generateClient custom operations can query with returnType of string 1`] = `
 [
   [

packages/api-graphql/__tests__/internals/generateClient.test.ts

@@ -162,7 +162,7 @@ describe('generateClient', () => {
 	});
 
 	describe('client `enums` property', () => {
-		const expectedEnumsProperties = ['Status'];
+		const expectedEnumsProperties = ['Status', 'ProductMetaStatus'];
 
 		it('generates `enums` property when Amplify.getConfig() returns valid GraphQL provider config', () => {
 			Amplify.configure(configFixture); // clear the resource config
@@ -5222,6 +5222,68 @@ describe('generateClient', () => {
 			});
 		});
 
+		test('can query with returnType of nested custom types', async () => {
+			const mockReturnData = {
+				note: 'test node',
+				productMeta: {
+					releaseDate: '2024-03-04',
+					status: 'in_production',
+					deepMeta: {
+						content: 'test deep meta content',
+					},
+				},
+			};
+			const spy = mockApiResponse({
+				data: {
+					echoNestedCustomTypes: mockReturnData,
+				},
+			});
+
+			const client = generateClient<Schema>({
+				amplify: Amplify,
+			});
+			const result = await client.queries.echoNestedCustomTypes({
+				input: 'test input',
+			});
+
+			expect(normalizePostGraphqlCalls(spy)).toMatchSnapshot();
+			expect(result?.data).toEqual(mockReturnData);
+		});
+
+		test('can query with returnType of a model that has nested custom types', async () => {
+			const mockReturnData = {
+				sku: 'sku',
+				factoryId: 'factoryId',
+				warehouseId: 'warehouseId',
+				description: 'description',
+				trackingMeta: {
+					productMeta: {
+						releaseDate: '2024-03-04',
+						status: 'discontinued',
+						deepMeta: {
+							content: 'test content',
+						},
+					},
+					note: 'test note',
+				},
+			};
+			const spy = mockApiResponse({
+				data: {
+					echoNestedCustomTypes: mockReturnData,
+				},
+			});
+
+			const client = generateClient<Schema>({
+				amplify: Amplify,
+			});
+			const result = await client.queries.echoModelHasNestedTypes({
+				input: 'test input',
+			});
+
+			expect(normalizePostGraphqlCalls(spy)).toMatchSnapshot();
+			expect(result?.data).toEqual(mockReturnData);
+		});
+
 		test('can query with returnType of string', async () => {
 			const spy = mockApiResponse({
 				data: {

packages/api-graphql/__tests__/internals/utils/generateEnumsProperty.test.ts

@@ -1,5 +1,5 @@
 import { GraphQLProviderConfig } from '@aws-amplify/core/internals/utils';
-import { generateEnumsProperty } from '../../../src/internals/utils/generateEnumsProperty';
+import { generateEnumsProperty } from '../../../src/internals/utils/clientProperties/generateEnumsProperty';
 
 describe('generateEnumsProperty()', () => {
 	it('returns an empty object when there is no valid `modelIntrospection`', () => {

packages/api-graphql/package.json

@@ -71,7 +71,7 @@
 	},
 	"homepage": "https://aws-amplify.github.io/",
 	"devDependencies": {
-		"@aws-amplify/data-schema": "^0.13.6",
+		"@aws-amplify/data-schema": "^0.13.9",
 		"@rollup/plugin-typescript": "11.1.5",
 		"rollup": "^4.9.6",
 		"typescript": "5.0.2"
@@ -86,7 +86,7 @@
 	"dependencies": {
 		"@aws-amplify/api-rest": "4.0.19",
 		"@aws-amplify/core": "6.0.19",
-		"@aws-amplify/data-schema-types": "^0.7.5",
+		"@aws-amplify/data-schema-types": "^0.7.6",
 		"@aws-sdk/types": "3.387.0",
 		"graphql": "15.8.0",
 		"rxjs": "^7.8.1",

packages/api-graphql/src/internals/APIClient.ts

@@ -314,7 +314,7 @@ type OperationPrefix =
 
 const SELECTION_SET_WILDCARD = '*';
 
-function defaultSelectionSetForNonModelWithIR(
+export function defaultSelectionSetForNonModelWithIR(
 	nonModelDefinition: SchemaNonModel,
 	modelIntrospection: ModelIntrospectionSchema,
 ): Record<string, unknown> {

packages/api-graphql/src/internals/generateClient.ts

@@ -8,9 +8,9 @@ import {
 	CustomQueries,
 	CustomMutations,
 } from '@aws-amplify/data-schema-types';
-import { generateEnumsProperty } from './utils/generateEnumsProperty';
-import { generateModelsProperty } from './utils/generateModelsProperty';
-import { isApiGraphQLConfig } from './utils/isApiGraphQLProviderConfig';
+import { generateEnumsProperty } from './utils/clientProperties/generateEnumsProperty';
+import { generateModelsProperty } from './utils/clientProperties/generateModelsProperty';
+import { isApiGraphQLConfig } from './utils/runtimeTypeGuards/isApiGraphQLProviderConfig';
 import {
 	generateCustomQueriesProperty,
 	generateCustomMutationsProperty,
@@ -23,7 +23,7 @@ import {
 	__headers,
 } from '../types';
 import { ClientGenerationParams } from './types';
-import { isConfigureEventWithResourceConfig } from './utils/isConfigureEventWithResourceConfig';
+import { isConfigureEventWithResourceConfig } from './utils/runtimeTypeGuards/isConfigureEventWithResourceConfig';
 
 /**
  * @private

packages/api-graphql/src/internals/operations/custom.ts

@@ -6,6 +6,9 @@ import {
 	flattenItems,
 	authModeParams,
 	getCustomHeaders,
+	generateSelectionSet,
+	selectionSetIRToString,
+	defaultSelectionSetForNonModelWithIR,
 } from '../APIClient';
 import {
 	AuthModeParams,
@@ -197,27 +200,6 @@ function innerArguments(operation: CustomOperation): string {
 	return args.length > 0 ? `(${args})` : '';
 }
 
-/**
- * Accepts a Model or NonModel from generated introspection schema and
- * generates the selection set string.
- *
- * @param nonModel Model or nonModel object from introspection schema.
- * @returns String selection set.
- */
-function modelishTypeSelectionSet(
-	nonModel: SchemaModel | SchemaNonModel,
-): string {
-	return Object.values(nonModel.fields)
-		.filter(
-			field =>
-				hasStringField(field, 'type') ||
-				hasStringField(field.type, 'nonModel') ||
-				hasStringField(field.type, 'enum'),
-		)
-		.map(field => field.name)
-		.join(' ');
-}
-
 /**
  * Generates the selection set string for a custom operation. This is slightly
  * different than the selection set generation for models. If the custom op returns
@@ -251,11 +233,12 @@ function operationSelectionSet(
 	) {
 		return '';
 	} else if (hasStringField(operation.type, 'nonModel')) {
-		const model = modelIntrospection.nonModels[operation.type.nonModel];
-		return `{${modelishTypeSelectionSet(model)}}`;
+		const nonModel = modelIntrospection.nonModels[operation.type.nonModel];
+		return `{${selectionSetIRToString(
+			defaultSelectionSetForNonModelWithIR(nonModel, modelIntrospection),
+		)}}`;
 	} else if (hasStringField(operation.type, 'model')) {
-		const nonModel = modelIntrospection.models[operation.type.model];
-		return `{${modelishTypeSelectionSet(nonModel)}}`;
+		return `{${generateSelectionSet(modelIntrospection, operation.type.model)}}`;
 	} else {
 		return '';
 	}