test(datastore): Lazy loading integration tests

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

@@ -11,6 +11,8 @@ public protocol LazyModelMarker {
     associatedtype Element: Model
 
     var element: Element? { get }
+
+    var identifiers: [String: String]? { get }
 }
 
 public struct AnyModelProvider<Element: Model>: ModelProvider {
@@ -41,6 +43,6 @@ public protocol ModelProvider {
 }
 
 public enum ModelProviderState<Element: Model> {
-    case notLoaded(identifiers: [String: String])
+    case notLoaded(identifiers: [String: String]?)
     case loaded(Element?)
 }

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

@@ -225,6 +225,15 @@ extension ModelField {
         }
         return true
     }
+
+    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/Lazy/DefaultModelProvider.swift

@@ -10,18 +10,31 @@ import Foundation
 // MARK: - DefaultModelProvider
 
 public struct DefaultModelProvider<Element: Model>: ModelProvider {
-
-    let element: Element?
+    enum LoadedState {
+        case notLoaded(identifiers: [String: String]?)
+        case loaded(model: Element?)
+    }
+
+    var loadedState: LoadedState
+
     public init(element: Element? = nil) {
-        self.element = element
+        self.loadedState = .loaded(model: element)
+    }
+
+    public init(identifiers: [String: String]?) {
+        self.loadedState = .notLoaded(identifiers: identifiers)
     }
 
     public func load() async throws -> Element? {
-        return element
+        return Fatal.preconditionFailure("DefaultModelProvider does not provide loading capabilities")
     }
 
     public func getState() -> ModelProviderState<Element> {
-        return .loaded(element)
+        switch loadedState {
+        case .notLoaded(let identifiers):
+            return .notLoaded(identifiers: identifiers)
+        case .loaded(let model):
+            return .loaded(model)
+        }
     }
-
 }

Amplify/Categories/DataStore/Model/Lazy/LazyModel.swift

@@ -12,7 +12,7 @@ public class LazyModel<Element: Model>: Codable, LazyModelMarker {
 
     /// Represents the data state of the `LazyModel`.
     enum LoadedState {
-        case notLoaded(identifiers: [String: String])
+        case notLoaded(identifiers: [String: String]?)
         case loaded(Element?)
     }
     var loadedState: LoadedState
@@ -42,6 +42,17 @@ public class LazyModel<Element: Model>: Codable, LazyModelMarker {
         }
     }
 
