Separate WalletDB from storage upgrades.

Make sure that WalletDB.data is always a StoredDict.
Perform db upgrades in a separate class, since they
operate on a dict object.

electrum/daemon.py

@@ -47,7 +47,7 @@ from .util import log_exceptions, ignore_exceptions, randrange, OldTaskGroup
 from .util import EventListener, event_listener
 from .wallet import Wallet, Abstract_Wallet
 from .storage import WalletStorage
-from .wallet_db import WalletDB
+from .wallet_db import WalletDB, WalletRequiresSplit, WalletRequiresUpgrade
 from .commands import known_commands, Commands
 from .simple_config import SimpleConfig
 from .exchange_rate import FxThread
@@ -512,10 +512,11 @@ class Daemon(Logger):
                 return
             storage.decrypt(password)
         # read data, pass it to db
-        db = WalletDB(storage.read(), storage=storage, manual_upgrades=manual_upgrades)
-        if db.requires_split():
+        try:
+            db = WalletDB(storage.read(), storage=storage, manual_upgrades=manual_upgrades)
+        except WalletRequiresSplit:
             return
-        if db.requires_upgrade():
+        except WalletRequiresUpgrade:
             return
         if db.get_action():
             return

electrum/gui/qml/qewalletdb.py

@@ -5,7 +5,7 @@ from PyQt5.QtCore import pyqtProperty, pyqtSignal, pyqtSlot, QObject
 from electrum.i18n import _
 from electrum.logging import get_logger
 from electrum.storage import WalletStorage
-from electrum.wallet_db import WalletDB
+from electrum.wallet_db import WalletDB, WalletRequiresSplit
 from electrum.wallet import Wallet
 from electrum.util import InvalidPassword, WalletFileException, send_exception_to_crash_reporter
 
@@ -143,7 +143,10 @@ class QEWalletDB(QObject):
         else:  # storage not encrypted; but it might still have a keystore pw
             # FIXME hack... load both db and full wallet, just to tell if it has keystore pw.
             #       this also completely ignores db.requires_split(), db.get_action(), etc
-            db = WalletDB(self._storage.read(), storage=self._storage, manual_upgrades=False)
+            try:
+                db = WalletDB(self._storage.read(), storage=self._storage, manual_upgrades=False)
+            except WalletRequiresSplit as e:
+                raise WalletFileException(_('This wallet requires to be split. This is currently not supported on mobile'))
             wallet = Wallet(db, config=self._config)
             self.needsPassword = wallet.has_password()
             if self.needsPassword:
@@ -162,18 +165,14 @@ class QEWalletDB(QObject):
     def _load_db(self):
         """can raise WalletFileException"""
         # needs storage accessible
-        self._db = WalletDB(self._storage.read(), storage=self._storage, manual_upgrades=True)
-        if self._db.requires_split():
+        try:
+            self._db = WalletDB(self._storage.read(), storage=self._storage, manual_upgrades=False)
+        except WalletRequiresSplit as e:
             self._logger.warning('wallet requires split')
             raise WalletFileException(_('This wallet needs splitting. This is not supported on mobile'))
         if self._db.get_action():
             self._logger.warning('action pending. QML version doesn\'t support continuation of wizard')
             raise WalletFileException(_('This wallet has an action pending. This is currently not supported on mobile'))
 
-        if self._db.requires_upgrade():
-            self._logger.warning('wallet requires upgrade, upgrading')
-            self._db.upgrade()
-            self._db.write()
-
         self._ready = True
         self.readyChanged.emit()

electrum/gui/qt/__init__.py

@@ -61,7 +61,7 @@ from electrum.plugin import run_hook
 from electrum.util import (UserCancelled, profiler, send_exception_to_crash_reporter,
                            WalletFileException, BitcoinException, get_new_wallet_name)
 from electrum.wallet import Wallet, Abstract_Wallet
-from electrum.wallet_db import WalletDB
+from electrum.wallet_db import WalletDB, WalletRequiresSplit, WalletRequiresUpgrade
 from electrum.logging import Logger
 from electrum.gui import BaseElectrumGui
 from electrum.simple_config import SimpleConfig
@@ -430,10 +430,11 @@ class ElectrumGui(BaseElectrumGui, Logger):
         if storage.is_encrypted_with_user_pw() or storage.is_encrypted_with_hw_device():
             storage.decrypt(d['password'])
 
-        db = WalletDB(storage.read(), storage=storage, manual_upgrades=True)
-        if db.requires_split() or db.requires_upgrade():
+        try:
+            db = WalletDB(storage.read(), storage=storage, manual_upgrades=True)
+        except WalletRequiresSplit as e:
             try:
