You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
52 lines
1.5 KiB
52 lines
1.5 KiB
from __future__ import absolute_import
|
|
|
|
import six
|
|
import rocksdb
|
|
|
|
from itertools import izip
|
|
from six.moves import map
|
|
|
|
from disco.util.serializer import Serializer
|
|
from .base import BaseProvider, SEP_SENTINEL
|
|
|
|
|
|
class RocksDBProvider(BaseProvider):
|
|
def __init__(self, config):
|
|
super(RocksDBProvider, self).__init__(config)
|
|
self.format = config.get('format', 'pickle')
|
|
self.path = config.get('path', 'storage.db')
|
|
self.db = None
|
|
|
|
@staticmethod
|
|
def k(k):
|
|
return bytes(k) if six.PY3 else str(k.encode('utf-8'))
|
|
|
|
def load(self):
|
|
self.db = rocksdb.DB(self.path, rocksdb.Options(create_if_missing=True))
|
|
|
|
def exists(self, key):
|
|
return self.db.get(self.k(key)) is not None
|
|
|
|
# TODO prefix extractor
|
|
def keys(self, other):
|
|
count = other.count(SEP_SENTINEL) + 1
|
|
it = self.db.iterkeys()
|
|
it.seek_to_first()
|
|
|
|
for key in it:
|
|
key = key.decode('utf-8')
|
|
if key.startswith(other) and key.count(SEP_SENTINEL) == count:
|
|
yield key
|
|
|
|
def get_many(self, keys):
|
|
for key, value in izip(keys, self.db.multi_get(list(map(self.k, keys)))):
|
|
yield (key, Serializer.loads(self.format, value.decode('utf-8')))
|
|
|
|
def get(self, key):
|
|
return Serializer.loads(self.format, self.db.get(self.k(key)).decode('utf-8'))
|
|
|
|
def set(self, key, value):
|
|
self.db.put(self.k(key), Serializer.dumps(self.format, value))
|
|
|
|
def delete(self, key):
|
|
self.db.delete(self.k(key))
|
|
|