+    public var identifiers: [String: String]? {
+        get {
+            switch loadedState {
+            case .notLoaded(let identifiers):
+                return identifiers
+            case .loaded:
+                return nil
+            }
+        }
+    }
+
     public init(modelProvider: AnyModelProvider<Element>) {
         self.modelProvider = modelProvider
         switch self.modelProvider.getState() {
@@ -52,11 +63,16 @@ public class LazyModel<Element: Model>: Codable, LazyModelMarker {
         }
     }
 
-    public convenience init(element: Element? = nil) {
+    public convenience init(element: Element?) {
         let modelProvider = DefaultModelProvider(element: element).eraseToAnyModelProvider()
         self.init(modelProvider: modelProvider)
     }
 
+    public convenience init(identifiers: [String: String]?) {
+        let modelProvider = DefaultModelProvider<Element>(identifiers: identifiers).eraseToAnyModelProvider()
+        self.init(modelProvider: modelProvider)
+    }
+
     required convenience public init(from decoder: Decoder) throws {
         for modelDecoder in ModelProviderRegistry.decoders.get() {
             if modelDecoder.shouldDecode(modelType: Element.self, decoder: decoder) {
@@ -70,7 +86,7 @@ public class LazyModel<Element: Model>: Codable, LazyModelMarker {
             let element = try Element(from: decoder)
             self.init(element: element)
         } else {
-            self.init()
+            self.init(identifiers: nil)
         }
     }
 

AmplifyPlugins/Core/AWSPluginsCore/Model/Decorator/ModelDecorator.swift

@@ -14,9 +14,11 @@ import Amplify
 public struct ModelDecorator: ModelBasedGraphQLDocumentDecorator {
 
     private let model: Model
-
-    public init(model: Model) {
+    private let mutationType: GraphQLMutationType
+
+    public init(model: Model, mutationType: GraphQLMutationType) {
         self.model = model
+        self.mutationType = mutationType
     }
 
     public func decorate(_ document: SingleDirectiveGraphQLDocument,
@@ -27,7 +29,7 @@ public struct ModelDecorator: ModelBasedGraphQLDocumentDecorator {
     public func decorate(_ document: SingleDirectiveGraphQLDocument,
                          modelSchema: ModelSchema) -> SingleDirectiveGraphQLDocument {
         var inputs = document.inputs
-        var graphQLInput = model.graphQLInputForMutation(modelSchema)
+        var graphQLInput = model.graphQLInputForMutation(modelSchema, mutationType: mutationType)
 
         if !modelSchema.authRules.isEmpty {
             modelSchema.authRules.forEach { authRule in

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

@@ -245,7 +245,7 @@ extension GraphQLRequest: ModelSyncGraphQLRequestFactory {
         var documentBuilder = ModelBasedGraphQLDocumentBuilder(modelName: modelSchema.name,
                                                                operationType: .mutation)
         documentBuilder.add(decorator: DirectiveNameDecorator(type: type))
-        documentBuilder.add(decorator: ModelDecorator(model: model))
+        documentBuilder.add(decorator: ModelDecorator(model: model, mutationType: type))
         if let filter = filter {
             documentBuilder.add(decorator: FilterDecorator(filter: filter))
         }

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

@@ -165,15 +165,15 @@ extension GraphQLRequest: ModelGraphQLRequestFactory {
 
         switch type {
         case .create:
-            documentBuilder.add(decorator: ModelDecorator(model: model))
+            documentBuilder.add(decorator: ModelDecorator(model: model, mutationType: type))
         case .delete:
             documentBuilder.add(decorator: ModelIdDecorator(model: model,
                                                             schema: modelSchema))
             if let predicate = predicate {
                 documentBuilder.add(decorator: FilterDecorator(filter: predicate.graphQLFilter(for: modelSchema)))
             }
         case .update:
-            documentBuilder.add(decorator: ModelDecorator(model: model))
+            documentBuilder.add(decorator: ModelDecorator(model: model, mutationType: type))
             if let predicate = predicate {
                 documentBuilder.add(decorator: FilterDecorator(filter: predicate.graphQLFilter(for: modelSchema)))
             }

AmplifyPlugins/Core/AWSPluginsCore/Model/Support/Model+GraphQL.swift

@@ -31,7 +31,7 @@ extension Model {
     /// Returns the input used for mutations
     /// - Parameter modelSchema: model's schema
     /// - Returns: A key-value map of the GraphQL mutation input
-    func graphQLInputForMutation(_ modelSchema: ModelSchema) -> GraphQLInput {
+    func graphQLInputForMutation(_ modelSchema: ModelSchema, mutationType: GraphQLMutationType) -> GraphQLInput {
         var input: GraphQLInput = [:]
 
         // filter existing non-readonly fields
@@ -43,7 +43,7 @@ extension Model {
             guard let value = modelFieldValue else {
                 // Special case for associated models when the value is `nil`, by setting all of the associated
                 // model's primary key fields (targetNames) to `nil`.
-                if case .model = modelField.type {
+                if case .model = modelField.type { // add it for "belongs-to"
                     let fieldNames = getFieldNameForAssociatedModels(modelField: modelField)
                     for fieldName in fieldNames {
                         // Only set to `nil` if it has not been set already. For hasOne relationships, where the
@@ -59,6 +59,8 @@ extension Model {
                             input.updateValue(nil, forKey: fieldName)
                         }
                     }
+                } else if case .collection = modelField.type { // skip all "has-many"
+                    continue
                 } else {
                     input.updateValue(nil, forKey: name)
                 }
@@ -82,9 +84,10 @@ extension Model {
                 // get the associated model target names and their values
                 let associatedModelIds = associatedModelIdentifierFields(fromModelValue: value,
                                                                          field: modelField,
-                                                                         associatedModelName: associateModelName)
+                                                                         associatedModelName: associateModelName,
+                                                                         mutationType: mutationType)
                 for (fieldName, fieldValue) in associatedModelIds {
-                    input[fieldName] = fieldValue
+                    input.updateValue(fieldValue, forKey: fieldName)
                 }
             case .embedded, .embeddedCollection:
                 if let encodable = value as? Encodable {
@@ -152,22 +155,28 @@ extension Model {
     ///            and `value` its value in the associated model
     private func associatedModelIdentifierFields(fromModelValue value: Any,
                                                  field: ModelField,
-                                                 associatedModelName: String) -> [(String, Persistable)] {
+                                                 associatedModelName: String,
+                                                 mutationType: GraphQLMutationType) -> [(String, Persistable?)] {
         guard let associateModelSchema = ModelRegistry.modelSchema(from: associatedModelName) else {
             preconditionFailure("Associated model \(associatedModelName) not found.")
         }
 
         let fieldNames = getFieldNameForAssociatedModels(modelField: field)
-        let values = getModelIdentifierValues(from: value, modelSchema: associateModelSchema)
+        var values = getModelIdentifierValues(from: value, modelSchema: associateModelSchema)
 
-        // if the field is required, the associated field keys and values should match
-        if fieldNames.count != values.count, field.isRequired {
-            preconditionFailure(
+        if fieldNames.count != values.count {
+            // if the field is required, the associated field keys and values should match
+            if field.isRequired {
+                preconditionFailure(
                 """
                 Associated model target names and values for field \(field.name) of model \(modelName) mismatch.
                 There is a possibility that is an issue with the generated models.
                 """
-            )
+                )
+            } else if mutationType == .update {
+                // otherwise, pad the values with `nil` to account for removals of associations on updates.
+                values = [Persistable?](repeating: nil, count: fieldNames.count)
+            }
         }
 
         return Array(zip(fieldNames, values))
@@ -179,7 +188,7 @@ extension Model {
     ///   - value: model value
     ///   - modelSchema: model's schema
     /// - Returns: array of values of its primary key
-    private func getModelIdentifierValues(from value: Any, modelSchema: ModelSchema) -> [Persistable] {
+    private func getModelIdentifierValues(from value: Any, modelSchema: ModelSchema) -> [Persistable?] {
         if let modelValue = value as? Model {
             return modelValue.identifier(schema: modelSchema).values
         } else if let optionalModel = value as? Model?,

AmplifyPlugins/Core/AWSPluginsCore/Model/Support/ModelSchema+GraphQL.swift

@@ -49,7 +49,7 @@ extension ModelSchema {
     /// The list of fields formatted for GraphQL usage.
     var graphQLFields: [ModelField] {
         sortedFields.filter { field in
-            !field.hasAssociation || field.isAssociationOwner
+            !field.hasAssociation || field.isBelongsToOrHasOne
         }
     }
 }

AmplifyPlugins/Core/AWSPluginsCore/Model/Support/SelectionSet.swift

@@ -34,18 +34,25 @@ extension SelectionSet {
         withModelFields(fields)
     }
 
-    func withModelFields(_ fields: [ModelField]) {
+    func withModelFields(_ fields: [ModelField], recursive: Bool = true) {
         fields.forEach { field in
             if field.isEmbeddedType, let embeddedTypeSchema = field.embeddedTypeSchema {
                 let child = SelectionSet(value: .init(name: field.name, fieldType: .embedded))
                 child.withEmbeddableFields(embeddedTypeSchema.sortedFields)
                 self.addChild(settingParentOf: child)
-            } else if field.isAssociationOwner,
-                let associatedModelName = field.associatedModelName,
-                let schema = ModelRegistry.modelSchema(from: associatedModelName) {
-                let child = SelectionSet(value: .init(name: field.name, fieldType: .model))
-                child.withModelFields(schema.graphQLFields)
-                self.addChild(settingParentOf: child)
+            } else if field.isBelongsToOrHasOne,
+                      let associatedModelName = field.associatedModelName,
+                      let schema = ModelRegistry.modelSchema(from: associatedModelName) {
+                if recursive {
+                    var recursive = recursive
+                    if field.isBelongsToOrHasOne {
+                        recursive = false
+                    }
+
+                    let child = SelectionSet(value: .init(name: field.name, fieldType: .model))
+                    child.withModelFields(schema.graphQLFields, recursive: recursive)
+                    self.addChild(settingParentOf: child)
+                }
             } else {
                 self.addChild(settingParentOf: .init(value: .init(name: field.graphQLName, fieldType: .value)))
             }

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

@@ -54,5 +54,5 @@ public struct DataStoreModelDecoder: ModelProviderDecoder {
 
 /// Metadata that contains the primary keys and values of a model
 public struct DataStoreModelIdentifierMetadata: Codable {
-    let identifier: String
+    let identifier: String?
 }

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

@@ -12,22 +12,26 @@ import Combine
 public class DataStoreModelProvider<ModelType: Model>: ModelProvider {
 
     enum LoadedState {
-        case notLoaded(identifiers: [String: String])
+        case notLoaded(identifiers: [String: String]?)
         case loaded(model: ModelType?)
     }
 
     var loadedState: LoadedState
 
     convenience init(metadata: DataStoreModelIdentifierMetadata) {
+        if let identifier = metadata.identifier {
+            self.init(identifiers: [ModelType.schema.primaryKey.sqlName: identifier])
+        } else {
+            self.init(identifiers: nil)
+        }
 
-        self.init(identifiers: [ModelType.schema.primaryKey.sqlName: metadata.identifier])
     }
 
     init(model: ModelType?) {
         self.loadedState = .loaded(model: model)
     }
 
-    init(identifiers: [String: String]) {
+    init(identifiers: [String: String]?) {
         self.loadedState = .notLoaded(identifiers: identifiers)
     }
 
@@ -36,7 +40,7 @@ public class DataStoreModelProvider<ModelType: Model>: ModelProvider {
     public func load() async throws -> ModelType? {
         switch loadedState {
         case .notLoaded(let identifiers):
-            guard let identifier = identifiers.first else {
+            guard let identifiers = identifiers, let identifier = identifiers.first else {
                 return nil
             }
             let queryPredicate: QueryPredicate = field(identifier.key).eq(identifier.value)

AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Storage/SQLite/Model+SQLite.swift

@@ -113,7 +113,9 @@ extension Model {
                 if let associatedModelValue = value as? Model {
                     return associatedModelValue.identifier
                 } else if let associatedLazyModel = value as? (any LazyModelMarker) {
-                    return associatedLazyModel.element?.identifier
+                    // The identifier (sometimes the FK), comes from the loaded model's identifier or
+                    // from the not loaded identifier's first and only value
+                    return associatedLazyModel.element?.identifier ?? associatedLazyModel.identifiers?.first?.value
                 } else if let associatedModelJSON = value as? [String: JSONValue] {
                     return associatedPrimaryKeyValue(fromJSON: associatedModelJSON,
                                                      associatedModelSchema: associatedModelSchema)

AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Storage/SQLite/Statement+Model.swift

@@ -102,7 +102,6 @@ extension Statement: StatementModelConvertible {
                 from: value,
                 fieldType: field.type
             )
-
             // Check if the value for the primary key is `nil`. This is when an associated model does not exist.
             // To create a decodable `modelDictionary` that can be decoded to the Model types, the entire
             // object at this particular key should be set to `nil`. The following code does that by dropping the last
@@ -135,7 +134,10 @@ extension Statement: StatementModelConvertible {
             // For example, when the value is the `id` of the Blog, then the field.isPrimaryKey is satisfied.
             // Every association of the Blog, such as the has-many Post is populated with the List with
             // associatedId == blog's id. This way, the list of post can be lazily loaded later using the associated id.
-            if let id = modelValue as? String, field.isPrimaryKey {
+            if let id = modelValue as? String,
+                (field.name == ModelIdentifierFormat.Custom.sqlColumnName || // this is the `@@primaryKey` (CPK)
+                 (schema.primaryKey.fields.count == 1 // or there's only one primary key (not composite key)
+                  && schema.primaryKey.indexOfField(named: field.name) != nil)) { // and this field is the primary key
                 let associations = schema.fields.values.filter {
                     $0.isArray && $0.hasAssociation
                 }

AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Sync/MutationSync/OutgoingMutationQueue/OutgoingMutationQueue.swift

@@ -355,7 +355,8 @@ final class OutgoingMutationQueue: OutgoingMutationQueueBehavior {
                                      data: outboxMutationProcessedEvent)
             Amplify.Hub.dispatch(to: .dataStore, payload: payload)
         } catch {
-            log.error("\(#function) Couldn't decode local model as \(mutationEvent.modelName)")
+            log.error("\(#function) Couldn't decode local model as \(mutationEvent.modelName) \(error)")
+            log.error("\(#function) Couldn't decode from \(mutationEvent.json)")
             return
         }
     }
@@ -370,7 +371,8 @@ final class OutgoingMutationQueue: OutgoingMutationQueueBehavior {
                                      data: outboxMutationEnqueuedEvent)
             Amplify.Hub.dispatch(to: .dataStore, payload: payload)
         } catch {
-            log.error("\(#function) Couldn't decode local model as \(mutationEvent.modelName)")
+            log.error("\(#function) Couldn't decode local model as \(mutationEvent.modelName) \(error)")
+            log.error("\(#function) Couldn't decode from \(mutationEvent.json)")
             return
         }
     }