-                wizard.run_upgrades(db)
+                wizard.run_split(storage, e._split_data)
             except UserCancelled:
                 return
 

electrum/gui/qt/wizard/wallet.py

@@ -2,6 +2,7 @@ import os
 import sys
 import threading
 import time
+import json
 
 from typing import TYPE_CHECKING
 
@@ -21,6 +22,7 @@ from electrum.wallet import wallet_types
 from .wizard import QEAbstractWizard, WizardComponent
 from electrum.logging import get_logger, Logger
 from electrum import WalletStorage, mnemonic, keystore
+from electrum.wallet_db import WalletDB
 from electrum.wizard import NewWalletWizard
 
 from electrum.gui.qt.bip39_recovery_dialog import Bip39RecoveryDialog
@@ -34,7 +36,6 @@ if TYPE_CHECKING:
     from electrum.simple_config import SimpleConfig
     from electrum.plugin import Plugins
     from electrum.daemon import Daemon
-    from electrum.wallet_db import WalletDB
     from electrum.gui.qt import QElectrumApplication
 
 WIF_HELP_TEXT = (_('WIF keys are typed in Electrum, based on script type.') + '\n\n' +
@@ -161,25 +162,20 @@ class QENewWalletWizard(NewWalletWizard, QEAbstractWizard, MessageBoxMixin):
         self._password = data['password']
         self.path = path
 
-    def run_upgrades(self, db: 'WalletDB') -> None:
-        storage = db.storage
-        path = storage.path
-        if db.requires_split():
-            msg = _(
-                "The wallet '{}' contains multiple accounts, which are no longer supported since Electrum 2.7.\n\n"
-                "Do you want to split your wallet into multiple files?").format(path)
-            if not self.question(msg):
-                return
-            file_list = db.split_accounts(path)
+    def run_split(self, storage, split_data) -> None:
+        root_path = storage.path
+        msg = _(
+            "The wallet '{}' contains multiple accounts, which are no longer supported since Electrum 2.7.\n\n"
+            "Do you want to split your wallet into multiple files?").format(root_path)
+        if self.question(msg):
+            file_list = WalletDB.split_accounts(root_path, split_data)
             msg = _('Your accounts have been moved to') + ':\n' + '\n'.join(file_list) + '\n\n' + _(
-                'Do you want to delete the old file') + ':\n' + path
+                'Do you want to delete the old file') + ':\n' + root_path
             if self.question(msg):
-                os.remove(path)
+                os.remove(root_path)
                 self.show_warning(_('The file was removed'))
-            # raise now, to avoid having the old storage opened
-            raise UserCancelled()
-        if db.requires_upgrade():
-            self.waiting_dialog(db.upgrade, _('Upgrading wallet format...'))
+        # raise now, to avoid having the old storage opened
+        raise UserCancelled()
 
     def is_finalized(self, wizard_data: dict) -> bool:
         # check decryption of existing wallet and keep wizard open if incorrect.

electrum/json_db.py

@@ -34,8 +34,6 @@ from .logging import Logger
 if TYPE_CHECKING:
     from .storage import WalletStorage
 
-JsonDBJsonEncoder = util.MyEncoder
-
 def modifier(func):
     def wrapper(self, *args, **kwargs):
         with self.lock:
@@ -168,24 +166,28 @@ class StoredDict(dict):
 
 class JsonDB(Logger):
 
-    def __init__(self, data, storage=None):
+    def __init__(self, s: str, storage=None, encoder=None):
         Logger.__init__(self)
         self.lock = threading.RLock()
         self.storage = storage
+        self.encoder = encoder
         self._modified = False
         # load data
-        if data:
-            self.load_data(data)
-        else:
-            self.data = {}
-
-    def load_data(self, s):
+        data = self.load_data(s)
+        # convert to StoredDict
+        self.data = StoredDict(data, self, [])
+
+    def load_data(self, s:str) -> dict:
+        """ overloaded in wallet_db """
+        if s == '':
+            return {}
         try:
-            self.data = json.loads(s)
+            data = json.loads(s)
         except Exception:
             raise WalletFileException("Cannot read wallet file. (parsing failed)")
-        if not isinstance(self.data, dict):
+        if not isinstance(data, dict):
             raise WalletFileException("Malformed wallet file (not dict)")
+        return data
 
     def set_modified(self, b):
         with self.lock:
@@ -204,8 +206,8 @@ class JsonDB(Logger):
     @modifier
     def put(self, key, value):
         try:
