commit df6e56f4b2aabc904ec2e35a293272191cf55429
parent 2f7508ad6e9ac79366548d2d9d5eeca00ba9e223
Author: lash <dev@holbrook.no>
Date: Wed, 20 Apr 2022 10:48:34 +0000
Add rocksdb backend
Diffstat:
6 files changed, 223 insertions(+), 1 deletion(-)
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,3 +1,5 @@
+- 0.2.1
+ * Add rocksdb backend
- 0.2.0
* Add redis backend
* UTC timestamp for modification time in core state
diff --git a/setup.py b/setup.py
@@ -1,3 +1,8 @@
from setuptools import setup
-setup()
+setup(
+ extras_require={
+ 'redis': 'redis==3.5.3',
+ 'rocksdb': 'lbry-rocksdb==0.8.2',
+ },
+ )
diff --git a/shep/store/redis.py b/shep/store/redis.py
@@ -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])
diff --git a/shep/store/rocksdb.py b/shep/store/rocksdb.py
@@ -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)
diff --git a/tests/test_file.py b/tests/test_file.py
@@ -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')
diff --git a/tests/test_rocksdb.py b/tests/test_rocksdb.py
@@ -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()
