feat: DataStore and API Lazy Loading with Custom Selection Set (#2583)

* adds selection set customization to api category

* fix selection set indent logic

* lazy load all commits

* Add PrimaryKeyOnly flag to GraphQLRequestBuilders

* API - Add initial PostComment4V2 tests

* API progress on PostComment4V2 tests

* API update and delete tests for PostComment4V2 tests

* cancel subscription in subscription tests

* Adding copy of DataStore tests and make buildable

* updated codegen for postcomment and minor testSave tests

* Setting up ProjectTeam codegen schemas and running initial test

* API LL3 LL4 testing

* API LL5 LL6 testing

* API LL2 testing

* disable unfinished tests

* Update DataStore integ models with modePath

* fix unit tests in API and DataStore

* test DS integration tests

* migrate LazyModel to LazyReference

* code clean up

* code clean up

* Update Amplify-Package.xcscheme

* add geo scheme

* code clean up

* address PR comments 1 - remove duplicate state enum, move primaryKeysOnly check within decorator

* fix primaryKey parameter

* add L12 codegenerated files

* API testing LL12

* DS LL12 test placeholder

* enable reduced selection set for DataStore for all operations

* disable reduced selection set for DataStore mutations

* disable reduced selection set for DataStore mutations 2

* remove print statments

* add typename from client side for DataStore MutationSync<AnyModel>

* add _deleted to new subscription selection sets only for nested children, more API CommentPost schema integ tests, fix AWSPluginsCore unit tests

* fix API integ tests for ProjectTeam 5/6

* optimize list and lazy reference decoding logic

* remove public loadingStrategy configuration

* fix race case for plugin resets wrt the ModelRegistry being reset

* update models and fix tests

* test stability - set up datastore only for selection set tests

* reduce logging in Decorator

* test stability - reduce syncMaxRecords to 100

* test stability - clearOnTearDown false for all tests

* API integ tests - subscribe team project LL8/9

* API integ tests - refactor LL112 into 3 separate test classes

* PhoneCall testing

* fix error handling
- remove DataError
- return API or DataStore error from respective plugin model providers
- return CoreError from core implementation

* code clean up
- rename shouldDecode to decode
- use _ prefix for public but intended internal

Co-authored-by: Daniel Rochetti <[email protected]>

Author: lawmicha

.swiftpm/xcode/xcshareddata/xcschemes/AWSLocationGeoPlugin.xcscheme

@@ -1,6 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <Scheme
-   LastUpgradeVersion = "1400"
+   LastUpgradeVersion = "1410"
    version = "1.3">
    <BuildAction
       parallelizeBuildables = "YES"
@@ -20,6 +20,20 @@
                ReferencedContainer = "container:">
             </BuildableReference>
          </BuildActionEntry>
+         <BuildActionEntry
+            buildForTesting = "YES"
+            buildForRunning = "NO"
+            buildForProfiling = "NO"
+            buildForArchiving = "NO"
+            buildForAnalyzing = "NO">
+            <BuildableReference
+               BuildableIdentifier = "primary"
+               BlueprintIdentifier = "AWSLocationGeoPluginTests"
+               BuildableName = "AWSLocationGeoPluginTests"
+               BlueprintName = "AWSLocationGeoPluginTests"
+               ReferencedContainer = "container:">
+            </BuildableReference>
+         </BuildActionEntry>
       </BuildActionEntries>
    </BuildAction>
    <TestAction

Amplify/Categories/API/Internal/APICategory+Resettable.swift

@@ -20,10 +20,6 @@ extension APICategory: Resettable {
             }
             await taskGroup.waitForAll()
         }
-        log.verbose("Resetting ModelRegistry and ModelListDecoderRegistry")
-        ModelRegistry.reset()
-        ModelListDecoderRegistry.reset()
-        log.verbose("Resetting ModelRegistry and ModelListDecoderRegistry: finished")
 
         isConfigured = false
     }

Amplify/Categories/DataStore/Internal/DataStoreCategory+Resettable.swift

@@ -20,10 +20,6 @@ extension DataStoreCategory: Resettable {
             }
             await taskGroup.waitForAll()
         }
