fix(datastore): Various integration bug fixes

packages/amplify_datastore/example/ios/unit_tests/GraphQLResponse+DecodeTests.swift

@@ -116,7 +116,7 @@ class GraphQLResponseDecodeTests: XCTestCase {
             "a": ["b"]
         ]
         let expectedJson = "{\"a\":[\"b\"]}"
-        let result: Result<String, APIError> = GraphQLResponse<String>.decodeDataPayload(json, modelName: nil)
+        let result: Result<String, APIError> = GraphQLResponse<String>.decodeDataPayload(json, modelName: "Post")
         XCTAssertNoThrow(try result.get())
         XCTAssertEqual(expectedJson, try result.get())
     }
@@ -137,7 +137,7 @@ class GraphQLResponseDecodeTests: XCTestCase {
             "author": "authorId",
         ], modelName: "Post"))
 
-        let result: Result<AnyModel, APIError> = GraphQLResponse<AnyModel>.decodeDataPayload(json, modelName: nil)
+        let result: Result<AnyModel, APIError> = GraphQLResponse<AnyModel>.decodeDataPayload(json, modelName: "Post")
         XCTAssertNoThrow(try result.get())
         XCTAssertEqual(expectedModel.modelName, try result.get().modelName)
         XCTAssertEqual(expectedModel.id, try result.get().id)
@@ -209,14 +209,16 @@ class GraphQLResponseDecodeTests: XCTestCase {
     func testFromAppSyncResponse_withOnlyData_decodePayload() {
         SchemaData.modelSchemaRegistry.registerModels(registry: ModelRegistry.self)
         let json: JSONValue = [
-            "onCreatePost": [
-                "__typename": "Post",
-                "id": "123",
-                "title": "test",
-                "author": "authorId",
-                "_version": 1,
-               "_deleted": nil,
-               "_lastChangedAt": 1707773705221
+            "data": [
+                "onCreatePost": [
+                    "__typename": "Post",
+                    "id": "123",
+                    "title": "test",
+                    "author": "authorId",
+                    "_version": 1,
+                   "_deleted": nil,
+                   "_lastChangedAt": 1707773705221
+                ]
             ]
         ]
 
@@ -227,7 +229,11 @@ class GraphQLResponseDecodeTests: XCTestCase {
             "author": "authorId",
         ], modelName: "Post"))
 
-        let result: Result<GraphQLResponse<MutationSync<AnyModel>>, APIError> = .fromAppSyncResponse(json: json, decodePath: "onCreatePost", modelName: nil)
+        let result: Result<GraphQLResponse<MutationSync<AnyModel>>, APIError> = .fromAppSyncResponse(
+            json: json,
+            decodePath: "onCreatePost",
+            modelName: "Post"
+        )
         XCTAssertNoThrow(try result.get())
         let mutationSync = try! result.get()
         XCTAssertNoThrow(try mutationSync.get())
@@ -244,7 +250,11 @@ class GraphQLResponseDecodeTests: XCTestCase {
             ]
         ]
 
