Onionr/src/db/__init__.py

69 lines
1.6 KiB
Python
Raw Normal View History

2022-01-06 20:48:22 +00:00
import dbm
import time
import os
timeout = 120
def _do_timeout(func, *args):
ts = 0
res = None
while True:
try:
res = func(*args)
except dbm.error:
if not ts:
ts = time.time()
continue
if time.time() - ts > timeout:
raise TimeoutError()
2022-02-03 18:55:07 +00:00
time.sleep(0.01)
2022-01-06 20:48:22 +00:00
else:
return res
def set(db_path, key, value):
2022-02-03 06:32:26 +00:00
"""Set a value in the db, open+timeout so not good for rapid use"""
2022-01-06 20:48:22 +00:00
def _set(key, value):
with dbm.open(db_path, "c") as my_db:
my_db[key] = value
_do_timeout(_set, key, value)
def get(db_path, key):
2022-02-03 06:32:26 +00:00
"""Get a value in the db, open+timeout so not good for rapid use"""
2022-01-06 20:48:22 +00:00
def _get(key):
2022-02-03 06:32:26 +00:00
with dbm.open(db_path, "cu") as my_db:
2022-01-06 20:48:22 +00:00
return my_db[key]
return _do_timeout(_get, key)
2022-02-03 06:32:26 +00:00
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)
2022-02-03 18:55:07 +00:00
return _do_timeout(_get_db)
2022-02-03 06:32:26 +00:00
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)