-        log.verbose("Resetting ModelRegistry and ModelListDecoderRegistry")
-        ModelRegistry.reset()
-        ModelListDecoderRegistry.reset()
-        log.verbose("Resetting ModelRegistry and ModelListDecoderRegistry: finished")
 
         isConfigured = false
     }

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

@@ -36,7 +36,5 @@ extension ModelListDecoderRegistry {
 /// application making any change to these `public` types should be backward compatible, otherwise it will be a breaking
 /// change.
 public protocol ModelListDecoder {
-    static func shouldDecode<ModelType: Model>(modelType: ModelType.Type, decoder: Decoder) -> Bool
-    static func makeListProvider<ModelType: Model>(
-        modelType: ModelType.Type, decoder: Decoder) throws -> AnyModelListProvider<ModelType>
+    static func decode<ModelType: Model>(modelType: ModelType.Type, decoder: Decoder) -> AnyModelListProvider<ModelType>?
 }

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

@@ -6,7 +6,6 @@
 //
 
 import Foundation
-import Combine
 
 /// Empty protocol used as a marker to detect when the type is a `List`
 ///
@@ -23,7 +22,13 @@ public protocol ModelListMarker { }
 /// application making any change to these `public` types should be backward compatible, otherwise it will be a breaking
 /// change.
 public enum ModelListProviderState<Element: Model> {
-    case notLoaded
+    /// If the list represents an association between two models, the `associatedIdentifiers` will
+    /// hold the information necessary to query the associated elements (e.g. comments of a post)
+    ///
+    /// The associatedField represents the field to which the owner of the `List` is linked to.
+    /// For example, if `Post.comments` is associated with `Comment.post` the `List<Comment>`
+    /// of `Post` will have a reference to the `post` field in `Comment`.
+    case notLoaded(associatedIdentifiers: [String], associatedField: String)
     case loaded([Element])
 }
 

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

@@ -0,0 +1,82 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+
+/// Protocol used as a marker to detect when the type is a `LazyReference`.
+/// Used to retrieve either the `reference` or the `identifiers` of the Model directly, without having load a not
+/// loaded LazyReference. This is useful when translating the model object over to the payload required for the
+/// underlying storage, such as storing the values in DataStore's local database or AppSync GraphQL request payload.
+///
+///
+/// - Warning: Although this has `public` access, it is intended for internal & codegen use and should not be used
+/// directly by host applications. The behavior of this may change without warning. Though it is not used by host
+/// application making any change to these `public` types should be backward compatible, otherwise it will be a breaking
+/// change.
+public protocol _LazyReferenceValue {
+    var _state: _LazyReferenceValueState { get }
+}
+
+public enum _LazyReferenceValueState {
+    case notLoaded(identifiers: [LazyReferenceIdentifier]?)
+    case loaded(model: Model?)
+}
+
+/// State of the ModelProvider
+///
+/// - Warning: Although this has `public` access, it is intended for internal & codegen use and should not be used
+/// directly by host applications. The behavior of this may change without warning. Though it is not used by host
+/// application making any change to these `public` types should be backward compatible, otherwise it will be a breaking
+/// change.
+public enum ModelProviderState<Element: Model> {
+    case notLoaded(identifiers: [LazyReferenceIdentifier]?)
+    case loaded(model: Element?)
+}
+
+/// - Warning: Although this has `public` access, it is intended for internal & codegen use and should not be used
+/// directly by host applications. The behavior of this may change without warning. Though it is not used by host
+/// application making any change to these `public` types should be backward compatible, otherwise it will be a breaking
+/// change.
+public protocol ModelProvider {
+    associatedtype Element: Model
+    
+    func load() async throws -> Element?
+    
+    func getState() -> ModelProviderState<Element>
+}
+
+/// - Warning: Although this has `public` access, it is intended for internal & codegen use and should not be used
+/// directly by host applications. The behavior of this may change without warning. Though it is not used by host
+/// application making any change to these `public` types should be backward compatible, otherwise it will be a breaking
+/// change.
+public struct AnyModelProvider<Element: Model>: ModelProvider {
+    
+    private let loadAsync: () async throws -> Element?
+    private let getStateClosure: () -> ModelProviderState<Element>
+    
+    public init<Provider: ModelProvider>(provider: Provider) where Provider.Element == Self.Element {
+        self.loadAsync = provider.load
+        self.getStateClosure = provider.getState
+    }
+    public func load() async throws -> Element? {
+        try await loadAsync()
+    }
+    
+    public func getState() -> ModelProviderState<Element> {
+        getStateClosure()
+    }
+}
+
+/// - Warning: Although this has `public` access, it is intended for internal & codegen use and should not be used
+/// directly by host applications. The behavior of this may change without warning. Though it is not used by host
+/// application making any change to these `public` types should be backward compatible, otherwise it will be a breaking
+/// change.
+public extension ModelProvider {
+    func eraseToAnyModelProvider() -> AnyModelProvider<Element> {
+        AnyModelProvider(provider: self)
+    }
+}

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

@@ -0,0 +1,40 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+
+/// Registry of `ModelProviderDecoder`'s used to retrieve decoders that can create `ModelProvider`s to perform
+/// LazyReference functionality.
+///
+/// - Warning: Although this has `public` access, it is intended for internal & codegen use and should not be used
+/// directly by host applications. The behavior of this may change without warning. Though it is not used by host
+/// application making any change to these `public` types should be backward compatible, otherwise it will be a breaking
+/// change.
+public struct ModelProviderRegistry {
+    static var decoders = AtomicValue(initialValue: [ModelProviderDecoder.Type]())
+
+    /// Register a decoder during plugin configuration time, to allow runtime retrievals of model providers.
+    public static func registerDecoder(_ decoder: ModelProviderDecoder.Type) {
+        decoders.append(decoder)
+    }
+}
+
+extension ModelProviderRegistry {
+    static func reset() {
+        decoders.set([ModelProviderDecoder.Type]())
+    }
+}
+
+/// `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
+/// directly by host applications. The behavior of this may change without warning. Though it is not used by host
+/// application making any change to these `public` types should be backward compatible, otherwise it will be a breaking
+/// change.
+public protocol ModelProviderDecoder {
+    static func decode<ModelType: Model>(modelType: ModelType.Type, decoder: Decoder) -> AnyModelProvider<ModelType>?
+}

Amplify/Categories/DataStore/Model/Internal/Schema/ModelField+Association.swift

@@ -225,6 +225,19 @@ extension ModelField {
         }
         return true
     }