-        let result: Result<GraphQLResponse<MutationSync<AnyModel>>, APIError> = .fromAppSyncResponse(json: json, decodePath: "onCreatePost", modelName: nil)
+        let result: Result<GraphQLResponse<MutationSync<AnyModel>>, APIError> = .fromAppSyncResponse(
+            json: json,
+            decodePath: "onCreatePost",
+            modelName: "Post"
+        )
         XCTAssertNoThrow(try result.get())
         let mutationSync = try! result.get()
         XCTAssertThrowsError(try mutationSync.get()) { error in
@@ -261,14 +271,16 @@ class GraphQLResponseDecodeTests: XCTestCase {
     func testFromAppSyncResponse_withDataAndErrors_decodePayload() {
         SchemaData.modelSchemaRegistry.registerModels(registry: ModelRegistry.self)
         let json: JSONValue = [
-            "onCreatePost": [
-                "__typename": "Post",
-                "id": "123",
-                "title": "test",
-                "author": "authorId",
-                "_version": 1,
-               "_deleted": nil,
-               "_lastChangedAt": 1707773705221
+            "data": [
+                "onCreatePost": [
+                    "__typename": "Post",
+                    "id": "123",
+                    "title": "test",
+                    "author": "authorId",
+                    "_version": 1,
+                   "_deleted": nil,
+                   "_lastChangedAt": 1707773705221
+                ],
             ],
             "errors": [
                 ["message": "error1"],
@@ -283,7 +295,11 @@ class GraphQLResponseDecodeTests: XCTestCase {
             "author": "authorId",
         ], modelName: "Post"))
 
-        let result: Result<GraphQLResponse<MutationSync<AnyModel>>, APIError> = .fromAppSyncResponse(json: json, decodePath: "onCreatePost", modelName: nil)
+        let result: Result<GraphQLResponse<MutationSync<AnyModel>>, APIError> = .fromAppSyncResponse(
+            json: json,
+            decodePath: "onCreatePost",
+            modelName: "Post"
+        )
         XCTAssertNoThrow(try result.get())
         let mutationSync = try! result.get()
         XCTAssertThrowsError(try mutationSync.get()) { error in
@@ -305,7 +321,11 @@ class GraphQLResponseDecodeTests: XCTestCase {
             "a": "b"
         ]
 
-        let result: Result<GraphQLResponse<MutationSync<AnyModel>>, APIError> = .fromAppSyncResponse(json: json, decodePath: "onCreatePost", modelName: nil)
+        let result: Result<GraphQLResponse<MutationSync<AnyModel>>, APIError> = .fromAppSyncResponse(
+            json: json,
+            decodePath: "onCreatePost",
+            modelName: "Post"
+        )
         XCTAssertThrowsError(try result.get()) { error in
             if case .unknown(let description, _, _) = (error as! APIError) {
                 XCTAssertEqual("Failed to get data object or errors from GraphQL response", description)
@@ -325,7 +345,7 @@ class GraphQLResponseDecodeTests: XCTestCase {
         ]
 
         let jsonString = String(data: try! JSONEncoder().encode(json), encoding: .utf8)!
-        let response: GraphQLResponse<AnyModel> = .fromAppSyncResponse(string: jsonString, decodePath: nil)
+        let response: GraphQLResponse<AnyModel> = .fromAppSyncResponse(string: jsonString, decodePath: nil, modelName: "Post")
         XCTAssertNoThrow(try response.get())
         XCTAssertEqual(json.id?.stringValue, try response.get().identifier)
         XCTAssertEqual(json.__typename?.stringValue, try response.get().modelName)
@@ -334,7 +354,7 @@ class GraphQLResponseDecodeTests: XCTestCase {
     func testFromAppSyncResponse_withBrokenJsonString_failWithTransformationError() {
         SchemaData.modelSchemaRegistry.registerModels(registry: ModelRegistry.self)
         let jsonString = "{"
-        let response: GraphQLResponse<AnyModel> = .fromAppSyncResponse(string: jsonString, decodePath: nil)
+        let response: GraphQLResponse<AnyModel> = .fromAppSyncResponse(string: jsonString, decodePath: nil, modelName: "Post")
         XCTAssertThrowsError(try response.get()) { error in
             guard case .transformationError = error as! GraphQLResponseError<AnyModel> else {
                 XCTFail("Should failed with transformationError")
@@ -346,29 +366,39 @@ class GraphQLResponseDecodeTests: XCTestCase {
     func testFromAppSyncSubscriptionResponse_withJsonString_decodeCorrectly() {
         SchemaData.modelSchemaRegistry.registerModels(registry: ModelRegistry.self)
         let payloadJson: JSONValue = [
-            "onCreatePost": [
-                "__typename": "Post",
-                "id": "123",
-                "title": "test",
-                "author": "authorId",
-                "_version": 1,
-               "_deleted": nil,
-               "_lastChangedAt": 1707773705221
+            "data": [
+                "onCreatePost": [
+                    "__typename": "Post",
+                    "id": "123",
+                    "title": "test",
+                    "author": "authorId",
+                    "_version": 1,
+                   "_deleted": nil,
+                   "_lastChangedAt": 1707773705221
+                ]
             ]
         ]
 
         let jsonString = String(data: try! JSONEncoder().encode(payloadJson), encoding: .utf8)!
-        let response: GraphQLResponse<MutationSync<AnyModel>> = .fromAppSyncResponse(string: jsonString, decodePath: "onCreatePost")
+        let response: GraphQLResponse<MutationSync<AnyModel>> = .fromAppSyncResponse(
+            string: jsonString,
+            decodePath: "onCreatePost",
+            modelName: "Post"
+        )
         XCTAssertNoThrow(try response.get())
         let mutationSync = try! response.get()
-        XCTAssertEqual(payloadJson.onCreatePost?.id?.stringValue, mutationSync.model.identifier)
-        XCTAssertEqual(payloadJson.onCreatePost?.__typename?.stringValue, mutationSync.model.modelName)
+        XCTAssertEqual(payloadJson.data?.onCreatePost?.id?.stringValue, mutationSync.model.identifier)
+        XCTAssertEqual(payloadJson.data?.onCreatePost?.__typename?.stringValue, mutationSync.model.modelName)
     }
 
     func testFromAppSyncSubscriptionResponse_withWrongJsonString_failWithTransformationError() {
         SchemaData.modelSchemaRegistry.registerModels(registry: ModelRegistry.self)
         let jsonString = "{"
-        let response: GraphQLResponse<MutationSync<AnyModel>> = .fromAppSyncSubscriptionResponse(string: jsonString, decodePath: nil)
+        let response: GraphQLResponse<MutationSync<AnyModel>> = .fromAppSyncSubscriptionResponse(
+            string: jsonString,
+            decodePath: nil,
+            modelName: "Post"
+        )
         XCTAssertThrowsError(try response.get()) { error in
             guard case .transformationError = error as! GraphQLResponseError<MutationSync<AnyModel>> else {
                 XCTFail("Should failed with transformationError")

packages/amplify_datastore/ios/Classes/FlutterApiPlugin.swift

@@ -2,15 +2,13 @@ import Foundation
 import Flutter
 import Combine
 
-
-
 public class FlutterApiPlugin: APICategoryPlugin
 {
     public var key: PluginKey = "awsAPIPlugin"
     private let apiAuthFactory: APIAuthProviderFactory
     private let nativeApiPlugin: NativeApiPlugin
     private let nativeSubscriptionEvents: PassthroughSubject<NativeGraphQLSubscriptionResponse, Never>
-    private var cancellables = Set<AnyCancellable>()
+    private var cancellables = AtomicDictionary<AnyCancellable, Void>()
 
     init(
         apiAuthProviderFactory: APIAuthProviderFactory,
@@ -24,6 +22,7 @@ public class FlutterApiPlugin: APICategoryPlugin
 
     public func query<R>(request: GraphQLRequest<R>) async throws -> GraphQLTask<R>.Success where R : Decodable {
         let response = await asyncQuery(nativeRequest: request.toNativeGraphQLRequest())
+
         return try decodeGraphQLPayloadJson(request: request, payload: response.payloadJson)
     }
 
@@ -49,6 +48,7 @@ public class FlutterApiPlugin: APICategoryPlugin
 
         // TODO: shouldn't there be a timeout if there is no start_ack returned in a certain period of time
         let (sequence, cancellable) = nativeSubscriptionEvents
+            .setFailureType(to: Error.self)
             .receive(on: DispatchQueue.global())
             .filter { $0.subscriptionId == subscriptionId }
             .handleEvents(receiveCompletion: {_ in
@@ -81,8 +81,19 @@ public class FlutterApiPlugin: APICategoryPlugin
                         return nil
                 }
             }
+            .flatMap { (event: GraphQLSubscriptionEvent<R>) -> AnyPublisher<GraphQLSubscriptionEvent<R>, Error> in
+                if case .data(.failure(let graphQLResponseError)) = event,
+                   case .error(let errors) = graphQLResponseError,
+                   errors.contains(where: self.isUnauthorizedError(graphQLError:)) {
+                    return Fail(error: APIError.operationError("Unauthorized", "", nil)).eraseToAnyPublisher()
+                }
+                return Just(event).setFailureType(to: Error.self).eraseToAnyPublisher()
+            }
+            .eraseToAnyPublisher()
             .toAmplifyAsyncThrowingSequence()
-        cancellables.insert(cancellable) // the subscription is bind with class instance lifecycle, it should be released when stream is finished or unsubscribed
+
+        cancellables.set(value: (), forKey: cancellable) // the subscription is bind with class instance lifecycle, it should be released when stream is finished or unsubscribed
+
         sequence.send(.connection(.connecting))
         DispatchQueue.main.async {
             self.nativeApiPlugin.subscribe(request: request.toNativeGraphQLRequest()) { response in
@@ -99,10 +110,15 @@ public class FlutterApiPlugin: APICategoryPlugin
         guard let payload else {
             throw DataStoreError.decodingError("Request payload could not be empty", "")
         }
+
+        guard let datastoreOptions = request.options?.pluginOptions as? AWSAPIPluginDataStoreOptions else {
+            throw DataStoreError.decodingError("Failed to decode the GraphQLRequest due to a missing options field.", "")
+        }
 
         return GraphQLResponse<R>.fromAppSyncResponse(
             string: payload,
-            decodePath: request.decodePath
+            decodePath: request.decodePath,
+            modelName: datastoreOptions.modelName
         )
     }
 
@@ -113,13 +129,25 @@ public class FlutterApiPlugin: APICategoryPlugin
         guard let payload else {
             throw DataStoreError.decodingError("Request payload could not be empty", "")
         }
+
+        guard let datastoreOptions = request.options?.pluginOptions as? AWSAPIPluginDataStoreOptions else {
+            throw DataStoreError.decodingError("Failed to decode the GraphQLRequest due to a missing options field.", "")
+        }
 
         return GraphQLResponse<R>.fromAppSyncSubscriptionResponse(
             string: payload,
-            decodePath: request.decodePath
+            decodePath: request.decodePath,
+            modelName: datastoreOptions.modelName
         )
     }
-
+
+    private func isUnauthorizedError(graphQLError: GraphQLError) -> Bool {
+        guard case let .string(errorTypeValue) = graphQLError.extensions?["errorType"] else {
+            return false
+        }
+        return errorTypeValue == "Unauthorized"
+    }
+
     func asyncQuery(nativeRequest: NativeGraphQLRequest) async -> NativeGraphQLResponse {
         await withCheckedContinuation { continuation in
             DispatchQueue.main.async {

packages/amplify_datastore/ios/Classes/api/GraphQLRequest+Extension.swift

@@ -13,13 +13,16 @@ extension GraphQLRequest {
         let variablesJson = self.variables
             .flatMap { try? JSONSerialization.data(withJSONObject: $0, options: []) }
             .flatMap { String(data: $0, encoding: .utf8) }
+
+        let datastoreOptions = self.options?.pluginOptions as? AWSAPIPluginDataStoreOptions
 
         return NativeGraphQLRequest(
             document: self.document,
             apiName: self.apiName,
             variablesJson: variablesJson ?? "{}",
             responseType: String(describing: self.responseType),
-            decodePath: self.decodePath
+            decodePath: self.decodePath,
+            authMode: datastoreOptions?.authType.map { String(describing: $0) }
         )
     }
 }