fix(DataStore): Remove from snapshot when item no longer matches predicate (#1522)

* fix(DataStore): Remove from snapshot when item no longer matches predicate

* address PR comments

Author: lawmicha

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Subscribe/DataStoreObserveQueryOperation.swift

@@ -234,9 +234,7 @@ public class AWSDataStoreObserveQueryOperation<M: Model>: AsynchronousOperation,
             finish()
             return
         }
-        if log.logLevel >= .debug {
-            stopwatch.start()
-        }
+        startSnapshotStopWatch()
         storageEngine.query(
             modelType,
             modelSchema: modelSchema,
@@ -252,7 +250,7 @@ public class AWSDataStoreObserveQueryOperation<M: Model>: AsynchronousOperation,
                 switch queryResult {
                 case .success(let queriedModels):
                     currentItems.sortedModels = queriedModels
-                    generateQuerySnapshot()
+                    sendSnapshot()
                 case .failure(let error):
                     self.passthroughPublisher.send(completion: .failure(error))
                     self.finish()
@@ -263,45 +261,50 @@ public class AWSDataStoreObserveQueryOperation<M: Model>: AsynchronousOperation,
 
     // MARK: Observe item changes
 
+    /// Subscribe to item changes with two subscribers (During Sync and After Sync). During Sync, the items are filtered
+    /// by name and predicate, then collected by the timeOrCount grouping, before sent for processing the snapshot.
+    /// After Sync, the item is only filtered by name, and not the predicate filter because updates to the item may
+    /// make it so that the item no longer matches the predicate and requires to be removed from `currentItems`.
+    /// This check is defered until `onItemChangedAfterSync` where the predicate is then used, and `currentItems` is
+    /// accessed under the serial queue.
     func subscribeToItemChanges() {
         batchItemsChangedSink = dataStorePublisher.publisher
             .filter { _ in !self.dispatchedModelSyncedEvent.get() }
-            .filter(onItemChangedFilter(mutationEvent:))
+            .filter(filterByModelName(mutationEvent:))
+            .filter(filterByPredicateMatch(mutationEvent:))
             .collect(.byTimeOrCount(serialQueue, itemsChangedPeriodicPublishTimeInSeconds, itemsChangedMaxSize))
             .sink(receiveCompletion: onReceiveCompletion(completed:),
-                  receiveValue: onItemsChange(mutationEvents:))
+                  receiveValue: onItemsChangeDuringSync(mutationEvents:))
 
         itemsChangedSink = dataStorePublisher.publisher
             .filter { _ in self.dispatchedModelSyncedEvent.get() }
-            .filter(onItemChangedFilter(mutationEvent:))
+            .filter(filterByModelName(mutationEvent:))
             .receive(on: serialQueue)
             .sink(receiveCompletion: onReceiveCompletion(completed:),
-                  receiveValue: onItemChange(mutationEvent:))
+                  receiveValue: onItemChangeAfterSync(mutationEvent:))
     }
 
