Add rocksdb backend

This commit is contained in:
lash 2022-04-20 10:48:34 +00:00
parent 2f7508ad6e
commit df6e56f4b2
Signed by: lash
GPG Key ID: 21D2E7BB88C2A746
6 changed files with 223 additions and 1 deletions

View File

@ -1,3 +1,5 @@
- 0.2.1
* Add rocksdb backend
- 0.2.0
* Add redis backend
* UTC timestamp for modification time in core state

View File

@ -1,3 +1,8 @@
from setuptools import setup
setup()
setup(
extras_require={
'redis': 'redis==3.5.3',
'rocksdb': 'lbry-rocksdb==0.8.2',
},
)

View File

@ -1,3 +1,6 @@
# standard imports
import datetime
# external imports
import redis
@ -9,6 +12,7 @@ class RedisStore:
self.__path = path
self.__binary = binary
def __to_path(self, k):
return '.'.join([self.__path, k])

121
shep/store/rocksdb.py Normal file
View File

@ -0,0 +1,121 @@
# standard imports
import datetime
# external imports
import rocksdb
class RocksDbStore:
def __init__(self, path, db, binary=False):
self.db = db
self.__path = path
self.__binary = binary
def __to_key(self, k):
return k.encode('utf-8')
def __to_contents(self, v):
if isinstance(v, bytes):
return v
return v.encode('utf-8')
def __to_path(self, k):
return '.'.join([self.__path, k])
def __from_path(self, s):
(left, right) = s.split('.', maxsplit=1)
return right
def __to_result(self, v):
if self.__binary:
return v
return v.decode('utf-8')
def add(self, k, contents=b''):
if contents == None:
contents = b''
else:
contents = self.__to_contents(contents)
k = self.__to_path(k)
k = self.__to_key(k)
self.db.put(k, contents)
def remove(self, k):
k = self.__to_path(k)
k = self.__to_key(k)
self.db.delete(k)
def get(self, k):
k = self.__to_path(k)
k = self.__to_key(k)
v = self.db.get(k)
return self.__to_result(v)
def list(self):
it = self.db.iteritems()
kb_start = self.__to_key(self.__path)
it.seek(kb_start)
r = []
l = len(self.__path)
for (k, v) in it:
if len(k) < l or k[:l] != self.__path:
break
k = self.__from_path(s)
kb = self.__to_key(k)
v = self.db.get(kb)
r.append((k, v,))
return r
def path(self):
return None
def replace(self, k, contents):
if contents == None:
contents = b''
else:
contents = self.__to_contents(contents)
k = self.__to_path(k)
k = self.__to_key(k)
v = self.db.get(k)
if v == None:
raise FileNotFoundError(k)
self.db.put(k, contents)
def modified(self, k):
k = self.__to_path(k)
k = '_mod' + k
v = self.db.get(k)
return int(v)
def register_modify(self, k):
k = self.__to_path(k)
k = '_mod' + k
ts = datetime.datetime.utcnow().timestamp()
self.db.set(k)
class RocksDbStoreFactory:
def __init__(self, path, binary=False):
self.db = rocksdb.DB(path, rocksdb.Options(create_if_missing=True))
self.__binary = binary
def add(self, k):
k = str(k)
return RocksDbStore(k, self.db, binary=self.__binary)

View File

@ -2,6 +2,7 @@
import unittest
import tempfile
import os
import shutil
# local imports
from shep.persist import PersistedState
@ -24,6 +25,10 @@ class TestFileStore(unittest.TestCase):
self.states.add('baz')
def tearDown(self):
shutil.rmtree(self.d)
def test_add(self):
self.states.put('abcd', state=self.states.FOO, contents='baz')
fp = os.path.join(self.d, 'FOO', 'abcd')

85
tests/test_rocksdb.py Normal file
View File

@ -0,0 +1,85 @@
# standard imports
import unittest
import os
import logging
import sys
import importlib
import tempfile
import shutil
# local imports
from shep.persist import PersistedState
from shep.error import (
StateExists,
StateInvalid,
StateItemExists,
StateItemNotFound,
)
logging.basicConfig(level=logging.DEBUG)
logg = logging.getLogger()
class TestRedisStore(unittest.TestCase):
def setUp(self):
from shep.store.rocksdb import RocksDbStoreFactory
self.d = tempfile.mkdtemp()
self.factory = RocksDbStoreFactory(self.d)
self.states = PersistedState(self.factory.add, 3)
self.states.add('foo')
self.states.add('bar')
self.states.add('baz')
def tearDown(self):
shutil.rmtree(self.d)
def test_add(self):
self.states.put('abcd', state=self.states.FOO, contents='baz')
v = self.states.get('abcd')
self.assertEqual(v, 'baz')
v = self.states.state('abcd')
self.assertEqual(v, self.states.FOO)
def test_next(self):
self.states.put('abcd')
self.states.next('abcd')
self.assertEqual(self.states.state('abcd'), self.states.FOO)
self.states.next('abcd')
self.assertEqual(self.states.state('abcd'), self.states.BAR)
self.states.next('abcd')
self.assertEqual(self.states.state('abcd'), self.states.BAZ)
with self.assertRaises(StateInvalid):
self.states.next('abcd')
v = self.states.state('abcd')
self.assertEqual(v, self.states.BAZ)
def test_replace(self):
with self.assertRaises(StateItemNotFound):
self.states.replace('abcd', contents='foo')
self.states.put('abcd', state=self.states.FOO, contents='baz')
self.states.replace('abcd', contents='bar')
v = self.states.get('abcd')
self.assertEqual(v, 'bar')
if __name__ == '__main__':
norocksdb = False
rocksdb = None
try:
importlib.import_module('rocksdb')
except ModuleNotFoundError:
logg.critical('rocksdb module not available, skipping tests.')
sys.exit(0)
unittest.main()