-            json.dumps(key, cls=JsonDBJsonEncoder)
-            json.dumps(value, cls=JsonDBJsonEncoder)
+            json.dumps(key, cls=self.encoder)
+            json.dumps(value, cls=self.encoder)
         except Exception:
             self.logger.info(f"json error: cannot save {repr(key)} ({repr(value)})")
             return False
@@ -235,7 +237,7 @@ class JsonDB(Logger):
             self.data,
             indent=4 if human_readable else None,
             sort_keys=bool(human_readable),
-            cls=JsonDBJsonEncoder,
+            cls=self.encoder,
         )
 
     def _should_convert_to_stored_dict(self, key) -> bool:

electrum/lnwatcher.py

@@ -325,7 +325,7 @@ class WatchTower(LNWatcher):
     LOGGING_SHORTCUT = 'W'
 
     def __init__(self, network: 'Network'):
-        adb = AddressSynchronizer(WalletDB({}, storage=None, manual_upgrades=False), network.config, name=self.diagnostic_name())
+        adb = AddressSynchronizer(WalletDB('', storage=None, manual_upgrades=False), network.config, name=self.diagnostic_name())
         adb.start_network(network)
         LNWatcher.__init__(self, adb, network)
         self.network = network

electrum/tests/test_storage_upgrade.py

@@ -7,7 +7,7 @@ import asyncio
 import inspect
 
 import electrum
-from electrum.wallet_db import WalletDB
+from electrum.wallet_db import WalletDBUpgrader, WalletDB, WalletRequiresUpgrade, WalletRequiresSplit
 from electrum.wallet import Wallet
 from electrum import constants
 from electrum import util
@@ -331,31 +331,30 @@ class TestStorageUpgrade(WalletTestCase):
     async def _upgrade_storage(self, wallet_json, accounts=1) -> Optional[WalletDB]:
         if accounts == 1:
             # test manual upgrades
-            db = self._load_db_from_json_string(wallet_json=wallet_json,
-                                                manual_upgrades=True)
-            self.assertFalse(db.requires_split())
-            if db.requires_upgrade():
-                db.upgrade()
+            try:
+                db = self._load_db_from_json_string(
+                    wallet_json=wallet_json,
+                    manual_upgrades=True)
+            except WalletRequiresUpgrade:
+                db = self._load_db_from_json_string(
+                    wallet_json=wallet_json,
+                    manual_upgrades=False)
                 await self._sanity_check_upgraded_db(db)
-            # test automatic upgrades
-            db2 = self._load_db_from_json_string(wallet_json=wallet_json,
-                                                 manual_upgrades=False)
-            await self._sanity_check_upgraded_db(db2)
-            return db2
+            return db
         else:
-            db = self._load_db_from_json_string(wallet_json=wallet_json,
-                                                manual_upgrades=True)
-            self.assertTrue(db.requires_split())
-            split_data = db.get_split_accounts()
-            self.assertEqual(accounts, len(split_data))
-            for item in split_data:
-                data = json.dumps(item)
-                new_db = WalletDB(data, storage=None, manual_upgrades=False)
-                await self._sanity_check_upgraded_db(new_db)
+            try:
+                db = self._load_db_from_json_string(
+                    wallet_json=wallet_json,
+                    manual_upgrades=True)
+            except WalletRequiresSplit as e:
+                split_data = e._split_data
+                self.assertEqual(accounts, len(split_data))
+                for item in split_data:
+                    data = json.dumps(item)
+                    new_db = WalletDB(data, storage=None, manual_upgrades=False)
+                    await self._sanity_check_upgraded_db(new_db)
 
     async def _sanity_check_upgraded_db(self, db):
-        self.assertFalse(db.requires_split())
-        self.assertFalse(db.requires_upgrade())
         wallet = Wallet(db, config=self.config)
         await wallet.stop()
 

electrum/tests/test_wallet.py

