Merge pull request #2015 from Expensify/main

sqlitecluster/SQLite.cpp

@@ -566,18 +566,35 @@ bool SQLite::write(const string& query) {
     }
 
     // This is literally identical to the idempotent version except for the check for _noopUpdateMode.
-    return _writeIdempotent(query);
+    SQResult ignore;
+    return _writeIdempotent(query, ignore);
+}
+
+bool SQLite::write(const string& query, SQResult& result) {
+    if (_noopUpdateMode) {
+        SALERT("Non-idempotent write in _noopUpdateMode. Query: " << query);
+        return true;
+    }
+
+    // This is literally identical to the idempotent version except for the check for _noopUpdateMode.
+    return _writeIdempotent(query, result);
 }
 
 bool SQLite::writeIdempotent(const string& query) {
-    return _writeIdempotent(query);
+    SQResult ignore;
+    return _writeIdempotent(query, ignore);
+}
+
+bool SQLite::writeIdempotent(const string& query, SQResult& result) {
+    return _writeIdempotent(query, result);
 }
 
 bool SQLite::writeUnmodified(const string& query) {
-    return _writeIdempotent(query, true);
+    SQResult ignore;
+    return _writeIdempotent(query, ignore, true);
 }
 
-bool SQLite::_writeIdempotent(const string& query, bool alwaysKeepQueries) {
+bool SQLite::_writeIdempotent(const string& query, SQResult& result, bool alwaysKeepQueries) {
     SASSERT(_insideTransaction);
     _queryCache.clear();
     _queryCount++;
@@ -600,7 +617,7 @@ bool SQLite::_writeIdempotent(const string& query, bool alwaysKeepQueries) {
     {
         shared_lock<shared_mutex> lock(_sharedData.writeLock);
         if (_enableRewrite) {
-            resultCode = SQuery(_db, "read/write transaction", query, 2'000'000, true);
+            resultCode = SQuery(_db, "read/write transaction", query, result, 2'000'000, true);
             if (resultCode == SQLITE_AUTH) {
                 // Run re-written query.
                 _currentlyRunningRewritten = true;
@@ -610,7 +627,7 @@ bool SQLite::_writeIdempotent(const string& query, bool alwaysKeepQueries) {
                 _currentlyRunningRewritten = false;
             }
         } else {
-            resultCode = SQuery(_db, "read/write transaction", query);
+            resultCode = SQuery(_db, "read/write transaction", query, result);
         }
     }
 

sqlitecluster/SQLite.h

@@ -1,5 +1,6 @@
 #pragma once
 #include <libstuff/sqlite3.h>
+#include <libstuff/SQResult.h>
 #include <libstuff/SPerformanceTimer.h>
 
 class SQLite {
@@ -100,15 +101,23 @@ class SQLite {
     bool addColumn(const string& tableName, const string& column, const string& columnType);
 
     // Performs a read/write query (eg, INSERT, UPDATE, DELETE). This is added to the current transaction's query list.
-    // Returns true  on success.
+    // Returns true on success.
     // If we're in noop-update mode, this call alerts and performs no write, but returns as if it had completed.
     bool write(const string& query);
 
+    // Performs a read/write query
+    // Designed for use with queries that include a RETURNING clause
+    bool write(const string& query, SQResult& result);
+
     // This is the same as `write` except it runs successfully without any warnings or errors in noop-update mode.
     // It's intended to be used for `mockRequest` enabled commands, such that we only run a version of them that's
     // known to be repeatable. What counts as repeatable is up to the individual command.
     bool writeIdempotent(const string& query);
 
+    // Executes a write query and retrieves the result.
+    // Designed for use with queries that include a RETURNING clause
+    bool writeIdempotent(const string& query, SQResult& result);
+
     // This runs a query completely unchanged, always adding it to the uncommitted query, such that it will be recorded
     // in the journal even if it had no effect on the database. This lets replicated or synchronized queries be added
     // to the journal *even if they have no effect* on the rest of the database.
@@ -400,7 +409,7 @@ class SQLite {
     static thread_local int64_t _conflictPage;
     static thread_local string _conflictTable;
 
-    bool _writeIdempotent(const string& query, bool alwaysKeepQueries = false);
+    bool _writeIdempotent(const string& query, SQResult& result, bool alwaysKeepQueries = false);
 
     // Constructs a UNION query from a list of 'query parts' over each of our journal tables.
     // Fore each table, queryParts will be joined with that table's name as a separator. I.e., if you have a tables

test/tests/LibStuffTest.cpp

@@ -35,6 +35,7 @@ struct LibStuff : tpunit::TestFixture {
                                     TEST(LibStuff::SREMatchTest),
                                     TEST(LibStuff::SREReplaceTest),
                                     TEST(LibStuff::SQResultTest),
+                                    TEST(LibStuff::testReturningClause),
                                     TEST(LibStuff::SRedactSensitiveValuesTest)
                                     )
     { }
@@ -746,6 +747,37 @@ struct LibStuff : tpunit::TestFixture {
         ASSERT_EQUAL(result[0]["coco"], "name1");
     }
 
+    void testReturningClause() {
+        // Given a sqlite DB with a table and pre inserted values
+        SQLite db(":memory:", 1000, 1000, 1);
+        db.beginTransaction(SQLite::TRANSACTION_TYPE::EXCLUSIVE);
+        db.write("CREATE TABLE testReturning(id INTEGER PRIMARY KEY, name STRING, value STRING, created DATE);");
+        db.write("INSERT INTO testReturning VALUES(11, 'name1', 'value1', '2024-12-02');");
+        db.write("INSERT INTO testReturning VALUES(21, 'name2', 'value2', '2024-12-03');");
+        db.prepare();
+        db.commit();
+
+        // When trying to delete a row by returning the deleted items
+        db.beginTransaction(SQLite::TRANSACTION_TYPE::SHARED);
+        SQResult result;
+        db.writeIdempotent("DELETE FROM testReturning WHERE id = 21 AND name = 'name2' RETURNING id, name;", result);
+        db.prepare();
+        db.commit();
+
+        // Verify that deleted items were returned as expected
+        ASSERT_EQUAL("21", result[0][0]);
+        ASSERT_EQUAL("name2", result[0][1]);
+
+        // Verify that the row was successfully deleted and now the table has only one row
+        db.beginTransaction(SQLite::TRANSACTION_TYPE::SHARED);
+        db.read("SELECT name, value FROM testReturning ORDER BY id;", result);
+        db.rollback();
+
+        ASSERT_EQUAL(1, result.size());
+        ASSERT_EQUAL(result[0]["name"], "name1");
+        ASSERT_EQUAL(result[0]["value"], "value1");
+    }
+
     void SRedactSensitiveValuesTest() {
         string logValue = R"({"edits":["test1", "test2", "test3"]})";
         SRedactSensitiveValues(logValue);