-    func onItemChangedFilter(mutationEvent: MutationEvent) -> Bool {
-        guard mutationEvent.modelName == modelSchema.name else {
-            return false
-        }
+    func filterByModelName(mutationEvent: MutationEvent) -> Bool {
+        // Filter in the model when it matches the model name for this operation
+        mutationEvent.modelName == modelSchema.name
+    }
 
+    func filterByPredicateMatch(mutationEvent: MutationEvent) -> Bool {
+        // Filter in the model when there is no predicate to check against.
         guard let predicate = self.predicate else {
             return true
         }
-
         do {
             let model = try mutationEvent.decodeModel(as: modelType)
+            // Filter in the model when the predicate matches, otherwise filter out
             return predicate.evaluate(target: model)
         } catch {
             log.error(error: error)
             return false
         }
     }
 
-    func onItemChange(mutationEvent: MutationEvent) {
-        onItemsChange(mutationEvents: [mutationEvent])
-    }
-
-    func onItemsChange(mutationEvents: [MutationEvent]) {
+    func onItemsChangeDuringSync(mutationEvents: [MutationEvent]) {
         serialQueue.async {
             if self.isCancelled || self.isFinished {
                 self.finish()
@@ -310,31 +313,77 @@ public class AWSDataStoreObserveQueryOperation<M: Model>: AsynchronousOperation,
             guard self.observeQueryStarted, !mutationEvents.isEmpty else {
                 return
             }
-            if self.log.logLevel >= .debug {
-                self.stopwatch.start()
-            }
-            self.generateQuerySnapshot(itemsChanged: mutationEvents)
+
+            self.startSnapshotStopWatch()
+            self.apply(itemsChanged: mutationEvents)
+            self.sendSnapshot()
         }
     }
 
-    private func generateQuerySnapshot(itemsChanged: [MutationEvent] = []) {
-        updateCurrentItems(with: itemsChanged)
-        passthroughPublisher.send(currentSnapshot)
-        if log.logLevel >= .debug {
-            let time = stopwatch.stop()
-            log.debug("Time to generate snapshot: \(time) seconds")
+    // Item changes after sync is more elaborate than item changes during sync because the item was never filtered out
+    // by the predicate (unlike during sync). An item that no longer matches the predicate may already exist in the
+    // snapshot and now needs to be removed. The evaluation is done here under the serial queue since checking to
+    // remove the item requires that check on `currentItems` and is required to be performed under the serial queue.
+    func onItemChangeAfterSync(mutationEvent: MutationEvent) {
+        serialQueue.async {
+            if self.isCancelled || self.isFinished {
+                self.finish()
+                return
+            }
+            guard self.observeQueryStarted else {
+                return
+            }
+            self.startSnapshotStopWatch()
+
+            do {
+                let model = try mutationEvent.decodeModel(as: self.modelType)
+                guard let mutationType = MutationEvent.MutationType(rawValue: mutationEvent.mutationType) else {
+                    return
+                }
+
+                guard let predicate = self.predicate else {
+                    // 1. If there is no predicate, this item should be applied to the snapshot
+                    if self.currentItems.apply(model: model, mutationType: mutationType) {
+                        self.sendSnapshot()
+                    }
+                    return
+                }
+
+                // 2. When there is a predicate, evaluate further
+                let modelMatchesPredicate = predicate.evaluate(target: model)
+
+                guard !modelMatchesPredicate else {
+                    // 3. When the item matchs the predicate, the item should be applied to the snapshot
+                    if self.currentItems.apply(model: model, mutationType: mutationType) {
+                        self.sendSnapshot()
+                    }
+                    return
+                }
+
+                // 4. When the item does not match the predicate, and is an update/delete, then the item needs to be
+                // removed from `currentItems` because it no longer should be in the snapshot. If removing it was
+                // was successfully, then send a new snapshot
+                if mutationType == .update || mutationType == .delete, self.currentItems.remove(model) {
+                    self.sendSnapshot()
+                }
+            } catch {
+                self.log.error(error: error)
+                return
+            }
+
         }
     }
 
     /// Update `curentItems` list with the changed items.
     /// This method is not thread safe unless executed under the serial DispatchQueue `serialQueue`.
-    private func updateCurrentItems(with itemsChanged: [MutationEvent]) {
+    private func apply(itemsChanged: [MutationEvent]) {
         for item in itemsChanged {
             do {
                 let model = try item.decodeModel(as: modelType)
                 guard let mutationType = MutationEvent.MutationType(rawValue: item.mutationType) else {
                     return
                 }
+
                 currentItems.apply(model: model, mutationType: mutationType)
             } catch {
                 log.error(error: error)
@@ -343,6 +392,20 @@ public class AWSDataStoreObserveQueryOperation<M: Model>: AsynchronousOperation,
         }
     }
 
+    private func startSnapshotStopWatch() {
+        if log.logLevel >= .debug {
+            stopwatch.start()
+        }
+    }
+
+    private func sendSnapshot() {
+        passthroughPublisher.send(currentSnapshot)
+        if log.logLevel >= .debug {
+            let time = stopwatch.stop()
+            log.debug("Time to generate snapshot: \(time) seconds")
+        }
+    }
+
     private func onReceiveCompletion(completed: Subscribers.Completion<DataStoreError>) {
         if isCancelled || isFinished {
             finish()

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Subscribe/Support/SortedList.swift

@@ -9,42 +9,47 @@ import Amplify
 
 class SortedList<ModelType: Model> {
     var sortedModels: [ModelType]
+    var modelIds: Set<Model.Identifier>
     private let sortInput: [QuerySortDescriptor]?
     private let modelSchema: ModelSchema
 
     init(sortInput: [QuerySortDescriptor]?, modelSchema: ModelSchema) {
         self.sortedModels = []
+        self.modelIds = []
         self.sortInput = sortInput
         self.modelSchema = modelSchema
     }
 
     /// Apply the incoming `model` to the sorted array based on the mutation type. This logic accounts for duplicate
     /// events since identical events may have different sources (local and remote). When the mutation type is delete,
-    /// remove it if it exists in the array. When create/update, replace it if it exists in the array. If it does not
-    /// exist in the array, then add the model to the array in the correct position.
-    func apply(model: ModelType, mutationType: MutationEvent.MutationType) {
+    /// remove it if it exists in the array. When create/update and it has an sort order, then remove and add it back in the
+    /// correct sort order. if there is no sort order, replace it.
+    /// Return `true` if something occured (added, replaced, deleted), otherwise `false`
+    @discardableResult func apply(model: ModelType, mutationType: MutationEvent.MutationType) -> Bool {
         if mutationType == MutationEvent.MutationType.delete {
-            if let index = sortedModels.firstIndex(where: { $0.id == model.id }) {
-                sortedModels.remove(at: index)
-            }
-        } else if let index = sortedModels.firstIndex(where: { $0.id == model.id }) {
-            sortedModels[index] = model
-        } else {
-            add(model: model)
+            return remove(model)
+        }
+
+        guard let sortInputs = sortInput else {
+            // If there is no sort order, check if it exists to replace, otherwise add it to the end
+            appendOrReplace(model)
+            return true
         }
+
+        // When there is a sort input, always attempt to remove it before adding to the correct position.
+        // If we had simply replaced it, and the update if applied to a field on the model that is the same as the
+        // sort field, then it may no longer be in the correct position.
+        _ = remove(model)
+        add(model: model, sortInputs: sortInputs)
+        return true
     }
 
     /// Add the incoming `model` to the sorted array based on the sort input, or at the end if none is provided.
     /// Search for the index by comparing the incoming model with the current model in the binary search traversal.
     /// If the models are equal in terms of their sort order (comparator returns `nil`), then move onto the next sort
     /// input. If all sort comparators return `nil`, then the incoming model is equal to the current model on all
     /// sort inputs,  and inserting at the index will maintain the overall sort order.
-    func add(model: ModelType) {
-        guard let sortInputs = sortInput else {
-            sortedModels.append(model)
-            return
-        }
-
+    func add(model: ModelType, sortInputs: [QuerySortDescriptor]) {
         let index = sortedModels.binarySearch { existingModel in
             var sortOrder: Bool?
             var sortIndex: Int = 0
@@ -59,6 +64,28 @@ class SortedList<ModelType: Model> {
         }
 
         sortedModels.insert(model, at: index)
+        modelIds.insert(model.id)
+    }
+
+    /// Tries to remove the `model`, if removed then return `true`, otherwise `false`
+    func remove(_ model: ModelType) -> Bool {
+        if modelIds.contains(model.id), let index = sortedModels.firstIndex(where: { $0.id == model.id }) {
+            sortedModels.remove(at: index)
+            modelIds.remove(model.id)
+            return true
+        } else {
+            return false
+        }
+    }
+
+    /// Tries to replace the model with `model` if it already exists, otherwise append it to at the end
+    func appendOrReplace(_ model: ModelType) {
+        if modelIds.contains(model.id), let index = sortedModels.firstIndex(where: { $0.id == model.id }) {
+            sortedModels[index] = model
+        } else {
+            sortedModels.append(model)
+            modelIds.insert(model.id)
+        }
     }
 }
 

AmplifyPlugins/DataStore/AWSDataStoreCategoryPluginIntegrationTests/DataStoreLocalStoreTests.swift

@@ -376,15 +376,61 @@ class DataStoreLocalStoreTests: LocalStoreIntegrationTestBase {
         }
     }
 
+    /// DataStore without sync to cloud enabled.
+    ///
+    /// - Given: DataStore is set up with models, and local store is populated with models
+    /// - When:
+    ///    - ObserveQuery is called, add more models
+    /// - Then:
+    ///    - The first snapshot will have initial models and may have additional models
+    ///    - There may be subsequent snapshots based on how the items are batched
+    ///    - The last snapshot will have a total of initial plus additional models
+    @available(iOS 13.0, *)
+    func testObserveQuery() throws {
+        let cleared = expectation(description: "DataStore cleared")
+        Amplify.DataStore.clear { result in
+            switch result {
+            case .success:
+                cleared.fulfill()
+            case .failure(let error):
+                XCTFail("\(error)")
+            }
+        }
+        wait(for: [cleared], timeout: 2)
+        _ = setUpLocalStore(numberOfPosts: 14)
+        var snapshotCount = 0
+        let allSnapshotsReceived = expectation(description: "query snapshots received")
+
+        let sink = Amplify.DataStore.observeQuery(for: Post.self).sink { completed in
+            switch completed {
+            case .finished:
+                break
+            case .failure(let error):
+                XCTFail("\(error)")
+            }
+        } receiveValue: { querySnapshot in
+            snapshotCount += 1
+            XCTAssertFalse(querySnapshot.isSynced)
+            if querySnapshot.items.count == 30 {
+                allSnapshotsReceived.fulfill()
+            }
+        }
+        _ = setUpLocalStore(numberOfPosts: 15)
+        wait(for: [allSnapshotsReceived], timeout: 100)
+        XCTAssertTrue(snapshotCount >= 2)
+        sink.cancel()
+    }
+
     func setUpLocalStore(numberOfPosts: Int) -> [Post] {
         var savedPosts = [Post]()
-        for _ in 0 ..< numberOfPosts {
+        for id in 0 ..< numberOfPosts {
             let saveSuccess = expectation(description: "Save post completed")
             let post = Post(title: "title\(Int.random(in: 0 ... 5))",
                             content: "content",
                             createdAt: .now(),
                             rating: Double(Int.random(in: 0 ... 5)))
             savedPosts.append(post)
+            print("\(id) \(post.id)")
             Amplify.DataStore.save(post) { result in
                 switch result {
                 case .success:
@@ -395,6 +441,7 @@ class DataStoreLocalStoreTests: LocalStoreIntegrationTestBase {
             }
             wait(for: [saveSuccess], timeout: 10)
         }
+
         return savedPosts
     }