lbry-sdk/lbry/wallet/server/db/interface.py

import struct
import typing

import rocksdb
from typing import Optional
from lbry.wallet.server.db import DB_PREFIXES
from lbry.wallet.server.db.revertable import RevertableOpStack, RevertablePut, RevertableDelete


class BasePrefixDB:
    """
    Base class for a revertable rocksdb database (a rocksdb db where each set of applied changes can be undone)
    """
    UNDO_KEY_STRUCT = struct.Struct(b'>Q32s')
    PARTIAL_UNDO_KEY_STRUCT = struct.Struct(b'>Q')

    def __init__(self, path, max_open_files=64, secondary_path='', max_undo_depth: int = 200, unsafe_prefixes=None):
        column_family_options = {
                prefix.value: rocksdb.ColumnFamilyOptions() for prefix in DB_PREFIXES
        }
        self.column_families: typing.Dict[bytes, 'rocksdb.ColumnFamilyHandle'] = {}
        self._db = rocksdb.DB(
            path, rocksdb.Options(
                create_if_missing=True, use_fsync=False, target_file_size_base=33554432,
                max_open_files=max_open_files if not secondary_path else -1, create_missing_column_families=True
            ), secondary_name=secondary_path, column_families=column_family_options
        )
        for prefix in DB_PREFIXES:
            cf = self._db.get_column_family(prefix.value)
            if cf is None and not secondary_path:
                self._db.create_column_family(prefix.value, rocksdb.ColumnFamilyOptions())
                cf = self._db.get_column_family(prefix.value)
            self.column_families[prefix.value] = cf

        self._op_stack = RevertableOpStack(self.get, unsafe_prefixes=unsafe_prefixes)
        self._max_undo_depth = max_undo_depth

    def unsafe_commit(self):
        """
        Write staged changes to the database without keeping undo information
        Changes written cannot be undone
        """
        try:
            if not len(self._op_stack):
                return
            with self._db.write_batch(sync=True) as batch:
                batch_put = batch.put
                batch_delete = batch.delete
                get_column_family = self.column_families.__getitem__
                for staged_change in self._op_stack:
                    column_family = get_column_family(DB_PREFIXES(staged_change.key[:1]).value)
                    if staged_change.is_put:
                        batch_put((column_family, staged_change.key), staged_change.value)
                    else:
                        batch_delete((column_family, staged_change.key))
        finally:
            self._op_stack.clear()

    def commit(self, height: int, block_hash: bytes):
        """
        Write changes for a block height to the database and keep undo information so that the changes can be reverted
        """
        undo_ops = self._op_stack.get_undo_ops()
        delete_undos = []
        if height > self._max_undo_depth:
            delete_undos.extend(self._db.iterator(
                start=DB_PREFIXES.undo.value + self.PARTIAL_UNDO_KEY_STRUCT.pack(0),
                iterate_upper_bound=DB_PREFIXES.undo.value + self.PARTIAL_UNDO_KEY_STRUCT.pack(height - self._max_undo_depth),
                include_value=False
            ))
        try:
            undo_c_f = self.column_families[DB_PREFIXES.undo.value]
            with self._db.write_batch(sync=True) as batch:
                batch_put = batch.put
                batch_delete = batch.delete
                get_column_family = self.column_families.__getitem__
                for staged_change in self._op_stack:
                    column_family = get_column_family(DB_PREFIXES(staged_change.key[:1]).value)
                    if staged_change.is_put:
                        batch_put((column_family, staged_change.key), staged_change.value)
                    else:
                        batch_delete((column_family, staged_change.key))
                for undo_to_delete in delete_undos:
                    batch_delete((undo_c_f, undo_to_delete))
                batch_put((undo_c_f, DB_PREFIXES.undo.value + self.UNDO_KEY_STRUCT.pack(height, block_hash)), undo_ops)
        finally:
            self._op_stack.clear()

    def rollback(self, height: int, block_hash: bytes):
        """
        Revert changes for a block height
        """
        undo_key = DB_PREFIXES.undo.value + self.UNDO_KEY_STRUCT.pack(height, block_hash)
        undo_c_f = self.column_families[DB_PREFIXES.undo.value]
        undo_info = self._db.get((undo_c_f, undo_key))
        self._op_stack.apply_packed_undo_ops(undo_info)
        try:
            with self._db.write_batch(sync=True) as batch:
                batch_put = batch.put
                batch_delete = batch.delete
                get_column_family = self.column_families.__getitem__
                for staged_change in self._op_stack:
                    column_family = get_column_family(DB_PREFIXES(staged_change.key[:1]).value)
                    if staged_change.is_put:
                        batch_put((column_family, staged_change.key), staged_change.value)
                    else:
                        batch_delete((column_family, staged_change.key))
                # batch_delete(undo_key)
        finally:
            self._op_stack.clear()

    def get(self, key: bytes, fill_cache: bool = True) -> Optional[bytes]:
        cf = self.column_families[key[:1]]
        return self._db.get((cf, key), fill_cache=fill_cache)

    def iterator(self, start: bytes, column_family: 'rocksdb.ColumnFamilyHandle' = None,
                 iterate_lower_bound: bytes = None, iterate_upper_bound: bytes = None,
                 reverse: bool = False, include_key: bool = True, include_value: bool = True,
                 fill_cache: bool = True, prefix_same_as_start: bool = False, auto_prefix_mode: bool = False):
        return self._db.iterator(
            start=start, column_family=column_family, iterate_lower_bound=iterate_lower_bound,
            iterate_upper_bound=iterate_upper_bound, reverse=reverse, include_key=include_key,
            include_value=include_value, fill_cache=fill_cache, prefix_same_as_start=prefix_same_as_start,
            auto_prefix_mode=auto_prefix_mode
        )

    def close(self):
        self._db.close()

    def try_catch_up_with_primary(self):
        self._db.try_catch_up_with_primary()

    @property
    def closed(self) -> bool:
        return self._db.is_closed

    def stage_raw_put(self, key: bytes, value: bytes):
        self._op_stack.append_op(RevertablePut(key, value))

    def stage_raw_delete(self, key: bytes, value: bytes):
        self._op_stack.append_op(RevertableDelete(key, value))

    def estimate_num_keys(self, column_family: 'rocksdb.ColumnFamilyHandle' = None):
        return int(self._db.get_property(b'rocksdb.estimate-num-keys', column_family).decode())