Merge pull request groue#1858 from groue/dev/issue-1855

Fix upsert for WITHOUT ROWID tables

Author: groue

GRDB/Record/MutablePersistableRecord+Upsert.swift

@@ -383,13 +383,10 @@ extension MutablePersistableRecord {
     mutating func upsertWithCallbacks(_ db: Database)
     throws -> InsertionSuccess
     {
-        let (inserted, _) = try upsertAndFetchWithCallbacks(
+        try upsertWithCallbacks(
             db, onConflict: [],
             updating: .allColumns,
-            doUpdate: nil,
-            selection: [],
-            decode: { _ in /* Nothing to decode */ })
-        return inserted
+            doUpdate: nil)
     }
 
     @inlinable // allow specialization so that empty callbacks are removed
@@ -421,6 +418,109 @@ extension MutablePersistableRecord {
         return success
     }
 
+    @inlinable // allow specialization so that empty callbacks are removed
+    mutating func upsertWithCallbacks(
+        _ db: Database,
+        onConflict conflictTarget: [String],
+        updating strategy: UpsertUpdateStrategy,
+        doUpdate assignments: ((_ excluded: TableAlias<Self>) -> [ColumnAssignment])?
+    ) throws -> InsertionSuccess {
+        try willInsert(db)
+        
+        var inserted: InsertionSuccess?
+        try aroundInsert(db) {
+            inserted = try upsertWithoutCallbacks(
+                db, onConflict: conflictTarget,
+                updating: strategy, doUpdate: assignments)
+            return inserted!
+        }
+        
+        guard let inserted else {
+            try persistenceCallbackMisuse("aroundInsert")
+        }
+        didInsert(inserted)
+        return inserted
+    }
+    
+    /// Executes an `INSERT` statement, and DOES NOT run
+    /// insertion callbacks.
+    @usableFromInline
+    func upsertWithoutCallbacks(
+        _ db: Database,
+        onConflict conflictTarget: [String],
+        updating strategy: UpsertUpdateStrategy,
+        doUpdate assignments: ((_ excluded: TableAlias<Self>) -> [ColumnAssignment])?)
+    throws -> InsertionSuccess
+    {
+        let primaryKey = try db.primaryKey(Self.databaseTableName)
+        
+        if primaryKey.tableHasRowID {
+            let dao = try DAO(db, self)
+            let statement = try dao.upsertStatement(
+                db, onConflict: conflictTarget,
+                updating: strategy,
+                doUpdate: assignments,
+                updateCondition: nil,
+                returning: [Column.rowID])
+            let cursor = try Row.fetchCursor(statement)
+            
+            // Keep cursor alive until we can process the fetched row
+            let rowid: Int64 = try withExtendedLifetime(cursor) { cursor in
+                guard let row = try cursor.next() else {
+                    throw DatabaseError(message: "Insertion failed")
+                }
+                
+                let rowid: Int64 = row[0]
+                
+                // Now that we have fetched the values we need, we could stop
+                // there. But let's make sure we fully consume the cursor
+                // anyway, until SQLITE_DONE. This is necessary, for example,
+                // for upserts in tables that are synchronized with an
+                // FTS5 table.
+                // See <https://github.com/groue/GRDB.swift/issues/1390>
+                while try cursor.next() != nil { }
+                
+                return rowid
+            }
+            
+            // Update the persistenceContainer with the inserted rowid.
+            // This allows the Record class to set its `hasDatabaseChanges` property
+            // to false in its `aroundInsert` callback.
+            var persistenceContainer = dao.persistenceContainer
+            let rowIDColumn = dao.primaryKey.rowIDColumn
+            if let rowIDColumn {
+                persistenceContainer[rowIDColumn] = rowid
+            }
+            
+            let inserted = InsertionSuccess(
+                rowID: rowid,
+                rowIDColumn: rowIDColumn,
+                persistenceContainer: persistenceContainer)
+            return inserted
+        } else {
+            let dao = try DAO(db, self)
+            let statement = try dao.upsertStatement(
+                db, onConflict: conflictTarget,
+                updating: strategy,
+                doUpdate: assignments,
+                updateCondition: nil,
+                returning: [])
+            let cursor = try Row.fetchCursor(statement)
+            
+            // Let's make sure we fully consume the cursor, until SQLITE_DONE.
+            // This is necessary, for example, for upserts in tables that
+            // are synchronized with an FTS5 table.
+            // See <https://github.com/groue/GRDB.swift/issues/1390>
+            while try cursor.next() != nil { }
+
+            let inserted = InsertionSuccess(
+                rowID: 0, // Meaningless in WITHOUT ROWID tables.
+                rowIDColumn: nil,
+                persistenceContainer: dao.persistenceContainer)
+            return inserted
+        }
+    }
+    
     /// Executes an `INSERT RETURNING` statement, and DOES NOT run
     /// insertion callbacks.
     @usableFromInline
@@ -433,52 +533,90 @@ extension MutablePersistableRecord {
         decode: (Row) throws -> T)
     throws -> (InsertionSuccess, T)
     {
-        // Append the rowID to the returned columns
-        let selection = selection + [Column.rowID]
-        
-        let dao = try DAO(db, self)
-        let statement = try dao.upsertStatement(
-            db, onConflict: conflictTarget,
-            updating: strategy,
-            doUpdate: assignments,
-            updateCondition: nil,
-            returning: selection)
-        let cursor = try Row.fetchCursor(statement)
+        let primaryKey = try db.primaryKey(Self.databaseTableName)
 
-        // Keep cursor alive until we can process the fetched row
-        let (rowid, returned): (Int64, T) = try withExtendedLifetime(cursor) { cursor in
-            guard let row = try cursor.next() else {
-                throw DatabaseError(message: "Insertion failed")
+        if primaryKey.tableHasRowID {
+            // Append the rowID to the returned columns
+            let selection = selection + [Column.rowID]
+            
+            let dao = try DAO(db, self)
+            let statement = try dao.upsertStatement(
+                db, onConflict: conflictTarget,
+                updating: strategy,
+                doUpdate: assignments,
+                updateCondition: nil,
+                returning: selection)
+            let cursor = try Row.fetchCursor(statement)
+            
+            // Keep cursor alive until we can process the fetched row
+            let (rowid, returned): (Int64, T) = try withExtendedLifetime(cursor) { cursor in
+                guard let row = try cursor.next() else {
+                    throw DatabaseError(message: "Insertion failed")
+                }
+                
+                // Rowid is the last column
+                let rowid: Int64 = row[row.count - 1]
+                let returned = try decode(row)
+                
+                // Now that we have fetched the values we need, we could stop
+                // there. But let's make sure we fully consume the cursor
+                // anyway, until SQLITE_DONE. This is necessary, for example,
+                // for upserts in tables that are synchronized with an
+                // FTS5 table.
+                // See <https://github.com/groue/GRDB.swift/issues/1390>
+                while try cursor.next() != nil { }
+                
+                return (rowid, returned)
             }
 
-            // Rowid is the last column
-            let rowid: Int64 = row[row.count - 1]
-            let returned = try decode(row)
+            // Update the persistenceContainer with the inserted rowid.
+            // This allows the Record class to set its `hasDatabaseChanges` property
+            // to false in its `aroundInsert` callback.
+            var persistenceContainer = dao.persistenceContainer
+            let rowIDColumn = dao.primaryKey.rowIDColumn
+            if let rowIDColumn {
+                persistenceContainer[rowIDColumn] = rowid
+            }
 
-            // Now that we have fetched the values we need, we could stop
-            // there. But let's make sure we fully consume the cursor
-            // anyway, until SQLITE_DONE. This is necessary, for example,
-            // for upserts in tables that are synchronized with an
-            // FTS5 table.
-            // See <https://github.com/groue/GRDB.swift/issues/1390>
-            while try cursor.next() != nil { }
+            let inserted = InsertionSuccess(
+                rowID: rowid,
+                rowIDColumn: rowIDColumn,
+                persistenceContainer: persistenceContainer)
+            return (inserted, returned)
+        } else {
+            let dao = try DAO(db, self)
+            let statement = try dao.upsertStatement(
+                db, onConflict: conflictTarget,
+                updating: strategy,
+                doUpdate: assignments,
+                updateCondition: nil,
+                returning: selection)
+            let cursor = try Row.fetchCursor(statement)
 
-            return (rowid, returned)
-        }
-        
-        // Update the persistenceContainer with the inserted rowid.
-        // This allows the Record class to set its `hasDatabaseChanges` property
-        // to false in its `aroundInsert` callback.
-        var persistenceContainer = dao.persistenceContainer
-        let rowIDColumn = dao.primaryKey.rowIDColumn
-        if let rowIDColumn {
-            persistenceContainer[rowIDColumn] = rowid
+            // Keep cursor alive until we can process the fetched row
+            let returned: T = try withExtendedLifetime(cursor) { cursor in
+                guard let row = try cursor.next() else {
+                    throw DatabaseError(message: "Insertion failed")
+                }
+                
+                let returned = try decode(row)
+                
+                // Now that we have fetched the values we need, we could stop
+                // there. But let's make sure we fully consume the cursor
+                // anyway, until SQLITE_DONE. This is necessary, for example,
+                // for upserts in tables that are synchronized with an
+                // FTS5 table.
+                // See <https://github.com/groue/GRDB.swift/issues/1390>
+                while try cursor.next() != nil { }
+                
+                return returned
+            }
+            
+            let inserted = InsertionSuccess(
+                rowID: 0, // Meaningless in WITHOUT ROWID tables.
+                rowIDColumn: nil,
+                persistenceContainer: dao.persistenceContainer)
+            return (inserted, returned)
         }
-        
-        let inserted = InsertionSuccess(
-            rowID: rowid,
-            rowIDColumn: rowIDColumn,
-            persistenceContainer: persistenceContainer)
-        return (inserted, returned)
     }
 }

GRDB/Record/PersistableRecord+Upsert.swift

@@ -383,13 +383,10 @@ extension PersistableRecord {
     func upsertWithCallbacks(_ db: Database)
     throws -> InsertionSuccess
     {
-        let (inserted, _) = try upsertAndFetchWithCallbacks(
+        try upsertWithCallbacks(
             db, onConflict: [],
             updating: .allColumns,
-            doUpdate: nil,
-            selection: [],
-            decode: { _ in /* Nothing to decode */ })
-        return inserted
+            doUpdate: nil)
     }
 
     @inlinable // allow specialization so that empty callbacks are removed
@@ -420,4 +417,28 @@ extension PersistableRecord {
         didInsert(success.inserted)
         return success
     }
+    
+    @inlinable // allow specialization so that empty callbacks are removed
+    func upsertWithCallbacks(
+        _ db: Database,
+        onConflict conflictTarget: [String],
+        updating strategy: UpsertUpdateStrategy,
+        doUpdate assignments: ((_ excluded: TableAlias<Self>) -> [ColumnAssignment])?
+    ) throws -> InsertionSuccess {
+        try willInsert(db)
+        
+        var inserted: InsertionSuccess?
+        try aroundInsert(db) {
+            inserted = try upsertWithoutCallbacks(
+                db, onConflict: conflictTarget,
+                updating: strategy, doUpdate: assignments)
+            return inserted!
+        }
+        
+        guard let inserted else {
+            try persistenceCallbackMisuse("aroundInsert")
+        }
+        didInsert(inserted)
+        return inserted
+    }
 }

Tests/GRDBTests/MutablePersistableRecordTests.swift

@@ -2334,6 +2334,60 @@ extension MutablePersistableRecordTests {
         }
     }
 
+    func test_upsert_WITHOUT_ROWID_optimization() throws {
+#if GRDBCUSTOMSQLITE || SQLITE_HAS_CODEC
+        guard Database.sqliteLibVersionNumber >= 3035000 else {
+            throw XCTSkip("UPSERT is not available")
+        }
+#else
+        guard #available(iOS 15, macOS 12, tvOS 15, watchOS 8, *) else {
+            throw XCTSkip("UPSERT is not available")
+        }
+#endif
+        
+        struct MyRecord: Codable, FetchableRecord, MutablePersistableRecord {
+            var id: String
+            var name: String
+            var score: Int?
+        }
+        
+        try makeDatabaseQueue().write { db in
+            try db.execute(sql: """
+                CREATE TABLE myRecord(
+                  id TEXT NOT NULL PRIMARY KEY,
+                  name TEXT NOT NULL UNIQUE,
+                  score INTEGER NOT NULL
+                ) WITHOUT ROWID;
+                """)
+            
+            do {
+                var record = MyRecord(id: "1", name: "foo", score: 1000)
+                try record.upsert(db)
+                
+                XCTAssertEqual(lastSQLQuery, """
+                    INSERT INTO "myRecord" ("id", "name", "score") \
+                    VALUES ('1','foo',1000) \
+                    ON CONFLICT DO UPDATE SET "name" = "excluded"."name", "score" = "excluded"."score"
+                    """)
+            }
+            do {
+                var record = MyRecord(id: "2", name: "foo", score: 2000)
+                let upserted = try record.upsertAndFetch(db)
+                
+                XCTAssertEqual(lastSQLQuery, """
+                    INSERT INTO "myRecord" ("id", "name", "score") \
+                    VALUES ('2','foo',2000) \
+                    ON CONFLICT DO UPDATE SET "name" = "excluded"."name", "score" = "excluded"."score" \
+                    RETURNING *
+                    """)
+                
+                XCTAssertEqual(upserted.id, "1")
+                XCTAssertEqual(upserted.name, "foo")
+                XCTAssertEqual(upserted.score, 2000)
+            }
+        }
+    }
+    
     func test_upsertAndFetch_do_update_set_where_with_default_strategy() throws {
 #if GRDBCUSTOMSQLITE || SQLITE_HAS_CODEC
         guard Database.sqliteLibVersionNumber >= 3035000 else {