def reload(self, opts: Dict = None, read_only: bool = None):
if opts is None:
opts = self._opts
if read_only is None:
read_only = self._read_only
rocks_opts = Options(create_if_missing=True)
for key, value in opts.items():
setattr(rocks_opts, key, value)
if hasattr(self, '_db'):
self.close()
self._db = DB(self._path, rocks_opts, read_only=read_only)
def _visible_items(self, db: DB) -> Iterator[Tuple[bytes, bytes]]:
it = db.iteritems() # noqa: B301
it.seek_to_first()
for i, x in enumerate(it):
print('IT %r %r' % (i, x))
for key, value in it:
if key != self.offset_key:
yield key, value
def _visible_items(self, db: DB) -> Iterator[Tuple[bytes, bytes]]:
it = db.iteritems() # noqa: B301
it.seek_to_first()
for key, value in it:
if key != self.offset_key:
yield key, value
def _visible_keys(self, db: DB) -> Iterator[bytes]:
it = db.iterkeys() # noqa: B301
it.seek_to_first()
for key in it:
if key != self.offset_key:
yield key
class RocksDB:
ALLOWED_KEY_TYPES = set([bytes, type(None)])
def __init__(self,
path: str = './rocksdb',
opts: Dict = None,
read_only: bool = False):
self._path = path
if opts is None:
opts = {}
self._opts = opts
self._read_only = read_only
self.reload(read_only=read_only)
@property
def path(self) -> str:
return self._path
@property
def opts(self) -> Dict:
return dict(self._opts)
@property
def read_only(self) -> bool:
return self._read_only
@property
def db(self) -> DB:
return self._db
def reload(self, opts: Dict = None, read_only: bool = None):
if opts is None:
opts = self._opts
if read_only is None:
read_only = self._read_only
rocks_opts = Options(create_if_missing=True)
for key, value in opts.items():
setattr(rocks_opts, key, value)
if hasattr(self, '_db'):
self.close()
self._db = DB(self._path, rocks_opts, read_only=read_only)
def put(self, key: bytes, value, write_batch=None):
if type(key) not in self.ALLOWED_KEY_TYPES:
raise TypeError
if value is None:
raise ValueError
value_bytes = utils.pack(value)
if write_batch is not None:
write_batch.put(key, value_bytes)
else:
self._db.put(key, value_bytes, sync=True)
def get(self, key: bytes):
if type(key) not in self.ALLOWED_KEY_TYPES:
raise TypeError
value_bytes = self._db.get(key)
if value_bytes is not None:
return utils.unpack(value_bytes)
def exists(self, key: bytes) -> bool:
if self.get(key) is not None:
return True
return False
def delete(self, key: bytes, write_batch: WriteBatch = None):
if type(key) not in self.ALLOWED_KEY_TYPES:
raise TypeError
if write_batch is not None:
write_batch.delete(key)
else:
self._db.delete(key, sync=True)
def commit(self, write_batch: WriteBatch):
if write_batch is None:
raise ValueError
self._db.write(write_batch, sync=True)
def scan(self,
prefix: bytes = None,
start_key: bytes = None,
stop_key: bytes = None,
reversed_scan: bool = False) -> Generator:
for key in [prefix, start_key, stop_key]:
if type(key) not in self.ALLOWED_KEY_TYPES:
raise TypeError
iterator = self._db.iterkeys()
if prefix is None and start_key is None:
if reversed_scan:
iterator.seek_to_last()
else:
iterator.seek_to_first()
elif prefix is not None:
iterator.seek(prefix)
else:
iterator.seek(start_key)
if reversed_scan:
iterator = reversed(iterator)
for key in iterator:
if prefix is not None and key[:len(prefix)] != prefix:
return
if stop_key is not None and key > stop_key:
return
value_bytes = self._db.get(key)
value = utils.unpack(value_bytes)
yield key, value
def close(self):
del self._db
gc.collect()
def backup(self, path: str = './rocksdb-backup'):
backup_engine = BackupEngine(path)
backup_engine.create_backup(self._db, flush_before_backup=True)
@staticmethod
def backup_info(path: str):
backup_engine = BackupEngine(path)
return backup_engine.get_backup_info()
@staticmethod
def delete_backup(path: str, backup_id: int = None):
backup_engine = BackupEngine(path)
return backup_engine.delete_backup(backup_id)
def restore_backup(self, path: str, backup_id: int = None):
self.close()
backup_engine = BackupEngine(path)
if backup_id is None:
backup_engine.restore_latest_backup(self._path, self._path)
else:
backup_engine.restore_backup(backup_id, self._path, self._path)
self.reload()
import os
import discord
from rocksdb import DB, Options
from characterhandler import CharacterHandler
from commandhandler import CommandHandler
from conditionhandler import ConditionHandler
from inspirehandler import InspireHandler
from kv import KV
from rollhandler import RollHandler
if __name__ == '__main__':
client = discord.Client()
rocks_db_location = os.getenv('ROCKS_DB_LOCATION', 'anathema.db')
kv = KV(DB(rocks_db_location, Options(create_if_missing=True)))
handlers = [
RollHandler(kv),
InspireHandler(),
CharacterHandler(kv),
ConditionHandler()
]
command_handler = CommandHandler(kv, handlers.copy())
handlers.append(command_handler)
@client.event
async def on_ready():
"""Handles bot start"""
