fix: Handle .all predicate for create/update (#1640)

Author: lawmicha

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Storage/SQLite/QueryPredicate+SQLite.swift

@@ -50,3 +50,13 @@ extension QueryOperator {
         }
     }
 }
+
+extension QueryPredicate {
+    var isAll: Bool {
+        if let allPredicate = self as? QueryPredicateConstant, allPredicate == .all {
+            return true
+        } else {
+            return false
+        }
+    }
+}

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Storage/SQLite/StorageEngineAdapter+SQLite.swift

@@ -128,7 +128,7 @@ final class SQLiteStorageEngineAdapter: StorageEngineAdapter {
             let modelExists = try exists(modelSchema, withId: model.id)
 
             if !modelExists {
-                if condition != nil {
+                if let condition = condition, !condition.isAll {
                     let dataStoreError = DataStoreError.invalidCondition(
                         "Cannot apply a condition on model which does not exist.",
                         "Save the model instance without a condition first.")
@@ -141,7 +141,7 @@ final class SQLiteStorageEngineAdapter: StorageEngineAdapter {
             }
 
             if modelExists {
-                if condition != nil {
+                if let condition = condition, !condition.isAll {
                     let modelExistsWithCondition = try exists(modelSchema, withId: model.id, predicate: condition)
                     if !modelExistsWithCondition {
                         let dataStoreError = DataStoreError.invalidCondition(

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Storage/StorageEngine.swift

@@ -193,7 +193,8 @@ final class StorageEngine: StorageEngineBehavior {
 
         let mutationType = modelExists ? MutationEvent.MutationType.update : .create
 
-        if mutationType == .create && condition != nil {
+        // If it is `create`, and there is a condition, and that condition is not `.all`, fail the request
+        if mutationType == .create, let condition = condition, !condition.isAll {
             let dataStoreError = DataStoreError.invalidCondition(
                 "Cannot apply a condition on model which does not exist.",
                 "Save the model instance without a condition first.")

AmplifyPlugins/DataStore/AWSDataStoreCategoryPluginTests/Core/SQLiteStorageEngineAdapterTests.swift

@@ -57,6 +57,45 @@ class SQLiteStorageEngineAdapterTests: BaseDataStoreTests {
         wait(for: [expectation], timeout: 5)
     }
 
+    /// - Given: a list a `Post` instance
+    /// - When:
+    ///   - the `save(post, condition = .all)` is called
+    /// - Then:
+    ///   - call `query(Post)` to check if the model was correctly inserted
+    func testInsertPostWithAll() {
+        let expectation = self.expectation(
+            description: "it should save and select a Post from the database")
+
+        // insert a post
+        let post = Post(title: "title", content: "content", createdAt: .now())
+        storageAdapter.save(post, condition: QueryPredicateConstant.all) { saveResult in
+            switch saveResult {
+            case .success:
+                self.storageAdapter.query(Post.self) { queryResult in
+                    switch queryResult {
+                    case .success(let posts):
+                        XCTAssert(posts.count == 1)
+                        if let savedPost = posts.first {
+                            XCTAssert(post.id == savedPost.id)
+                            XCTAssert(post.title == savedPost.title)
+                            XCTAssert(post.content == savedPost.content)
+                            XCTAssertEqual(post.createdAt.iso8601String, savedPost.createdAt.iso8601String)
+                        }
+                        expectation.fulfill()
+                    case .failure(let error):
+                        XCTFail(String(describing: error))
+                        expectation.fulfill()
+                    }
+                }
+            case .failure(let error):
+                XCTFail(String(describing: error))
+                expectation.fulfill()
+            }
+        }
+
+        wait(for: [expectation], timeout: 5)
+    }
+
     /// - Given: a list a `Post` instance
     /// - When:
     ///   - the `save(post)` is called
@@ -198,6 +237,55 @@ class SQLiteStorageEngineAdapterTests: BaseDataStoreTests {
         wait(for: [expectation], timeout: 5)
     }
 
+    /// - Given: A Post instance
+    /// - When:
+    ///    - The `save(post)` is called
+    /// - Then:
+    ///    - call `update(post, condition = .all)` with `post.title` updated and condition `.all`
+    ///    - a successful update for `update(post, condition)`
+    ///    - call `query(Post)` to check if the model was correctly updated
+    func testInsertPostAndThenUpdateItWithConditionAll() {
+        let expectation = self.expectation(
+            description: "it should insert and update a Post")
+
+        func checkSavedPost(id: String) {
+            storageAdapter.query(Post.self) {
+                switch $0 {
+                case .success(let posts):
+                    XCTAssertEqual(posts.count, 1)
+                    if let post = posts.first {
+                        XCTAssertEqual(post.id, id)
+                        XCTAssertEqual(post.title, "title updated")
+                    }
+                    expectation.fulfill()
+                case .failure(let error):
+                    XCTFail(String(describing: error))
+                    expectation.fulfill()
+                }
+            }
+        }
+
+        var post = Post(title: "title", content: "content", createdAt: .now())
+        storageAdapter.save(post) { insertResult in
+            switch insertResult {
+            case .success:
+                post.title = "title updated"
+                self.storageAdapter.save(post, condition: QueryPredicateConstant.all) { updateResult in
+                    switch updateResult {
+                    case .success:
+                        checkSavedPost(id: post.id)
+                    case .failure(let error):
+                        XCTFail(error.errorDescription)
+                    }
+                }
+            case .failure(let error):
+                XCTFail(String(describing: error))
+            }
+        }
+
+        wait(for: [expectation], timeout: 5)
+    }
+
     /// - Given: A Post instance
     /// - When:
     ///    - The `save(post, condition)` is called, condition is passed in.

AmplifyPlugins/DataStore/AWSDataStoreCategoryPluginTests/Storage/StorageEngineTestsHasOne.swift

@@ -48,6 +48,28 @@ class StorageEngineTestsHasOne: StorageEngineTestsBase {
         }
     }
 
+    func testSaveModelWithPredicateAll() {
+        let team = Team(name: "Team")
+        let saveFinished = expectation(description: "Save finished")
+        var result: DataStoreResult<Team>?
+        storageEngine.save(team, condition: QueryPredicateConstant.all) { sResult in
+            result = sResult
+            saveFinished.fulfill()
+        }
+        wait(for: [saveFinished], timeout: defaultTimeout)
+
+        guard let saveResult = result else {
+            XCTFail("Save operation timed out")
+            return
+        }
+
+        guard case .success =
+            querySingleModelSynchronous(modelType: Team.self, predicate: Team.keys.id == team.id) else {
+                XCTFail("Failed to query Team")
+                return
+        }
+    }
+
     func testBelongsToRelationshipWithoutOwner() {
         let teamA = Team(name: "A-Team")
         let projectA = Project(name: "ProjectA", team: teamA)