diff --git a/src/blockdb/__init__.py b/src/blockdb/__init__.py new file mode 100644 index 00000000..83d84ea0 --- /dev/null +++ b/src/blockdb/__init__.py @@ -0,0 +1,16 @@ +from typing import Generator + +from onionrblocks import Block + +import db +from .. import identifyhome + +block_db_path = identifyhome.identify_home() + 'blocks.db' + + +def get_blocks_by_type(block_type) -> Generator[Block]: + block_db = db.get_db_obj(block_db_path, 'u') + for block_hash in db.list_keys(block_db_path): + block = Block(block_hash, block_db[block_hash], auto_verify=False) + if block.type == block_type: + yield block diff --git a/src/db/__init__.py b/src/db/__init__.py index ef5daed9..1e9a4a86 100644 --- a/src/db/__init__.py +++ b/src/db/__init__.py @@ -23,6 +23,7 @@ def _do_timeout(func, *args): def set(db_path, key, value): + """Set a value in the db, open+timeout so not good for rapid use""" def _set(key, value): with dbm.open(db_path, "c") as my_db: my_db[key] = value @@ -30,8 +31,38 @@ def set(db_path, key, value): def get(db_path, key): + """Get a value in the db, open+timeout so not good for rapid use""" def _get(key): - with dbm.open(db_path, "c") as my_db: + with dbm.open(db_path, "cu") as my_db: return my_db[key] return _do_timeout(_get, key) + + +def get_db_obj(db_path, extra_flag=''): + """For when you should keep a db obj open""" + def _get_db(): + return dbm.open(db_path, "c" + extra_flag) + return _do_timeout(_get_db, db_path) + + +def list_keys(db_path): + """Generator of all keys in the db. + + Uses a lot of mem if no firstkey supported""" + db_obj = _do_timeout(dbm.open, db_path, "cu") + if not hasattr(db_obj, "firstkey"): + for i in db_obj.keys(): + yield i + db_obj.close() + return + + + def _list_keys(db_obj): + with db_obj as my_db: + k = my_db.firstkey() + while k is not None: + yield k + k = my_db.nextkey(k) + yield from _list_keys(db_obj) +