+    
+    /// - Warning: Although this has `public` access, it is intended for internal & codegen use and should not be used
+    ///   directly by host applications. The behavior of this may change without warning. Though it is not used by host
+    ///   application making any change to these `public` types should be backward compatible, otherwise it will be a
+    ///   breaking change.
+    public var _isBelongsToOrHasOne: Bool {
+        switch association {
+        case .belongsTo, .hasOne:
+            return true
+        case .hasMany, .none:
+            return false
+        }
+    }
 
     /// - Warning: Although this has `public` access, it is intended for internal & codegen use and should not be used
     ///   directly by host applications. The behavior of this may change without warning. Though it is not used by host

Amplify/Categories/DataStore/Model/Collection/ArrayLiteralListProvider.swift renamed to Amplify/Categories/DataStore/Model/Lazy/ArrayLiteralListProvider.swift

@@ -0,0 +1,36 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+
+// MARK: - DefaultModelProvider
+
+public struct DefaultModelProvider<Element: Model>: ModelProvider {
+    
+    var loadedState: ModelProviderState<Element>
+    
+    public init(element: Element? = nil) {
+        self.loadedState = .loaded(model: element)
+    }
+    
+    public init(identifiers: [LazyReferenceIdentifier]?) {
+        self.loadedState = .notLoaded(identifiers: identifiers)
+    }
+    
+    public func load() async throws -> Element? {
+        switch loadedState {
+        case .notLoaded:
+            throw CoreError.clientValidation("DefaultModelProvider does not provide loading capabilities", "")
+        case .loaded(let model):
+            return model
+        }
+    }
+    
+    public func getState() -> ModelProviderState<Element> {
+        loadedState
+    }
+}