wallet: Add DatabaseBatch::ErasePrefix method

This new function is not used yet this commit, but next commit adds usages and
test coverage for both BDB and sqlite.

Author: ryanofsky

src/wallet/bdb.cpp

@@ -665,12 +665,14 @@ void BerkeleyDatabase::ReloadDbEnv()
     env->ReloadDbEnv();
 }
 
-BerkeleyCursor::BerkeleyCursor(BerkeleyDatabase& database)
+BerkeleyCursor::BerkeleyCursor(BerkeleyDatabase& database, BerkeleyBatch* batch)
 {
     if (!database.m_db.get()) {
         throw std::runtime_error(STR_INTERNAL_BUG("BerkeleyDatabase does not exist"));
     }
-    int ret = database.m_db->cursor(nullptr, &m_cursor, 0);
+    // Transaction argument to cursor is only needed when using the cursor to
+    // write to the database. Read-only cursors do not need a txn pointer.
+    int ret = database.m_db->cursor(batch ? batch->txn() : nullptr, &m_cursor, 0);
     if (ret != 0) {
         throw std::runtime_error(STR_INTERNAL_BUG(strprintf("BDB Cursor could not be created. Returned %d", ret)));
     }
@@ -817,6 +819,25 @@ bool BerkeleyBatch::HasKey(DataStream&& key)
     return ret == 0;
 }
 
+bool BerkeleyBatch::ErasePrefix(Span<const std::byte> prefix)
+{
+    if (!TxnBegin()) return false;
+    auto cursor{std::make_unique<BerkeleyCursor>(m_database, this)};
+    // const_cast is safe below even though prefix_key is an in/out parameter,
+    // because we are not using the DB_DBT_USERMEM flag, so BDB will allocate
+    // and return a different output data pointer
+    Dbt prefix_key{const_cast<std::byte*>(prefix.data()), static_cast<uint32_t>(prefix.size())}, prefix_value{};
+    int ret{cursor->dbc()->get(&prefix_key, &prefix_value, DB_SET_RANGE)};
+    for (int flag{DB_CURRENT}; ret == 0; flag = DB_NEXT) {
+        SafeDbt key, value;
+        ret = cursor->dbc()->get(key, value, flag);
+        if (ret != 0 || key.get_size() < prefix.size() || memcmp(key.get_data(), prefix.data(), prefix.size()) != 0) break;
+        ret = cursor->dbc()->del(0);
+    }
+    cursor.reset();
+    return TxnCommit() && (ret == 0 || ret == DB_NOTFOUND);
+}
+
 void BerkeleyDatabase::AddRef()
 {
     LOCK(cs_db);

src/wallet/bdb.h

@@ -191,10 +191,11 @@ class BerkeleyCursor : public DatabaseCursor
     Dbc* m_cursor;
 
 public:
-    explicit BerkeleyCursor(BerkeleyDatabase& database);
+    explicit BerkeleyCursor(BerkeleyDatabase& database, BerkeleyBatch* batch=nullptr);
     ~BerkeleyCursor() override;
 
     Status Next(DataStream& key, DataStream& value) override;
+    Dbc* dbc() const { return m_cursor; }
 };
 
 /** RAII class that provides access to a Berkeley database */
@@ -205,6 +206,7 @@ class BerkeleyBatch : public DatabaseBatch
     bool WriteKey(DataStream&& key, DataStream&& value, bool overwrite = true) override;
     bool EraseKey(DataStream&& key) override;
     bool HasKey(DataStream&& key) override;
+    bool ErasePrefix(Span<const std::byte> prefix) override;
 
 protected:
     Db* pdb{nullptr};
@@ -229,6 +231,7 @@ class BerkeleyBatch : public DatabaseBatch
     bool TxnBegin() override;
     bool TxnCommit() override;
     bool TxnAbort() override;
+    DbTxn* txn() const { return activeTxn; }
 };
 
 std::string BerkeleyDatabaseVersion();

src/wallet/db.h

@@ -110,6 +110,7 @@ class DatabaseBatch
 
         return HasKey(std::move(ssKey));
     }
+    virtual bool ErasePrefix(Span<const std::byte> prefix) = 0;
 
     virtual std::unique_ptr<DatabaseCursor> GetNewCursor() = 0;
     virtual bool TxnBegin() = 0;
