def __init__(self, db, queue_size=1000, dumper=None):
"""
The db can be the db file name and also can be the lmdb db object
which is already been opened.
So, it can be a 'str' and also can be a db object
The queue_size is the size of the queue
dumper: If set, the content of the input will be converted to
string by the dumper function, e.g. pickle.dumps or yaml.dump or str
If not set, it will be assumed the content already been string type
"""
# Call the parents __init__ first
super(DbWriteThread, self).__init__()
if type(db) == str:
self.db = lmdb_tools.open(db)
self.need_close = True
else:
self.db = db
self.need_close = False
# Create the queue
self.queue = Queue.Queue(maxsize=queue_size)
# When set to True, the thread will keep runing
# When set to False, the thread will exit
self._thread_state = True
self.dumper = dumper
# Indicate whether the IO of this object is balance
# When put a number to inqueue, +1
# When write a number to outqueue, -1
self.io_balance = 0
# Lock to avoid the disorder of the io_balance
self.mutex = threading.Lock()
def delete_keylist_from_db(db_file, key_list):
"""
Open the db file, delete the keys in the key list
"""
db = lmdb_tools.open(db_file, append=True)
log.info('Deleting entries from \033[0;33m%s\033[0m, \
Db size: \033[0;32m%d\033[0m, list size: \033[0;32m%d\033[0m' %
(db_file, lmdb_tools.get_entries(db), len(key_list)))
lmdb_tools.delete_keylist(db, key_list)
log.info('Finish. DB size: \033[0;32m%d\033[0m' %
lmdb_tools.get_entries(db))
lmdb_tools.close(db)
def append_db(db_file_src, db_file_dst):
"""
Append all contains of db_file_src to db_file_dst
"""
db_src = lmdb_tools.open_ro(db_file_src)
log.info('Src Db \033[0;33m%s\033[0m size: \033[0;32m%d\033[0m' %
(db_file_src, lmdb_tools.get_entries(db_src)))
db_dst = lmdb_tools.open(db_file_dst, append=True)
log.info('Dst Db \033[0;33m%s\033[0m size: \033[0;32m%d\033[0m' %
(db_file_dst, lmdb_tools.get_entries(db_dst)))
with db_src.begin(write=False) as txn_src:
with txn_src.cursor() as cur:
for key, val in cur:
with db_dst.begin(write=True) as txn_dst:
txn_dst.put(key, val)
log.info('Append Finished. Dst Size: \033[0;32m%d\033[0m' %
lmdb_tools.get_entries(db_dst))
lmdb_tools.close(db_src)
lmdb_tools.close(db_dst)
def save_dict_to_db(db_file, key_dict, operation=None, append=False):
"""
Open a new db (if exist, ask whether to delete the old one)
save the dict to db and close.
NOTE:
1. If the key is not string, use str to convert it
2. If the val is not string, try str first, if not work, use pickle
"""
log.info('Saving dict to \033[0;33m%s\033[0m, the size of dict \
is \033[0;32m%d\033[0m' % (db_file, len(key_dict)))
db = lmdb_tools.open(db_file, append)
try:
lmdb_tools.write_dict(db, key_dict, operation)
except:
try:
lmdb_tools.write_dict_str(db, key_dict)
except:
lmdb_tools.write_dict_pickle(db, key_dict)
log.info('Finished, the Db size: \033[0;32m%d\033[0m' %
lmdb_tools.get_entries(db))
lmdb_tools.close(db)
def __init__(self, db_file, readonly=True, echo=True, append=False):
self.readonly = readonly
self.db_file = db_file
self._echo = echo
self._warn = True
# If max commit interval time is 60 sec, it the commit interval
# is more than it, the buffer will be commited.
self._max_commit_interval = 60
# The timer
self._timer = timer_lib.timer(start=True)
# Define the dumper of the key and value
self._key_dumper = None
self._val_dumper = None
self._key_parser = None
self._val_parser = None
# The size of the commit buffer
self._buf_size = 100
self._buf_counter = 0
self._cur = None
self._iter = None
if readonly:
self.db = lmdb_tools.open_ro(db_file)
else:
self.db = lmdb_tools.open(db_file, append=append)
if self.db is None:
log.error('\033[01;31mERROR\033[0m: Can not open the \
db file \033[32m%s\033[0m' % db_file)
return
self._txn = self.db.begin(write=not self.readonly)
if echo:
log.info('Open %s, size: \033[01;31m%d\033[0m' %
(db_file, self.get_entries()))
def open(db_file):
log.info('Opening lmdb file: \033[0;33m%s\033[0m' % db_file)
db = lmdb_tools.open(db_file)
log.info('Db size: \033[0;32m%d\033[0m' % lmdb_tools.get_entries(db))
return db