From bb5fd125feb647eea0842ff75b16d7ed62329888 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=A9my=20Jourdin?= <jeremy.jourdin@gmail.com>
Date: Sat, 24 Feb 2024 09:51:36 +0100
Subject: [PATCH 1/2] Fix subscription result & improve types

---
 .../src/internals/operations/get.ts           |  5 ++--
 .../src/internals/operations/subscription.ts  | 27 +++++++++++--------
 packages/api-graphql/src/types/index.ts       | 12 ++++++---
 3 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/packages/api-graphql/src/internals/operations/get.ts b/packages/api-graphql/src/internals/operations/get.ts
index fc727f22c21..f0a87a88759 100644
--- a/packages/api-graphql/src/internals/operations/get.ts
+++ b/packages/api-graphql/src/internals/operations/get.ts
@@ -13,9 +13,10 @@ import {
 import {
 	AuthModeParams,
 	ClientWithModels,
+	ConfigurableHeadersParams,
+	FieldSelectionParams,
 	GraphQLOptionsV6,
 	GraphQLResult,
-	ListArgs,
 	QueryArgs,
 	V6Client,
 	V6ClientSSRRequest,
@@ -60,7 +61,7 @@ async function _get(
 	modelIntrospection: ModelIntrospectionSchema,
 	model: SchemaModel,
 	arg: QueryArgs,
-	options: AuthModeParams & ListArgs,
+	options: AuthModeParams & ConfigurableHeadersParams & FieldSelectionParams,
 	operation: ModelOperation,
 	context?: AmplifyServer.ContextSpec,
 ) {
diff --git a/packages/api-graphql/src/internals/operations/subscription.ts b/packages/api-graphql/src/internals/operations/subscription.ts
index 03947238fb6..50d43651663 100644
--- a/packages/api-graphql/src/internals/operations/subscription.ts
+++ b/packages/api-graphql/src/internals/operations/subscription.ts
@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import { map } from 'rxjs';
-import { V6Client, GraphqlSubscriptionResult } from '../../types';
+import { V6Client, GraphqlSubscriptionResult, ListArgs, AuthModeParams } from '../../types';
 import {
 	initializeModel,
 	generateGraphQLDocument,
@@ -24,7 +24,7 @@ export function subscriptionFactory(
 ) {
 	const { name } = model as any;
 
-	const subscription = (args?: any) => {
+	const subscription = (args?: ListArgs & AuthModeParams) => {
 		const query = generateGraphQLDocument(
 			modelIntrospection,
 			name,
@@ -56,15 +56,20 @@ export function subscriptionFactory(
 			map(value => {
 				const [key] = Object.keys(value.data);
 				const data = (value.data as any)[key];
-				const [initialized] = initializeModel(
-					client as V6Client<Record<string, any>>,
-					name,
-					[data],
-					modelIntrospection,
-					auth.authMode,
-					auth.authToken,
-				);
-				return initialized;
+
+				if (args?.selectionSet) {
+					return data;
+				} else {
+					const [initialized] = initializeModel(
+						client as V6Client<Record<string, any>>,
+						name,
+						[data],
+						modelIntrospection,
+						auth.authMode,
+						auth.authToken,
+					);
+					return initialized;
+				}
 			}),
 		);
 	};
diff --git a/packages/api-graphql/src/types/index.ts b/packages/api-graphql/src/types/index.ts
index 66645e7ab51..6d5a16ef6c0 100644
--- a/packages/api-graphql/src/types/index.ts
+++ b/packages/api-graphql/src/types/index.ts
@@ -447,11 +447,17 @@ export type ServerClientGenerationParams = {
 
 export type QueryArgs = Record<string, unknown>;
 
-export type ListArgs = {
+export type FieldSelectionParams = {
 	selectionSet?: string[];
-	filter?: {};
+}
+
+export type ConfigurableHeadersParams = {
 	headers?: CustomHeaders;
-};
+}
+
+export type ListArgs = {
+	filter?: {};
+} & ConfigurableHeadersParams & FieldSelectionParams;
 
 export type AuthModeParams = {
 	authMode?: GraphQLAuthMode;

From bea1cef6139f7266b3a60a4f9cafe31c38d3bba2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=A9my=20Jourdin?= <jeremy.jourdin@gmail.com>
Date: Thu, 7 Mar 2024 09:58:45 +0100
Subject: [PATCH 2/2] Fix lint

---
 packages/api-graphql/src/types/index.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/api-graphql/src/types/index.ts b/packages/api-graphql/src/types/index.ts
index 14c70c66226..8daafb4fd71 100644
--- a/packages/api-graphql/src/types/index.ts
+++ b/packages/api-graphql/src/types/index.ts
@@ -459,11 +459,11 @@ export type QueryArgs = Record<string, unknown>;
 
 export type FieldSelectionParams = {
 	selectionSet?: string[];
-}
+};
 
 export type ConfigurableHeadersParams = {
 	headers?: CustomHeaders;
-}
+};
 
 export type ListArgs = {
 	filter?: {};