コード例 #1
0
ファイル: rocksdb.py プロジェクト: pwilczynskiclearcode/faust
 def _visible_keys(self, db: DB) -> Iterator[bytes]:
     it = db.iterkeys()  # noqa: B301
     it.seek_to_first()
     for key in it:
         if key != self.offset_key:
             yield key
コード例 #2
0
ファイル: easyrocks.py プロジェクト: labteral/easyrocks
class RocksDB:

    ALLOWED_KEY_TYPES = set([bytes, type(None)])

    def __init__(self,
                 path: str = './rocksdb',
                 opts: Dict = None,
                 read_only: bool = False):
        self._path = path

        if opts is None:
            opts = {}
        self._opts = opts

        self._read_only = read_only
        self.reload(read_only=read_only)

    @property
    def path(self) -> str:
        return self._path

    @property
    def opts(self) -> Dict:
        return dict(self._opts)

    @property
    def read_only(self) -> bool:
        return self._read_only

    @property
    def db(self) -> DB:
        return self._db

    def reload(self, opts: Dict = None, read_only: bool = None):
        if opts is None:
            opts = self._opts

        if read_only is None:
            read_only = self._read_only

        rocks_opts = Options(create_if_missing=True)
        for key, value in opts.items():
            setattr(rocks_opts, key, value)

        if hasattr(self, '_db'):
            self.close()

        self._db = DB(self._path, rocks_opts, read_only=read_only)

    def put(self, key: bytes, value, write_batch=None):
        if type(key) not in self.ALLOWED_KEY_TYPES:
            raise TypeError

        if value is None:
            raise ValueError

        value_bytes = utils.pack(value)

        if write_batch is not None:
            write_batch.put(key, value_bytes)
        else:
            self._db.put(key, value_bytes, sync=True)

    def get(self, key: bytes):
        if type(key) not in self.ALLOWED_KEY_TYPES:
            raise TypeError

        value_bytes = self._db.get(key)

        if value_bytes is not None:
            return utils.unpack(value_bytes)

    def exists(self, key: bytes) -> bool:
        if self.get(key) is not None:
            return True
        return False

    def delete(self, key: bytes, write_batch: WriteBatch = None):
        if type(key) not in self.ALLOWED_KEY_TYPES:
            raise TypeError

        if write_batch is not None:
            write_batch.delete(key)
        else:
            self._db.delete(key, sync=True)

    def commit(self, write_batch: WriteBatch):
        if write_batch is None:
            raise ValueError
        self._db.write(write_batch, sync=True)

    def scan(self,
             prefix: bytes = None,
             start_key: bytes = None,
             stop_key: bytes = None,
             reversed_scan: bool = False) -> Generator:

        for key in [prefix, start_key, stop_key]:
            if type(key) not in self.ALLOWED_KEY_TYPES:
                raise TypeError

        iterator = self._db.iterkeys()
        if prefix is None and start_key is None:
            if reversed_scan:
                iterator.seek_to_last()
            else:
                iterator.seek_to_first()
        elif prefix is not None:
            iterator.seek(prefix)
        else:
            iterator.seek(start_key)

        if reversed_scan:
            iterator = reversed(iterator)

        for key in iterator:
            if prefix is not None and key[:len(prefix)] != prefix:
                return

            if stop_key is not None and key > stop_key:
                return

            value_bytes = self._db.get(key)
            value = utils.unpack(value_bytes)

            yield key, value

    def close(self):
        del self._db
        gc.collect()

    def backup(self, path: str = './rocksdb-backup'):
        backup_engine = BackupEngine(path)
        backup_engine.create_backup(self._db, flush_before_backup=True)

    @staticmethod
    def backup_info(path: str):
        backup_engine = BackupEngine(path)
        return backup_engine.get_backup_info()

    @staticmethod
    def delete_backup(path: str, backup_id: int = None):
        backup_engine = BackupEngine(path)
        return backup_engine.delete_backup(backup_id)

    def restore_backup(self, path: str, backup_id: int = None):
        self.close()
        backup_engine = BackupEngine(path)
        if backup_id is None:
            backup_engine.restore_latest_backup(self._path, self._path)
        else:
            backup_engine.restore_backup(backup_id, self._path, self._path)
        self.reload()