feat: DataStore and API Lazy Loading with Custom Selection Set (data-dev-preview)

.github/workflows/integ_test_api.yml

@@ -163,7 +163,6 @@ jobs:
           scheme: AWSAPIPluginGraphQLLambdaAuthTests
 
   api-lazy-load-test:
-    if: ${{ false }}
     needs: prepare-for-test
     runs-on: macos-12
     environment: IntegrationTest

.github/workflows/integ_test_datastore_lazy_load.yml

@@ -1,6 +1,8 @@
 name: DataStore Lazy Load Tests
 on:
   workflow_dispatch:
+  push:
+    branches: [main]
 
 permissions:
     id-token: write
@@ -27,7 +29,7 @@ jobs:
           aws_region: ${{ secrets.AWS_REGION }}
           aws_s3_bucket: ${{ secrets.AWS_S3_BUCKET_INTEG }}
 
-  datastore-integration-v2-test:
+  datastore-integration-lazy-load-test:
     timeout-minutes: 30
     needs: prepare-for-test
     runs-on: macos-12
@@ -53,5 +55,3 @@ jobs:
         with:
           project_path: ./AmplifyPlugins/DataStore/Tests/DataStoreHostApp
           scheme: AWSDataStorePluginLazyLoadTests
-
-

.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 associatedFields 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], associatedFields: [String])
     case loaded([Element])
 }
 
@@ -47,6 +52,9 @@ public protocol ModelListProvider {
     /// Asynchronously retrieve the next page as a new in-memory List object. Returns a failure if there
     /// is no next page of results. You can validate whether the list has another page with `hasNextPage()`.
     func getNextPage() async throws -> List<Element>
+
+    /// Custom encoder
+    func encode(to encoder: Encoder) throws
 }
 
 /// - Warning: Although this has `public` access, it is intended for internal & codegen use and should not be used
@@ -58,14 +66,16 @@ public struct AnyModelListProvider<Element: Model>: ModelListProvider {
     private let loadAsync: () async throws -> [Element]
     private let hasNextPageClosure: () -> Bool
     private let getNextPageAsync: () async throws -> List<Element>
-
+    private let encodeClosure: (Encoder) throws -> Void
+
     public init<Provider: ModelListProvider>(
         provider: Provider
     ) where Provider.Element == Self.Element {
         self.getStateClosure = provider.getState
         self.loadAsync = provider.load
         self.hasNextPageClosure = provider.hasNextPage
         self.getNextPageAsync = provider.getNextPage
+        self.encodeClosure = provider.encode
     }
 
     public func getState() -> ModelListProviderState<Element> {
@@ -83,6 +93,10 @@ public struct AnyModelListProvider<Element: Model>: ModelListProvider {
     public func getNextPage() async throws -> List<Element> {
         try await getNextPageAsync()
     }
+
+    public func encode(to encoder: Encoder) throws {
+        try encodeClosure(encoder)
+    }
 }
 
 /// - Warning: Although this has `public` access, it is intended for internal & codegen use and should not be used

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

@@ -0,0 +1,91 @@
+//
+// 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>
+
+    func encode(to encoder: Encoder) throws
+}
+
+/// - 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>
+    private let encodeClosure: (Encoder) throws -> Void
+
+    public init<Provider: ModelProvider>(provider: Provider) where Provider.Element == Self.Element {
+        self.loadAsync = provider.load
+        self.getStateClosure = provider.getState
+        self.encodeClosure = provider.encode
+    }
+
+    public func load() async throws -> Element? {
+        try await loadAsync()
+    }
+
+    public func getState() -> ModelProviderState<Element> {
+        getStateClosure()
+    }
+
+    public func encode(to encoder: Encoder) throws {
+        try encodeClosure(encoder)
+    }
+}
+
+/// - 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

@@ -87,7 +87,7 @@ import Foundation
 /// - 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.
 public enum ModelAssociation {
-    case hasMany(associatedFieldName: String?)
+    case hasMany(associatedFieldName: String?, associatedFieldNames: [String] = [])
     case hasOne(associatedFieldName: String?, targetNames: [String])
     case belongsTo(associatedFieldName: String?, targetNames: [String])
 
@@ -98,8 +98,8 @@ public enum ModelAssociation {
         return .belongsTo(associatedFieldName: nil, targetNames: targetNames)
     }
 
-    public static func hasMany(associatedWith: CodingKey?) -> ModelAssociation {
-        return .hasMany(associatedFieldName: associatedWith?.stringValue)
+    public static func hasMany(associatedWith: CodingKey? = nil, associatedFields: [CodingKey] = []) -> ModelAssociation {
+        return .hasMany(associatedFieldName: associatedWith?.stringValue, associatedFieldNames: associatedFields.map { $0.stringValue })
     }
 
     @available(*, deprecated, message: "Use hasOne(associatedWith:targetNames:)")
@@ -226,6 +226,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
     ///   application making any change to these `public` types should be backward compatible, otherwise it will be a
@@ -234,13 +247,9 @@ extension ModelField {
         if hasAssociation {
             let associatedModel = requiredAssociatedModelName
             switch association {
-            case .belongsTo(let associatedKey, _):
-                // TODO handle modelName casing (convert to camelCase)
-                let key = associatedKey ?? associatedModel
-                let schema = ModelRegistry.modelSchema(from: associatedModel)
-                return schema?.field(withName: key)
-            case .hasOne(let associatedKey, _),
-                 .hasMany(let associatedKey):
+            case .belongsTo(let associatedKey, _),
+                    .hasOne(let associatedKey, _),
+                    .hasMany(let associatedKey, _):
                 // TODO handle modelName casing (convert to camelCase)
                 let key = associatedKey ?? associatedModel
                 let schema = ModelRegistry.modelSchema(from: associatedModel)
@@ -252,6 +261,29 @@ extension ModelField {
         return nil
     }
 
+    /// - 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 associatedFieldNames: [String] {
+        switch association {
+        case .hasMany(let associatedKey, let associatedKeys):
+            if associatedKeys.isEmpty, let associatedKey = associatedKey {
+                return [associatedKey]
+            }
+            return associatedKeys
+
+        case .hasOne, .belongsTo:
+            return ModelRegistry.modelSchema(from: requiredAssociatedModelName)?
+                .primaryKey
+                .fields
+                .map(\.name) ?? []
+
+        case .none:
+            return []
+        }
+    }
+
     /// - 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

Amplify/Categories/DataStore/Model/Internal/Schema/ModelSchema+Definition.swift

@@ -241,6 +241,18 @@ public enum ModelFieldDefinition {
                       association: .hasMany(associatedWith: associatedKey))
     }
 
+    public static func hasMany(_ key: CodingKey,
+                               is nullability: ModelFieldNullability = .required,
+                               isReadOnly: Bool = false,
+                               ofType type: Model.Type,
+                               associatedFields associatedKeys: [CodingKey]) -> ModelFieldDefinition {
+        return .field(key,
+                      is: nullability,
+                      isReadOnly: isReadOnly,
+                      ofType: .collection(of: type),
+                      association: .hasMany(associatedWith: associatedKeys.first, associatedFields: associatedKeys))
+    }
+
     public static func hasOne(_ key: CodingKey,
                               is nullability: ModelFieldNullability = .required,
                               isReadOnly: Bool = false,