Add rocksdb backend

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

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',
+        },
+        )

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])
 

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)

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')

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()