class Counters(object):
def __init__(self, prefix="counter", host=None, port=None, track_counters=False):
self.c = get_client(host, port, False)
self.prefix = prefix
if track_counters:
self.tracker = Hash("c-tracker-%s" % prefix, host=host, port=port)
else:
self.tracker = None
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.delete()
def inc(self, name, value=1, track_id=None):
if self.tracker:
self.tracker.add(track_id or name, now_as_iso())
return retry_call(self.c.incr, "%s-%s" % (self.prefix, name), value)
def dec(self, name, value=1, track_id=None):
if self.tracker:
self.tracker.pop(str(track_id or name))
return retry_call(self.c.decr, "%s-%s" % (self.prefix, name), value)
def get_queues_sizes(self):
out = {}
for queue in retry_call(self.c.keys, "%s-*" % self.prefix):
queue_size = int(retry_call(self.c.get, queue))
out[queue] = queue_size
return {k.decode('utf-8'): v for k, v in out.items()}
def get_queues(self):
return [k.decode('utf-8') for k in retry_call(self.c.keys, "%s-*" % self.prefix)]
def ready(self):
try:
self.c.ping()
except ConnectionError:
return False
return True
def reset_queues(self):
if self.tracker:
self.tracker.delete()
for queue in retry_call(self.c.keys, "%s-*" % self.prefix):
retry_call(self.c.set, queue, "0")
def delete(self):
if self.tracker:
self.tracker.delete()
for queue in retry_call(self.c.keys, "%s-*" % self.prefix):
retry_call(self.c.delete, queue)
def do_ui(self, args):
"""
Perform UI related operations
Usage:
ui show_sessions [username]
ui clear_sessions [username]
actions:
show_sessions show all active sessions
clear_sessions Removes all active sessions
Parameters:
username User use to filter sessions
[optional]
Examples:
# Clear sessions for user bob
ui clear_sessions bob
# Show all current sessions
ui show_sessions
"""
valid_func = ['clear_sessions', 'show_sessions']
args = self._parse_args(args)
if len(args) not in [1, 2]:
self._print_error("Wrong number of arguments for restore command.")
return
func = args[0]
if func not in valid_func:
self._print_error(f"Invalid action '{func}' for ui command.")
return
if func == 'clear_sessions':
username = None
if len(args) == 2:
username = args[1]
flsk_sess = Hash("flask_sessions",
host=config.core.redis.nonpersistent.host,
port=config.core.redis.nonpersistent.port)
if not username:
flsk_sess.delete()
self.logger.info("All sessions where cleared.")
else:
for k, v in flsk_sess.items().items():
if v.get('username', None) == username:
self.logger.info(f"Removing session: {v}")
flsk_sess.pop(k)
self.logger.info(
f"All sessions for user '{username}' removed.")
if func == 'show_sessions':
username = None
if len(args) == 2:
username = args[1]
flsk_sess = Hash("flask_sessions",
host=config.core.redis.nonpersistent.host,
port=config.core.redis.nonpersistent.port)
if not username:
for k, v in flsk_sess.items().items():
self.logger.info(f"{v.get('username', None)} => {v}")
else:
self.logger.info(f'Showing sessions for user {username}:')
for k, v in flsk_sess.items().items():
if v.get('username', None) == username:
self.logger.info(f" {v}")
class DistributedBackup(object):
def __init__(self,
working_dir,
worker_count=50,
spawn_workers=True,
use_threading=False,
logger=None):
self.working_dir = working_dir
self.datastore = forge.get_datastore(archive_access=True)
self.logger = logger
self.plist = []
self.use_threading = use_threading
self.instance_id = get_random_id()
self.worker_queue = NamedQueue(f"r-worker-{self.instance_id}",
ttl=1800)
self.done_queue = NamedQueue(f"r-done-{self.instance_id}", ttl=1800)
self.hash_queue = Hash(f"r-hash-{self.instance_id}")
self.bucket_error = []
self.VALID_BUCKETS = sorted(list(
self.datastore.ds.get_models().keys()))
self.worker_count = worker_count
self.spawn_workers = spawn_workers
self.total_count = 0
self.error_map_count = {}
self.missing_map_count = {}
self.map_count = {}
self.last_time = 0
self.last_count = 0
self.error_count = 0
def cleanup(self):
self.worker_queue.delete()
self.done_queue.delete()
self.hash_queue.delete()
for p in self.plist:
p.terminate()
def done_thread(self, title):
t0 = time.time()
self.last_time = t0
running_threads = self.worker_count
while running_threads > 0:
msg = self.done_queue.pop(timeout=1)
if msg is None:
continue
if "stopped" in msg:
running_threads -= 1
continue
bucket_name = msg.get('bucket_name', 'unknown')
if msg.get('success', False):
self.total_count += 1
if msg.get("missing", False):
if bucket_name not in self.missing_map_count:
self.missing_map_count[bucket_name] = 0
self.missing_map_count[bucket_name] += 1
else:
if bucket_name not in self.map_count:
self.map_count[bucket_name] = 0
self.map_count[bucket_name] += 1
new_t = time.time()
if (new_t - self.last_time) > 5:
if self.logger:
self.logger.info(
"%s (%s at %s keys/sec) ==> %s" %
(self.total_count, new_t - self.last_time,
int((self.total_count - self.last_count) /
(new_t - self.last_time)), self.map_count))
self.last_count = self.total_count
self.last_time = new_t
else:
self.error_count += 1
if bucket_name not in self.error_map_count:
self.error_map_count[bucket_name] = 0
self.error_map_count[bucket_name] += 1
# Cleanup
self.cleanup()
summary = ""
summary += "\n########################\n"
summary += "####### SUMMARY #######\n"
summary += "########################\n"
summary += "%s items - %s errors - %s secs\n\n" % \
(self.total_count, self.error_count, time.time() - t0)
for k, v in self.map_count.items():
summary += "\t%15s: %s\n" % (k.upper(), v)
if len(self.missing_map_count.keys()) > 0:
summary += "\n\nMissing data:\n\n"
for k, v in self.missing_map_count.items():
summary += "\t%15s: %s\n" % (k.upper(), v)
if len(self.error_map_count.keys()) > 0:
summary += "\n\nErrors:\n\n"
for k, v in self.error_map_count.items():
summary += "\t%15s: %s\n" % (k.upper(), v)
if len(self.bucket_error) > 0:
summary += f"\nThese buckets failed to {title.lower()} completely: {self.bucket_error}\n"
if self.logger:
self.logger.info(summary)
# noinspection PyBroadException,PyProtectedMember
def backup(self, bucket_list, follow_keys=False, query=None):
if query is None:
query = 'id:*'
for bucket in bucket_list:
if bucket not in self.VALID_BUCKETS:
if self.logger:
self.logger.warn(
"\n%s is not a valid bucket.\n\n"
"The list of valid buckets is the following:\n\n\t%s\n"
% (bucket.upper(), "\n\t".join(self.VALID_BUCKETS)))
return
targets = ', '.join(bucket_list)
try:
if self.logger:
self.logger.info("\n-----------------------")
self.logger.info("----- Data Backup -----")
self.logger.info("-----------------------")
self.logger.info(f" Deep: {follow_keys}")
self.logger.info(f" Buckets: {targets}")
self.logger.info(f" Workers: {self.worker_count}")
self.logger.info(f" Target directory: {self.working_dir}")
self.logger.info(f" Filtering query: {query}")
# Start the workers
for x in range(self.worker_count):
if self.use_threading:
t = threading.Thread(target=backup_worker,
args=(x, self.instance_id,
self.working_dir))
t.setDaemon(True)
t.start()
else:
p = Process(target=backup_worker,
args=(x, self.instance_id, self.working_dir))
p.start()
self.plist.append(p)
# Start done thread
dt = threading.Thread(target=self.done_thread,
args=('Backup', ),
name="Done thread")
dt.setDaemon(True)
dt.start()
# Process data buckets
for bucket_name in bucket_list:
try:
collection = self.datastore.get_collection(bucket_name)
for item in collection.stream_search(query,
fl="id",
item_buffer_size=500,
as_obj=False):
self.worker_queue.push({
"bucket_name": bucket_name,
"key": item['id'],
"follow_keys": follow_keys
})
except Exception as e:
self.cleanup()
if self.logger:
self.logger.execption(e)
self.logger.error(
"Error occurred while processing bucket %s." %
bucket_name)
self.bucket_error.append(bucket_name)
for _ in range(self.worker_count):
self.worker_queue.push({"stop": True})
dt.join()
except Exception as e:
if self.logger:
self.logger.execption(e)
def restore(self):
try:
if self.logger:
self.logger.info("\n------------------------")
self.logger.info("----- Data Restore -----")
self.logger.info("------------------------")
self.logger.info(f" Workers: {self.worker_count}")
self.logger.info(f" Target directory: {self.working_dir}")
for x in range(self.worker_count):
if self.use_threading:
t = threading.Thread(target=restore_worker,
args=(x, self.instance_id,
self.working_dir))
t.setDaemon(True)
t.start()
else:
p = Process(target=restore_worker,
args=(x, self.instance_id, self.working_dir))
p.start()
self.plist.append(p)
# Start done thread
dt = threading.Thread(target=self.done_thread,
args=('Restore', ),
name="Done thread")
dt.setDaemon(True)
dt.start()
# Wait for workers to finish
dt.join()
except Exception as e:
if self.logger:
self.logger.execption(e)
