feat(api-graphql): custom subscriptions

.vscode/settings.json

@@ -1,7 +1,6 @@
 {
 	"editor.detectIndentation": false,
 	"editor.insertSpaces": false,
-	"editor.tabSize": 4,
 	"prettier.requireConfig": true,
 	"typescript.tsdk": "node_modules/typescript/lib",
 	"formattingToggle.affects": [

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

@@ -1580,6 +1580,32 @@ const amplifyConfig = {
 				},
 			},
 		},
+		subscriptions: {
+			onPostLiked: {
+				name: 'onPostLiked',
+				isArray: false,
+				type: {
+					model: 'Post',
+				},
+				isRequired: false,
+			},
+			onPostUpdated: {
+				name: 'onPostUpdated',
+				isArray: false,
+				type: {
+					model: 'Post',
+				},
+				isRequired: false,
+				arguments: {
+					postId: {
+						name: 'postId',
+						isArray: false,
+						type: 'String',
+						isRequired: false,
+					},
+				},
+			},
+		},
 	},
 };
 export default amplifyConfig;

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

@@ -120,7 +120,7 @@ const schema = a.schema({
 			argumentContent: a.string().required(),
 		})
 		.returns(a.ref('EchoResult'))
-		.function('echoFunction')
+		.handler(a.handler.function('echoFunction'))
 		.authorization([a.allow.public()]),
 
 	// custom query returning a primitive type
@@ -130,23 +130,23 @@ const schema = a.schema({
 			inputString: a.string().required(),
 		})
 		.returns(a.string())
-		.function('echoFunction')
+		.handler(a.handler.function('echoFunction'))
 		.authorization([a.allow.public()]),
 	echoNestedCustomTypes: a
 		.query()
 		.arguments({
 			input: a.string().required(),
 		})
 		.returns(a.ref('ProductTrackingMeta'))
-		.function('echoFunction')
+		.handler(a.handler.function('echoFunction'))
 		.authorization([a.allow.public()]),
 	echoModelHasNestedTypes: a
 		.query()
 		.arguments({
 			input: a.string().required(),
 		})
 		.returns(a.ref('Product'))
-		.function('echoFunction')
+		.handler(a.handler.function('echoFunction'))
 		.authorization([a.allow.public()]),
 	// custom mutation returning a non-model type
 	PostLikeResult: a.customType({
@@ -158,7 +158,7 @@ const schema = a.schema({
 			postId: a.id().required(),
 		})
 		.returns(a.ref('PostLikeResult'))
-		.function('echoFunction')
+		.handler(a.handler.function('echoFunction'))
 		.authorization([a.allow.private()]),
 
 	// custom mutation returning a model type
@@ -182,9 +182,21 @@ const schema = a.schema({
 			postId: a.id().required(),
 		})
 		.returns(a.ref('Post'))
-		.function('echoFunction')
+		.handler(a.handler.function('echoFunction'))
 		.authorization([a.allow.private()]),
 
+	onPostLiked: a
+		.subscription()
+		.for(a.ref('likePostReturnPost'))
+		.returns(a.ref('Post'))
+		.handler(a.handler.custom({ entry: './jsResolver_base.js' })),
+
+	onPostUpdated: a
+		.subscription()
+		.for(a.ref('Post').mutations(['update']))
+		.arguments({ postId: a.string() })
+		.returns(a.ref('Post'))
+		.handler(a.handler.custom({ entry: './jsResolver_base.js' })),
 	//#endregion
 });
 

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

@@ -10,7 +10,6 @@ import {
 	expectSubWithlibraryConfigHeaders,
 } from '../utils/expects';
 import { Observable, from } from 'rxjs';