@@ -186,6 +187,7 @@ class DummyBatch : public DatabaseBatch
     bool WriteKey(DataStream&& key, DataStream&& value, bool overwrite = true) override { return true; }
     bool EraseKey(DataStream&& key) override { return true; }
     bool HasKey(DataStream&& key) override { return true; }
+    bool ErasePrefix(Span<const std::byte> prefix) override { return true; }
 
 public:
     void Flush() override {}

src/wallet/sqlite.cpp

@@ -125,6 +125,7 @@ void SQLiteBatch::SetupSQLStatements()
         {&m_insert_stmt, "INSERT INTO main VALUES(?, ?)"},
         {&m_overwrite_stmt, "INSERT or REPLACE into main values(?, ?)"},
         {&m_delete_stmt, "DELETE FROM main WHERE key = ?"},
+        {&m_delete_prefix_stmt, "DELETE FROM main WHERE instr(key, ?) = 1"},
     };
 
     for (const auto& [stmt_prepared, stmt_text] : statements) {
@@ -375,6 +376,7 @@ void SQLiteBatch::Close()
         {&m_insert_stmt, "insert"},
         {&m_overwrite_stmt, "overwrite"},
         {&m_delete_stmt, "delete"},
+        {&m_delete_prefix_stmt, "delete prefix"},
     };
 
     for (const auto& [stmt_prepared, stmt_description] : statements) {
@@ -441,24 +443,34 @@ bool SQLiteBatch::WriteKey(DataStream&& key, DataStream&& value, bool overwrite)
     return res == SQLITE_DONE;
 }
 
-bool SQLiteBatch::EraseKey(DataStream&& key)
+bool SQLiteBatch::ExecStatement(sqlite3_stmt* stmt, Span<const std::byte> blob)
 {
     if (!m_database.m_db) return false;
-    assert(m_delete_stmt);
+    assert(stmt);
 
     // Bind: leftmost parameter in statement is index 1
-    if (!BindBlobToStatement(m_delete_stmt, 1, key, "key")) return false;
+    if (!BindBlobToStatement(stmt, 1, blob, "key")) return false;
 
     // Execute
-    int res = sqlite3_step(m_delete_stmt);
-    sqlite3_clear_bindings(m_delete_stmt);
-    sqlite3_reset(m_delete_stmt);
+    int res = sqlite3_step(stmt);
+    sqlite3_clear_bindings(stmt);
+    sqlite3_reset(stmt);
     if (res != SQLITE_DONE) {
         LogPrintf("%s: Unable to execute statement: %s\n", __func__, sqlite3_errstr(res));
     }
     return res == SQLITE_DONE;
 }
 
+bool SQLiteBatch::EraseKey(DataStream&& key)
+{
+    return ExecStatement(m_delete_stmt, key);
+}
+
+bool SQLiteBatch::ErasePrefix(Span<const std::byte> prefix)
+{
+    return ExecStatement(m_delete_prefix_stmt, prefix);
+}
+
 bool SQLiteBatch::HasKey(DataStream&& key)
 {
     if (!m_database.m_db) return false;

src/wallet/sqlite.h

@@ -36,13 +36,16 @@ class SQLiteBatch : public DatabaseBatch
     sqlite3_stmt* m_insert_stmt{nullptr};
     sqlite3_stmt* m_overwrite_stmt{nullptr};
     sqlite3_stmt* m_delete_stmt{nullptr};
+    sqlite3_stmt* m_delete_prefix_stmt{nullptr};
 
     void SetupSQLStatements();
+    bool ExecStatement(sqlite3_stmt* stmt, Span<const std::byte> blob);
 
     bool ReadKey(DataStream&& key, DataStream& value) override;
     bool WriteKey(DataStream&& key, DataStream&& value, bool overwrite = true) override;
     bool EraseKey(DataStream&& key) override;
     bool HasKey(DataStream&& key) override;
+    bool ErasePrefix(Span<const std::byte> prefix) override;
 
 public:
     explicit SQLiteBatch(SQLiteDatabase& database);

src/wallet/test/wallet_tests.cpp

@@ -922,6 +922,7 @@ class FailBatch : public DatabaseBatch
     bool WriteKey(DataStream&& key, DataStream&& value, bool overwrite = true) override { return m_pass; }
     bool EraseKey(DataStream&& key) override { return m_pass; }
     bool HasKey(DataStream&& key) override { return m_pass; }
+    bool ErasePrefix(Span<const std::byte> prefix) override { return m_pass; }
 
 public:
     explicit FailBatch(bool pass) : m_pass(pass) {}