@@ -15,7 +15,7 @@ from electrum.wallet import (Abstract_Wallet, Standard_Wallet, create_new_wallet
 from electrum.exchange_rate import ExchangeBase, FxThread
 from electrum.util import TxMinedInfo, InvalidPassword
 from electrum.bitcoin import COIN
-from electrum.wallet_db import WalletDB
+from electrum.wallet_db import WalletDB, JsonDB
 from electrum.simple_config import SimpleConfig
 from electrum import util
 
@@ -60,14 +60,14 @@ class TestWalletStorage(WalletTestCase):
             contents = f.write(contents)
 
         storage = WalletStorage(self.wallet_path)
-        db = WalletDB(storage.read(), storage=storage, manual_upgrades=True)
+        db = JsonDB(storage.read(), storage=storage)
         self.assertEqual("b", db.get("a"))
         self.assertEqual("d", db.get("c"))
 
     def test_write_dictionary_to_file(self):
 
         storage = WalletStorage(self.wallet_path)
-        db = WalletDB('', storage=storage, manual_upgrades=True)
+        db = JsonDB('', storage=storage)
 
         some_dict = {
             u"a": u"b",
@@ -110,7 +110,7 @@ class FakeWallet:
     def __init__(self, fiat_value):
         super().__init__()
         self.fiat_value = fiat_value
-        self.db = WalletDB("{}", storage=None, manual_upgrades=True)
+        self.db = WalletDB('', storage=None, manual_upgrades=True)
         self.adb = FakeADB()
         self.db.transactions = self.db.verified_tx = {'abc':'Tx'}
 

electrum/wallet.py

@@ -309,8 +309,8 @@ class Abstract_Wallet(ABC, Logger, EventListener):
 
     def __init__(self, db: WalletDB, *, config: SimpleConfig):
 
-        if not db.is_ready_to_be_used_by_wallet():
-            raise Exception("storage not ready to be used by Abstract_Wallet")
+        #if not db.is_ready_to_be_used_by_wallet():
+        #    raise Exception("storage not ready to be used by Abstract_Wallet")
 
         self.config = config
         assert self.config is not None, "config must not be None"

electrum/wallet_db.py

@@ -36,7 +36,7 @@ import time
 import attr
 
 from . import util, bitcoin
-from .util import profiler, WalletFileException, multisig_type, TxMinedInfo, bfh
+from .util import profiler, WalletFileException, multisig_type, TxMinedInfo, bfh, MyEncoder
 from .invoices import Invoice, Request
 from .keystore import bip44_derivation
 from .transaction import Transaction, TxOutpoint, tx_from_any, PartialTransaction, PartialTxOutput
@@ -49,6 +49,11 @@ from .plugin import run_hook, plugin_loaders
 from .version import ELECTRUM_VERSION
 
 
+class WalletRequiresUpgrade(WalletFileException):
+    pass
+class WalletRequiresSplit(WalletFileException):
+    def __init__(self, split_data):
+        self._split_data = split_data
 
 # seed_version is now used for the version of the wallet file
 
@@ -100,46 +105,21 @@ for key in ['locked_in', 'fails', 'settles']:
     json_db.register_parent_key(key, lambda x: HTLCOwner(int(x)))
 
 
-class WalletDB(JsonDB):
 
-    def __init__(self, data, *, storage=None, manual_upgrades: bool):
-        JsonDB.__init__(self, data, storage)
-        if not data:
-            # create new DB
-            self.put('seed_version', FINAL_SEED_VERSION)
-            self._add_db_creation_metadata()
-            self._after_upgrade_tasks()
-        self._manual_upgrades = manual_upgrades
-        self._called_after_upgrade_tasks = False
-        if not self._manual_upgrades and self.requires_split():
-            raise WalletFileException("This wallet has multiple accounts and must be split")
-        if not self.requires_upgrade():
-            self._after_upgrade_tasks()
-        elif not self._manual_upgrades:
-            self.upgrade()
-        # load plugins that are conditional on wallet type
-        self.load_plugins()
+class WalletDBUpgrader(Logger):
 
-    def load_data(self, s):
-        try:
-            JsonDB.load_data(self, s)
-        except Exception:
-            try:
-                d = ast.literal_eval(s)
-                labels = d.get('labels', {})
-            except Exception as e:
-                raise WalletFileException("Cannot read wallet file. (parsing failed)")
-            self.data = {}
-            for key, value in d.items():
-                try:
-                    json.dumps(key)
-                    json.dumps(value)
-                except Exception:
-                    self.logger.info(f'Failed to convert label to json format: {key}')
-                    continue
-                self.data[key] = value
-        if not isinstance(self.data, dict):
-            raise WalletFileException("Malformed wallet file (not dict)")
+    def __init__(self, data):
+        Logger.__init__(self)
+        self.data = data
+
+    def get(self, key, default=None):
+        return self.data.get(key, default)
+
+    def put(self, key, value):
+        if value is not None:
+            self.data[key] = value
+        else:
+            self.data.pop(key, None)
 
     def requires_split(self):
         d = self.get('accounts', {})
@@ -192,9 +172,6 @@ class WalletDB(JsonDB):
     @profiler
     def upgrade(self):
         self.logger.info('upgrading wallet format')
-        if self._called_after_upgrade_tasks:
-            # we need strict ordering between upgrade() and after_upgrade_tasks()
-            raise Exception("'after_upgrade_tasks' must NOT be called before 'upgrade'")
         self._convert_imported()
         self._convert_wallet_type()
         self._convert_account()
@@ -242,12 +219,6 @@ class WalletDB(JsonDB):
         self._convert_version_54()
         self.put('seed_version', FINAL_SEED_VERSION)  # just to be sure
 
-        self._after_upgrade_tasks()
-
-    def _after_upgrade_tasks(self):
-        self._called_after_upgrade_tasks = True
-        self._load_transactions()
-
     def _convert_wallet_type(self):
         if not self._is_upgrade_method_needed(0, 13):
             return
@@ -1140,7 +1111,6 @@ class WalletDB(JsonDB):
         else:
             return True
 
-    @locked
     def get_seed_version(self):
         seed_version = self.get('seed_version')
         if not seed_version:
@@ -1191,14 +1161,62 @@ class WalletDB(JsonDB):
         # generic exception
         raise WalletFileException(msg)
 
-    def _add_db_creation_metadata(self):
-        # store this for debugging purposes
-        v = DBMetadata(
-            creation_timestamp=int(time.time()),
-            first_electrum_version_used=ELECTRUM_VERSION,
-        )
-        assert self.get("db_metadata", None) is None
-        self.put("db_metadata", v)
+
+class WalletDB(JsonDB):
+
+    def __init__(self, s, *, storage=None, manual_upgrades=True):
+        self._upgrade = not manual_upgrades
+        JsonDB.__init__(self, s, storage, encoder=MyEncoder)
+        # create pointers
+        self.load_transactions()
+        # load plugins that are conditional on wallet type
+        self.load_plugins()
+
+    def load_data(self, s):
+        try:
+            data = JsonDB.load_data(self, s)
+        except Exception:
+            try:
+                d = ast.literal_eval(s)
+                labels = d.get('labels', {})
+            except Exception as e:
+                raise WalletFileException("Cannot read wallet file. (parsing failed)")
+            data = {}
+            for key, value in d.items():
+                try:
+                    json.dumps(key)
+                    json.dumps(value)
+                except Exception:
+                    self.logger.info(f'Failed to convert label to json format: {key}')
+                    continue
+                data[key] = value
+        if not isinstance(data, dict):
+            raise WalletFileException("Malformed wallet file (not dict)")
+
+        if len(data) == 0:
+            # create new DB
+            data['seed_version'] = FINAL_SEED_VERSION
+            # store this for debugging purposes
+            v = DBMetadata(
+                creation_timestamp=int(time.time()),
+                first_electrum_version_used=ELECTRUM_VERSION,
+            )
+            assert data.get("db_metadata", None) is None
+            data["db_metadata"] = v
+
+        dbu = WalletDBUpgrader(data)
+        if dbu.requires_split():
+            raise WalletRequiresSplit(dbu.get_split_accounts())
+        if self._upgrade:
+            dbu.upgrade()
+        if dbu.requires_upgrade():
+            raise WalletRequiresUpgrade()
+        return dbu.data
+
+
+    @locked
+    def get_seed_version(self):
+        return self.get('seed_version')
 
     def get_db_metadata(self) -> Optional[DBMetadata]:
         # field only present for wallet files created with ver 4.4.0 or later
@@ -1560,8 +1578,7 @@ class WalletDB(JsonDB):
                 self._addr_to_addr_index[addr] = (1, i)
 
     @profiler
-    def _load_transactions(self):
-        self.data = StoredDict(self.data, self, [])
+    def load_transactions(self):
         # references in self.data
         # TODO make all these private
         # txid -> address -> prev_outpoint -> value
@@ -1608,21 +1625,19 @@ class WalletDB(JsonDB):
         return True
 
     def is_ready_to_be_used_by_wallet(self):
-        return not self.requires_upgrade() and self._called_after_upgrade_tasks
+        return not self._requires_upgrade
 
-    def split_accounts(self, root_path):
+    @classmethod
+    def split_accounts(klass, root_path, split_data):
         from .storage import WalletStorage
-        out = []
-        result = self.get_split_accounts()
-        for data in result:
+        file_list = []
+        for data in split_data:
             path = root_path + '.' + data['suffix']
-            storage = WalletStorage(path)
-            db = WalletDB(json.dumps(data), storage=storage, manual_upgrades=False)
-            db._called_after_upgrade_tasks = False
-            db.upgrade()
+            item_storage = WalletStorage(path)
+            db = WalletDB(json.dumps(data), storage=item_storage, manual_upgrades=False)
             db.write()
-            out.append(path)
-        return out
+            file_list.append(path)
+        return file_list
 
     def get_action(self):
         action = run_hook('get_action', self)