stored_dict: add support for write batches

Add unit test of atomicity.
Use a write batch for DB upgrades.

Author: ecdsa

electrum/json_db.py

@@ -26,6 +26,7 @@
 import copy
 import json
 from typing import TYPE_CHECKING, Optional, Sequence, List, Union, Dict, Any
+from contextlib import contextmanager
 
 import jsonpatch
 import jsonpointer
@@ -94,6 +95,7 @@ def __init__(
         self._is_closed = True
         self.lock = threading.RLock()
         self.pending_changes = []  # type: List[str]
+        self._write_batch = False
         self._modified = False
         if self.path:
             self.storage = FileStorage(path, allow_partial_writes=allow_partial_writes)
@@ -318,8 +320,20 @@ def dump(self, *, human_readable: bool = True) -> str:
             sort_keys=bool(human_readable),
         )
 
+    @contextmanager
+    def write_batch(self):
+        # FIXME: if an exception is raised during the context, changes will not
+        # be written to disk, but they will be added to the in-memory dict.
+        assert self._write_batch is False
+        self._write_batch = True
+        yield
+        self._write_batch = False
+        if self.storage:
+            self.write()
+
     @locked
     def write(self):
+        assert self._write_batch is False
         if not self.storage:
             return
         if self.storage.should_do_full_write_next():
@@ -328,7 +342,7 @@ def write(self):
             self._append_pending_changes()
 
     def close(self):
-        # do not call write
+        # do not call write, because we may need to close the DB after an exception was raised during a batch write
         self._is_closed = True
 
     def is_closed(self):

electrum/stored_dict.py

@@ -290,6 +290,9 @@ def __init__(self, db: BaseDB, key: _FLEX_KEY, parent):
     def should_convert(self):
         return self._db._should_convert
 
+    def write_batch(self):
+        return self._db.write_batch()
+
     def dump(self) -> dict:
         data = {}
         for k, v in self.items():

electrum/wallet.py

@@ -4510,8 +4510,6 @@ def restore_wallet_from_text(
         if gap_limit_for_change is not None:
             db.put('gap_limit_for_change', gap_limit_for_change)
         wallet = wallet_factory(db, config=config)
-    if storage:
-        assert not storage.file_exists(), "file was created too soon! plaintext keys might have been written to disk"
     wallet.synchronize()
     msg = ("This wallet was restored offline. It may contain more addresses than displayed. "
            "Start a daemon and use load_wallet to sync its history.")

electrum/wallet_db.py

@@ -1562,14 +1562,15 @@ def upgrade_wallet_db(data: 'StoredDict', do_upgrade: bool) -> Tuple[dict, bool]
         )
 
     data._db._should_convert = False
-    dbu = WalletDBUpgrader(data)
-    if dbu.requires_split():
-        raise WalletRequiresSplit(dbu.get_split_accounts())
-    if dbu.requires_upgrade() and do_upgrade:
-        dbu.upgrade()
-        was_upgraded = True
-    if dbu.requires_upgrade():
-        raise WalletRequiresUpgrade()
+    with data.write_batch():
+        dbu = WalletDBUpgrader(data)
+        if dbu.requires_split():
+            raise WalletRequiresSplit(dbu.get_split_accounts())
+        if dbu.requires_upgrade() and do_upgrade:
+            dbu.upgrade()
+            was_upgraded = True
+        if dbu.requires_upgrade():
+            raise WalletRequiresUpgrade()
     data._db._should_convert = True
     return was_upgraded
 

tests/test_stored_dict.py

@@ -58,6 +58,28 @@ def test_db_iterators(self):
         for i, v in enumerate(sl):
             self.assertEqual(i, v)
 
+    def test_write_batch(self):
+        # test that batches are written atomically
+        sd = DictStorage(self.path)
+        with sd.write_batch():
+            sd['a'] = 0
+        self.assertEqual(len(sd), 1)
+        with sd.write_batch():
+            sd['a'] = 1
+        self.assertEqual(len(sd), 1)
+        try:
+            with sd.write_batch():
+                sd['b'] = 1
+                raise Exception('blah')
+        except Exception as e:
+            pass
+        # at this point, the StoredDict length is 2
+        self.assertEqual(len(sd), 2)
+        sd.close()
+        # check that changes have not been written to disk
+        sd = DictStorage(self.path)
+        self.assertEqual(len(sd), 1)
+
     async def test_dangling_dict(self):
         storage = DictStorage(self.path)
         storage['a'] = {'b': {'c': 0}}