fix(DataStore): pass DataStore source for API decoding (#3515)

Amplify/Categories/DataStore/Model/Internal/ModelProviderDecoder.swift

@@ -29,6 +29,17 @@ extension ModelProviderRegistry {
     }
 }
 
+/// Extension to hold the decoder sources
+public extension ModelProviderRegistry {
+
+    /// Static decoder sources that will be referenced to initialize different type of decoders having source as
+    /// a metadata. 
+    struct DecoderSource {
+        public static let dataStore = "DataStore"
+        public static let appSync = "AppSync"
+    }
+}
+
 /// `ModelProviderDecoder` provides decoding and lazy reference functionality.
 ///
 /// - Warning: Although this has `public` access, it is intended for internal & codegen use and should not be used

AmplifyPlugins/API/Sources/AWSAPIPlugin/Core/AppSyncModelDecoder.swift

@@ -13,15 +13,13 @@ import Amplify
 /// can be decoded to the Model, then the model provider is created as a "loaded" reference.
 public struct AppSyncModelDecoder: ModelProviderDecoder {
 
-    public static let AppSyncSource = "AppSync"
-
     /// Metadata that contains metadata of a model, specifically the identifiers used to hydrate the model.
     struct Metadata: Codable {
         let identifiers: [LazyReferenceIdentifier]
         let apiName: String?
         let source: String
 
-        init(identifiers: [LazyReferenceIdentifier], apiName: String?, source: String = AppSyncSource) {
+        init(identifiers: [LazyReferenceIdentifier], apiName: String?, source: String = ModelProviderRegistry.DecoderSource.appSync) {
             self.identifiers = identifiers
             self.apiName = apiName
             self.source = source
@@ -30,7 +28,7 @@ public struct AppSyncModelDecoder: ModelProviderDecoder {
 
     public static func decode<ModelType: Model>(modelType: ModelType.Type, decoder: Decoder) -> AnyModelProvider<ModelType>? {
         if let metadata = try? Metadata(from: decoder) {
-            if metadata.source == AppSyncSource {
+            if metadata.source == ModelProviderRegistry.DecoderSource.appSync {
                 log.verbose("Creating not loaded model \(modelType.modelName) with metadata \(metadata)")
                 return AppSyncModelProvider<ModelType>(metadata: metadata).eraseToAnyModelProvider()
             } else {

AmplifyPlugins/API/Sources/AWSAPIPlugin/Core/AppSyncModelMetadata.swift

@@ -31,8 +31,11 @@ public struct AppSyncModelMetadataUtils {
         }
     }
 
-    static func addMetadata(toModel graphQLData: JSONValue,
-                            apiName: String?) -> JSONValue {
+    static func addMetadata(
+        toModel graphQLData: JSONValue,
+        apiName: String?,
+        source: String = ModelProviderRegistry.DecoderSource.appSync) -> JSONValue {
+
         guard case var .object(modelJSON) = graphQLData else {
             Amplify.API.log.debug("Not an model object: \(graphQLData)")
             return graphQLData
@@ -85,7 +88,8 @@ public struct AppSyncModelMetadataUtils {
                 // Scenario: Belongs-To Primary Keys only are available for lazy loading
                 if let modelIdentifierMetadata = createModelIdentifierMetadata(associatedModelType,
                                                                                modelObject: modelObject,
-                                                                               apiName: apiName) {
+                                                                               apiName: apiName,
+                                                                               source: source) {
                     if let serializedMetadata = try? encoder.encode(modelIdentifierMetadata),
                        let metadataJSON = try? decoder.decode(JSONValue.self, from: serializedMetadata) {
                         Amplify.API.log.verbose("Adding [\(modelField.name): \(metadataJSON)]")
@@ -166,7 +170,8 @@ public struct AppSyncModelMetadataUtils {
     /// are more keys in the `modelObject` which means it was eager loaded.
     static func createModelIdentifierMetadata(_ associatedModel: Model.Type,
                                               modelObject: [String: JSONValue],
-                                              apiName: String?) -> AppSyncModelDecoder.Metadata? {
+                                              apiName: String?,
+                                              source: String) -> AppSyncModelDecoder.Metadata? {
         let primarykeys = associatedModel.schema.primaryKey
         var identifiers = [LazyReferenceIdentifier]()
         for identifierField in primarykeys.fields {
@@ -180,7 +185,10 @@ public struct AppSyncModelMetadataUtils {
         modelObject["_deleted"] = nil
         modelObject["_version"] = nil
         if !identifiers.isEmpty && (identifiers.count) == modelObject.keys.count {
-            return AppSyncModelDecoder.Metadata(identifiers: identifiers, apiName: apiName)
+            return AppSyncModelDecoder.Metadata(
+                identifiers: identifiers,
+                apiName: apiName,
+                source: source)
         } else {
             return nil
         }

AmplifyPlugins/API/Sources/AWSAPIPlugin/Core/AppSyncModelProvider.swift

@@ -12,19 +12,21 @@ import AWSPluginsCore
 public class AppSyncModelProvider<ModelType: Model>: ModelProvider {
 
     let apiName: String?
-
+    let source: String
     var loadedState: ModelProviderState<ModelType>
 
     // Creates a "not loaded" provider
     init(metadata: AppSyncModelDecoder.Metadata) {
         self.loadedState = .notLoaded(identifiers: metadata.identifiers)
         self.apiName = metadata.apiName
+        self.source = metadata.source
     }
 
     // Creates a "loaded" provider
     init(model: ModelType?) {
         self.loadedState = .loaded(model: model)
         self.apiName = nil
+        self.source = ModelProviderRegistry.DecoderSource.appSync
     }
 
     // MARK: - APIs
@@ -62,8 +64,12 @@ public class AppSyncModelProvider<ModelType: Model>: ModelProvider {
     public func encode(to encoder: Encoder) throws {
         switch loadedState {
         case .notLoaded(let identifiers):
-            var container = encoder.singleValueContainer()
-            try container.encode(identifiers)
+            let metadata = AppSyncModelDecoder.Metadata(
+                identifiers: identifiers ?? [],
+                apiName: apiName,
+                source: source)
+            try metadata.encode(to: encoder)
+
         case .loaded(let element):
             try element.encode(to: encoder)
         }

AmplifyPlugins/API/Sources/AWSAPIPlugin/Operation/AWSGraphQLOperation.swift

@@ -153,7 +153,7 @@ final public class AWSGraphQLOperation<R: Decodable>: GraphQLOperation<R> {
     private func getEndpointInterceptors(from request: GraphQLOperationRequest<R>) -> Result<AWSAPIEndpointInterceptors?, APIError> {
         getEndpointConfig(from: request).flatMap { endpointConfig in
             do {
-                if let pluginOptions = request.options.pluginOptions as? AWSPluginOptions,
+                if let pluginOptions = request.options.pluginOptions as? AWSAPIPluginDataStoreOptions,
                    let authType = pluginOptions.authType {
                     return .success(try pluginConfig.interceptorsForEndpoint(
                         withConfig: endpointConfig,

AmplifyPlugins/API/Sources/AWSAPIPlugin/Operation/AWSGraphQLSubscriptionTaskRunner.swift

@@ -81,7 +81,7 @@ public class AWSGraphQLSubscriptionTaskRunner<R: Decodable>: InternalTaskRunner,
 
         // Retrieve request plugin option and
         // auth type in case of a multi-auth setup
-        let pluginOptions = request.options.pluginOptions as? AWSPluginOptions
+        let pluginOptions = request.options.pluginOptions as? AWSAPIPluginDataStoreOptions
         let urlRequest = generateSubscriptionURLRequest(from: endpointConfig)
 
         // Retrieve the subscription connection
@@ -280,7 +280,7 @@ final public class AWSGraphQLSubscriptionOperation<R: Decodable>: GraphQLSubscri
 
         // Retrieve request plugin option and
         // auth type in case of a multi-auth setup
-        let pluginOptions = request.options.pluginOptions as? AWSPluginOptions
+        let pluginOptions = request.options.pluginOptions as? AWSAPIPluginDataStoreOptions
         let urlRequest = generateSubscriptionURLRequest(from: endpointConfig)
 
         // Retrieve the subscription connection

AmplifyPlugins/API/Sources/AWSAPIPlugin/Support/Decode/GraphQLResponseDecoder+DecodeData.swift

@@ -47,8 +47,17 @@ extension GraphQLResponseDecoder {
                   case .object(var graphQLDataObject) = graphQLData,
                   case .array(var graphQLDataArray) = graphQLDataObject["items"] {
             for (index, item) in graphQLDataArray.enumerated() {
-                let modelJSON = AppSyncModelMetadataUtils.addMetadata(toModel: item,
-                                                                      apiName: request.apiName)
+                let modelJSON: JSONValue
+                if let _ = (request.options.pluginOptions as? AWSAPIPluginDataStoreOptions) {
+                    modelJSON = AppSyncModelMetadataUtils.addMetadata(
+                        toModel: item,
+                        apiName: request.apiName,
+                        source: ModelProviderRegistry.DecoderSource.dataStore)
+                } else {
+                    modelJSON = AppSyncModelMetadataUtils.addMetadata(
+                        toModel: item,
+                        apiName: request.apiName)
+                }
                 graphQLDataArray[index] = modelJSON
             }
             graphQLDataObject["items"] = JSONValue.array(graphQLDataArray)
@@ -107,7 +116,7 @@ extension GraphQLResponseDecoder {
     // latest version of the developer's app will continue to work because the the mutation request sent from the
     // latest library continues to have the typename field.
     private func shouldAddTypename(to graphQLData: JSONValue) -> JSONValue? {
-        if let modelName = modelName,
+        if let modelName = dataStorePluginOptions?.modelName,
            request.responseType == MutationSync<AnyModel>.self,
            case var .object(modelJSON) = graphQLData,
            // No need to replace existing response payloads that have it already

AmplifyPlugins/API/Sources/AWSAPIPlugin/Support/Decode/GraphQLResponseDecoder.swift

@@ -15,18 +15,17 @@ class GraphQLResponseDecoder<R: Decodable> {
     var response: Data
     let decoder = JSONDecoder()
     let encoder = JSONEncoder()
-    let modelName: String?
+    let dataStorePluginOptions: AWSAPIPluginDataStoreOptions?
 
     public init(request: GraphQLOperationRequest<R>, response: Data = Data()) {
         self.request = request
         self.response = response
         decoder.dateDecodingStrategy = ModelDateFormatting.decodingStrategy
         encoder.dateEncodingStrategy = ModelDateFormatting.encodingStrategy
-        if let pluginOptions = request.options.pluginOptions as? AWSPluginOptions,
-           let modelName = pluginOptions.modelName {
-            self.modelName = modelName
+        if let pluginOptions = request.options.pluginOptions as? AWSAPIPluginDataStoreOptions {
+            self.dataStorePluginOptions = pluginOptions
         } else {
-            self.modelName = nil
+            self.dataStorePluginOptions = nil
         }
     }
 

AmplifyPlugins/API/Tests/AWSAPIPluginTests/Support/Decode/GraphQLResponseDecoderLazyPostComment4V2Tests.swift

@@ -376,7 +376,12 @@ class GraphQLResponseDecoderLazyPostComment4V2Tests: XCTestCase, SharedTestCases
             ]
         ]
 
+        let expectedEncodedCommentModel = """
+        {"content":"content","createdAt":null,"id":"id","post":{"identifiers":[{"name":"id","value":"\(post.id)"}],"source":"AppSync"},"updatedAt":null}
+        """
+
         let commentWithLazyLoadPost = try decoder.decodeToResponseType(graphQLData)
+        XCTAssertEqual(try commentWithLazyLoadPost.toJSON(), expectedEncodedCommentModel)
         XCTAssertEqual(commentWithLazyLoadPost.id, "id")
         XCTAssertEqual(commentWithLazyLoadPost.content, "content")
         switch commentWithLazyLoadPost._post.modelProvider.getState() {

AmplifyPlugins/Core/AWSPluginsCore/AWSAPIPluginDataStoreOptions.swift

@@ -0,0 +1,28 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+import Amplify
+
+/// Plugin specific options type
+///
+/// - Warning: Although this has `public` access, it is intended for internal use and should not be used directly
+///   by host applications. The behavior of this may change without warning.
+public struct AWSAPIPluginDataStoreOptions {
+
+    /// authorization type
+    public let authType: AWSAuthorizationType?
+
+    /// name of the model
+    public let modelName: String
+
+    public init(authType: AWSAuthorizationType?,
+                modelName: String) {
+        self.authType = authType
+        self.modelName = modelName
+    }
+}

AmplifyPlugins/Core/AWSPluginsCore/AWSPluginOptions.swift

@@ -9,6 +9,26 @@ import Foundation
 import Amplify
 
 /// Plugin specific options type
+///
+/// - Warning: Although this has `public` access, it is intended for internal use and should not be used directly
+///   by host applications. The behavior of this may change without warning.
+///
+/// This method was used internally by DataStore to pass information to APIPlugin, it
+/// has since been renamed to `AWSDataStorePluginOptions`. For customers
+/// looking to use the runtime authType parameter, this is a feature that should result in
+/// an options object on APIPlugin as something like `AWSAPIPluginOptions`, ie.
+///
+///```swift
+///public struct AWSAPIPluginOptions {
+///    /// authorization type
+///    public let authType: AWSAuthorizationType?
+///
+///    public init(authType: AWSAuthorizationType?) {
+///        self.authType = authType
+///    }
+///}
+///```
+@available(*, deprecated, message: "Intended for internal use.")
 public struct AWSPluginOptions {
 
     /// authorization type
@@ -17,7 +37,8 @@ public struct AWSPluginOptions {
     /// name of the model
     public let modelName: String?
 
-    public init(authType: AWSAuthorizationType?, modelName: String?) {
+    public init(authType: AWSAuthorizationType?,
+                modelName: String) {
         self.authType = authType
         self.modelName = modelName
     }

AmplifyPlugins/Core/AWSPluginsCore/Model/GraphQLRequest/GraphQLRequest+AnyModelWithSync.swift

@@ -70,7 +70,7 @@ extension GraphQLRequest: ModelSyncGraphQLRequestFactory {
         documentBuilder.add(decorator: AuthRuleDecorator(.query, authType: authType))
         let document = documentBuilder.build()
 
-        let awsPluginOptions = AWSPluginOptions(authType: authType, modelName: modelName)
+        let awsPluginOptions = AWSAPIPluginDataStoreOptions(authType: authType, modelName: modelName)
         let requestOptions = GraphQLRequest<MutationSyncResult?>.Options(pluginOptions: awsPluginOptions)
 
         return GraphQLRequest<MutationSyncResult?>(document: document.stringValue,
@@ -158,7 +158,7 @@ extension GraphQLRequest: ModelSyncGraphQLRequestFactory {
         documentBuilder.add(decorator: AuthRuleDecorator(.mutation, authType: authType))
         let document = documentBuilder.build()
 
-        let awsPluginOptions = AWSPluginOptions(authType: authType, modelName: modelSchema.name)
+        let awsPluginOptions = AWSAPIPluginDataStoreOptions(authType: authType, modelName: modelSchema.name)
         let requestOptions = GraphQLRequest<MutationSyncResult>.Options(pluginOptions: awsPluginOptions)
 
         return GraphQLRequest<MutationSyncResult>(document: document.stringValue,
@@ -180,7 +180,7 @@ extension GraphQLRequest: ModelSyncGraphQLRequestFactory {
         documentBuilder.add(decorator: AuthRuleDecorator(.subscription(subscriptionType, nil), authType: authType))
         let document = documentBuilder.build()
 
-        let awsPluginOptions = AWSPluginOptions(authType: authType, modelName: modelSchema.name)
+        let awsPluginOptions = AWSAPIPluginDataStoreOptions(authType: authType, modelName: modelSchema.name)
         let requestOptions = GraphQLRequest<MutationSyncResult>.Options(pluginOptions: awsPluginOptions)
         return GraphQLRequest<MutationSyncResult>(document: document.stringValue,
                                                   variables: document.variables,
@@ -202,7 +202,10 @@ extension GraphQLRequest: ModelSyncGraphQLRequestFactory {
         documentBuilder.add(decorator: AuthRuleDecorator(.subscription(subscriptionType, claims), authType: authType))
         let document = documentBuilder.build()
 
-        let awsPluginOptions = AWSPluginOptions(authType: authType, modelName: modelSchema.name)
+        let awsPluginOptions = AWSAPIPluginDataStoreOptions(
+            authType: authType,
+            modelName: modelSchema.name
+        )
         let requestOptions = GraphQLRequest<MutationSyncResult>.Options(pluginOptions: awsPluginOptions)
         return GraphQLRequest<MutationSyncResult>(document: document.stringValue,
                                                   variables: document.variables,
@@ -229,7 +232,7 @@ extension GraphQLRequest: ModelSyncGraphQLRequestFactory {
         documentBuilder.add(decorator: AuthRuleDecorator(.query, authType: authType))
         let document = documentBuilder.build()
 
-        let awsPluginOptions = AWSPluginOptions(authType: authType, modelName: modelSchema.name)
+        let awsPluginOptions = AWSAPIPluginDataStoreOptions(authType: authType, modelName: modelSchema.name)
         let requestOptions = GraphQLRequest<SyncQueryResult>.Options(pluginOptions: awsPluginOptions)
 
         return GraphQLRequest<SyncQueryResult>(document: document.stringValue,
@@ -259,7 +262,7 @@ extension GraphQLRequest: ModelSyncGraphQLRequestFactory {
         documentBuilder.add(decorator: AuthRuleDecorator(.mutation, authType: authType))
         let document = documentBuilder.build()
 
-        let awsPluginOptions = AWSPluginOptions(authType: authType, modelName: modelSchema.name)
+        let awsPluginOptions = AWSAPIPluginDataStoreOptions(authType: authType, modelName: modelSchema.name)
         let requestOptions = GraphQLRequest<MutationSyncResult>.Options(pluginOptions: awsPluginOptions)
 
         return GraphQLRequest<MutationSyncResult>(document: document.stringValue,

AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Core/DataStoreModelDecoder.swift

@@ -8,17 +8,17 @@
 import Foundation
 import Amplify
 import SQLite
+import AWSPluginsCore
 
 public struct DataStoreModelDecoder: ModelProviderDecoder {
 
-    public static let DataStoreSource = "DataStore"
-
     /// Metadata that contains the foreign key value of a parent model, which is the primary key of the model to be loaded.
     struct Metadata: Codable {
         let identifiers: [LazyReferenceIdentifier]
         let source: String
 
-        init(identifiers: [LazyReferenceIdentifier], source: String = DataStoreSource) {
+        init(identifiers: [LazyReferenceIdentifier],
+             source: String = ModelProviderRegistry.DecoderSource.dataStore) {
             self.identifiers = identifiers
             self.source = source
         }
@@ -38,7 +38,7 @@ public struct DataStoreModelDecoder: ModelProviderDecoder {
 
     public static func decode<ModelType: Model>(modelType: ModelType.Type, decoder: Decoder) -> AnyModelProvider<ModelType>? {
         if let metadata = try? DataStoreModelDecoder.Metadata(from: decoder) {
-            if metadata.source == DataStoreSource {
+            if metadata.source == ModelProviderRegistry.DecoderSource.dataStore {
                 return DataStoreModelProvider<ModelType>(metadata: metadata).eraseToAnyModelProvider()
             } else {
                 return nil

AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Core/DataStoreModelProvider.swift

@@ -55,9 +55,7 @@ public class DataStoreModelProvider<ModelType: Model>: ModelProvider {
         switch loadedState {
         case .notLoaded(let identifiers):
             let metadata = DataStoreModelDecoder.Metadata(identifiers: identifiers ?? [])
-            var container = encoder.singleValueContainer()
-            try container.encode(metadata)
-
+            try metadata.encode(to: encoder)
         case .loaded(let element):
             try element.encode(to: encoder)
         }