-import * as internals from '../../src/internals';
 
 const serverManagedFields = {
 	id: 'some-id',
@@ -1996,7 +1995,9 @@ describe('generateClient', () => {
 			const spy = jest.fn(() => from([graphqlMessage]));
 			(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
 
-			client.models.Note.onCreate().subscribe({
+			const onC = client.models.Note.onCreate();
+
+			onC.subscribe({
 				next(value) {
 					expect(spy).toHaveBeenCalledWith(
 						expect.objectContaining({
@@ -5181,6 +5182,8 @@ describe('generateClient', () => {
 		beforeEach(() => {
 			jest.clearAllMocks();
 			jest.resetAllMocks();
+			jest.restoreAllMocks();
+
 			Amplify.configure(configFixture as any);
 
 			jest
@@ -5392,6 +5395,121 @@ describe('generateClient', () => {
 			expect(comments[0]).toEqual(expect.objectContaining(listCommentItem));
 		});
 
+		test('can subscribe to custom subscription', done => {
+			const postToSend = {
+				__typename: 'Post',
+				...serverManagedFields,
+				content: 'a lovely post',
+			};
+
+			const graphqlMessage = {
+				data: {
+					onPostLiked: postToSend,
+				},
+			};
+
+			const client = generateClient<Schema>({ amplify: Amplify });
+
+			const spy = jest.fn(() => from([graphqlMessage]));
+			(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
+
+			const spyGql = jest.spyOn(raw.GraphQLAPI as any, 'graphql');
+
+			const expectedOperation = 'subscription';
+			const expectedFieldAndSelectionSet =
+				'onPostLiked {id content owner createdAt updatedAt}';
+
+			const sub = client.subscriptions.onPostLiked().subscribe({
+				next(value) {
+					expect(value).toEqual(expect.objectContaining(postToSend));
+
+					expect(spyGql).toHaveBeenCalledWith(
+						expect.anything(),
+						expect.objectContaining({
+							query: expect.stringContaining(expectedOperation),
+							variables: {},
+						}),
+						expect.anything(),
+					);
+					expect(spyGql).toHaveBeenCalledWith(
+						expect.anything(),
+						expect.objectContaining({
+							query: expect.stringContaining(expectedFieldAndSelectionSet),
+							variables: {},
+						}),
+						expect.anything(),
+					);
+					done();
+				},
+				error(error) {
+					expect(error).toBeUndefined();
+					done('bad news!');
+				},
+			});
+
+			sub.unsubscribe();
+		});
+
+		test('can subscribe to custom subscription with args', done => {
+			const postToSend = {
+				__typename: 'Post',
+				...serverManagedFields,
+				postId: 'abc123',
+				content: 'a lovely post',
+			};
+
+			const graphqlMessage = {
+				data: {
+					onPostLiked: postToSend,
+				},
+			};
+
+			const client = generateClient<Schema>({ amplify: Amplify });
+
+			const spy = jest.fn(() => from([graphqlMessage]));
+			(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
+
+			const spyGql = jest.spyOn(raw.GraphQLAPI as any, 'graphql');
+
+			const expectedOperation = 'subscription($postId: String)';
+			const expectedFieldAndSelectionSet =
+				'onPostUpdated(postId: $postId) {id content owner createdAt updatedAt}';
+
+			const sub = client.subscriptions
+				.onPostUpdated({ postId: 'abc123' })
+				.subscribe({
+					next(value) {
+						expect(value).toEqual(expect.objectContaining(postToSend));
+
+						expect(spyGql).toHaveBeenCalledWith(
+							expect.anything(),
+							expect.objectContaining({
+								query: expect.stringContaining(expectedOperation),
+
+								variables: { postId: 'abc123' },
+							}),
+							expect.anything(),
+						);
+						expect(spyGql).toHaveBeenCalledWith(
+							expect.anything(),
+							expect.objectContaining({
+								query: expect.stringContaining(expectedFieldAndSelectionSet),
+
+								variables: { postId: 'abc123' },
+							}),
+							expect.anything(),
+						);
+						done();
+					},
+					error(error) {
+						expect(error).toBeUndefined();
+						done('bad news!');
+					},
+				});
+
+			sub.unsubscribe();
+		});
+
 		test('includes client level headers', async () => {
 			const spy = mockApiResponse({
 				data: {

packages/api-graphql/__tests__/internals/server/generateClientWithAmplifyInstance.test.ts

@@ -425,7 +425,7 @@ describe('server generateClient', () => {
 				return result;
 			});
 
-			const mockContextSpec = {};
+			const mockContextSpec = { token: { value: Symbol('test') } };
 
 			const result = await client.queries.echo(mockContextSpec, {
 				argumentContent: 'echo argumentContent value',

packages/api-graphql/package.json

@@ -72,7 +72,7 @@
 	},
 	"homepage": "https://aws-amplify.github.io/",
 	"devDependencies": {
-		"@aws-amplify/data-schema": "^0.13.9",
+		"@aws-amplify/data-schema": "^0.13.19",
 		"@rollup/plugin-typescript": "11.1.5",
 		"rollup": "^4.9.6",
 		"typescript": "5.0.2"
@@ -87,7 +87,7 @@
 	"dependencies": {
 		"@aws-amplify/api-rest": "4.0.21",
 		"@aws-amplify/core": "6.0.21",
-		"@aws-amplify/data-schema-types": "^0.7.6",
+		"@aws-amplify/data-schema-types": "^0.7.10",
 		"@aws-sdk/types": "3.387.0",
 		"graphql": "15.8.0",
 		"rxjs": "^7.8.1",

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

@@ -4,6 +4,7 @@ import { Hub } from '@aws-amplify/core';
 import {
 	CustomMutations,
 	CustomQueries,
+	CustomSubscriptions,
 	EnumTypes,
 	ModelTypes,
 } from '@aws-amplify/data-schema-types';
@@ -23,6 +24,7 @@ import { isApiGraphQLConfig } from './utils/runtimeTypeGuards/isApiGraphQLProvid
 import {
 	generateCustomMutationsProperty,
 	generateCustomQueriesProperty,
+	generateCustomSubscriptionsProperty,
 } from './generateCustomOperationsProperty';
 import { ClientGenerationParams } from './types';
 import { isConfigureEventWithResourceConfig } from './utils/runtimeTypeGuards/isConfigureEventWithResourceConfig';
@@ -51,6 +53,7 @@ export function generateClient<T extends Record<any, any> = never>(
 		enums: emptyProperty as EnumTypes<never>,
 		queries: emptyProperty as CustomQueries<never>,
 		mutations: emptyProperty as CustomMutations<never>,
+		subscriptions: emptyProperty as CustomSubscriptions<never>,
 	} as any;
 
 	const apiGraphqlConfig = params.amplify.getConfig().API?.GraphQL;
@@ -63,6 +66,10 @@ export function generateClient<T extends Record<any, any> = never>(
 			client,
 			apiGraphqlConfig,
 		);
+		client.subscriptions = generateCustomSubscriptionsProperty(
+			client,
+			apiGraphqlConfig,
+		);
 	} else {
 		// This happens when the `Amplify.configure()` call gets evaluated after the `generateClient()` call.
 		//
@@ -101,6 +108,10 @@ const generateModelsPropertyOnAmplifyConfigure = (clientRef: any) => {
 				clientRef,
 				apiGraphQLConfig,
 			);
+			clientRef.subscriptions = generateCustomSubscriptionsProperty(
+				clientRef,
+				apiGraphQLConfig,
+			);
 		}
 	});
 };

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

@@ -1,6 +1,10 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-import { CustomMutations, CustomQueries } from '@aws-amplify/data-schema-types';
+import {
+	CustomMutations,
+	CustomQueries,
+	CustomSubscriptions,
+} from '@aws-amplify/data-schema-types';
 import {
 	GraphQLProviderConfig,
 	ModelIntrospectionSchema,
@@ -10,16 +14,21 @@ import { V6Client, __amplify } from '../types';
 
 import { customOpFactory } from './operations/custom';
 
-type OpTypes = 'queries' | 'mutations';
+type OpTypes = 'queries' | 'mutations' | 'subscriptions';
 
 type CustomOpsProperty<
 	T extends Record<any, any>,
 	OpType extends OpTypes,
-> = OpType extends 'queries' ? CustomQueries<T> : CustomMutations<T>;
+> = OpType extends 'queries'
+	? CustomQueries<T>
+	: OpType extends 'mutations'
+		? CustomMutations<T>
+		: CustomSubscriptions<T>;
 
 const operationTypeMap = {
 	queries: 'query',
 	mutations: 'mutation',
+	subscriptions: 'subscription',
 } as const;
 
 export function generateCustomOperationsProperty<
@@ -88,3 +97,14 @@ export function generateCustomQueriesProperty<T extends Record<any, any>>(
 		'queries',
 	);
 }
+
+export function generateCustomSubscriptionsProperty<T extends Record<any, any>>(
+	client: V6Client<Record<string, any>>,
+	config: GraphQLProviderConfig['GraphQL'],
+) {
+	return generateCustomOperationsProperty<T, 'subscriptions'>(
+		client,
+		config,
+		'subscriptions',
+	);
+}