shep

Multi-state key stores using bit masks for python3
git clone git://git.defalsify.org/shep.git
Log | Files | Refs | LICENSE

rocksdb.py (3110B)

      1 # standard imports
      2 import datetime
      3 import os
      4 
      5 # external imports
      6 import rocksdb
      7 
      8 # local imports
      9 from .base import StoreFactory
     10 
     11 
     12 class RocksDbStore:
     13 
     14     def __init__(self, path, db, binary=False):
     15         self.db = db
     16         self.__path = path
     17         self.__binary = binary
     18 
     19 
     20     def __to_key(self, k):
     21         return k.encode('utf-8')
     22 
     23 
     24     def __to_contents(self, v):
     25         if isinstance(v, bytes):
     26             return v
     27         return v.encode('utf-8')
     28 
     29 
     30     def __to_path(self, k):
     31         return '.'.join([self.__path, k])
     32 
     33 
     34     def __from_path(self, s):
     35         (left, right) = s.split('.', maxsplit=1)
     36         return right
     37 
     38 
     39     def __to_result(self, v):
     40         if self.__binary:
     41             return v
     42         return v.decode('utf-8')
     43 
     44 
     45     def put(self, k, contents=b''):
     46         if contents == None:
     47             contents = b''
     48         else:
     49             contents = self.__to_contents(contents)
     50         k = self.__to_path(k)
     51         k = self.__to_key(k)
     52         self.db.put(k, contents)
     53 
     54 
     55     def remove(self, k):
     56         k = self.__to_path(k)
     57         k = self.__to_key(k)
     58         self.db.delete(k)
     59 
     60 
     61     def get(self, k):
     62         k = self.__to_path(k)
     63         k = self.__to_key(k)
     64         v = self.db.get(k)
     65         return self.__to_result(v)
     66 
     67  
     68     def list(self):
     69         it = self.db.iteritems()
     70         kb_start = self.__to_key(self.__path)
     71         it.seek(kb_start)
     72 
     73         r = []
     74         l = len(self.__path)
     75         for (kb, v) in it:
     76             k = kb.decode('utf-8') 
     77             if len(k) < l or k[:l] != self.__path:
     78                 break
     79             k = self.__from_path(k)
     80             v = self.db.get(kb)
     81             r.append((k, v,))
     82 
     83         return r
     84 
     85 
     86     def path(self):
     87         return None
     88 
     89 
     90     def replace(self, k, contents):
     91         if contents == None:
     92             contents = b''
     93         else:
     94             contents = self.__to_contents(contents)
     95         k = self.__to_path(k)
     96         k = self.__to_key(k)
     97         v = self.db.get(k)
     98         if v == None:
     99             raise FileNotFoundError(k)
    100         self.db.put(k, contents)
    101 
    102 
    103     def modified(self, k):
    104         k = self.__to_path(k)
    105         k = '_mod' + k
    106         v = self.db.get(k)
    107         return int(v)
    108 
    109 
    110     def register_modify(self, k):
    111         k = self.__to_path(k)
    112         k = '_mod' + k
    113         ts = datetime.datetime.utcnow().timestamp()
    114         self.db.set(k)
    115 
    116 
    117 class RocksDbStoreFactory(StoreFactory):
    118 
    119     def __init__(self, path, binary=False):
    120         try:
    121             os.stat(path)
    122         except FileNotFoundError:
    123             os.makedirs(path)
    124         self.db = rocksdb.DB(path, rocksdb.Options(create_if_missing=True))
    125         self.__binary = binary
    126 
    127 
    128     def add(self, k):
    129         k = str(k)
    130         return RocksDbStore(k, self.db, binary=self.__binary)
    131 
    132 
    133     def close(self):
    134         self.db.close()
    135 
    136 
    137     def ls(self):
    138         it = self.db.iterkeys()
    139         it.seek_to_first()
    140         r = []
    141         for k in it:
    142             v = k.rsplit(b'.', maxsplit=1)
    143             if v != k:
    144                 v = v[0].decode('utf-8')
    145                 if v not in r:
    146                     r.append(v)
    147         return r