diff --git a/README.md b/README.md index 369368b..fb30141 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,7 @@ optional arguments: --host HOST Redis host --port PORT Redis port --password PASSWORD Redis password + --ssl Enable SSL connection --cluster Enable cluster mode --batch_size BATCH_SIZE Batch size for SCAN command diff --git a/rks/main.py b/rks/main.py index fa317dd..69e3fec 100644 --- a/rks/main.py +++ b/rks/main.py @@ -2,13 +2,16 @@ import redis import rediscluster from datetime import datetime + from .redis_utils import get_redis_keys, get_redis_cluster_keys + def main(): parser = argparse.ArgumentParser(description='Analyze Redis Instance Key Statistics.') parser.add_argument('--host', required=True, help='Redis host') parser.add_argument('--port', required=True, help='Redis port') parser.add_argument('--password', default=None, help='Redis password') + parser.add_argument('--ssl', action='store_true', help='Enable SSL connection') parser.add_argument('--cluster', action='store_true', help='Enable cluster mode') parser.add_argument('--batch_size', type=int, default=1000, help='Batch size for SCAN command') parser.add_argument('--replica_only', action='store_true', help='Execute only on replica instances') @@ -22,6 +25,7 @@ def main(): rc = rediscluster.RedisCluster(startup_nodes=[{"host": args.host, "port": int(args.port)}], skip_full_coverage_check=True, password=args.password, + ssl=args.ssl, socket_timeout=CONNECTION_TIMEOUT) redis_info = rc.info() @@ -32,7 +36,13 @@ def main(): process_start_time = datetime.now() - if get_redis_cluster_keys(rc, args.batch_size, args.replica_only, args.pretty_format) == -1: + if get_redis_cluster_keys(rc, + args.batch_size, + args.replica_only, + args.pretty_format, + args.password, + args.ssl, + CONNECTION_TIMEOUT) == -1: print(f"Aborted the operation on non-readonly redis at host: {args.host}") return @@ -56,6 +66,7 @@ def main(): try: r = redis.Redis(host=args.host, port=int(args.port), password=args.password, + ssl=args.ssl, socket_timeout=CONNECTION_TIMEOUT) redis_info = r.info() diff --git a/rks/redis_utils.py b/rks/redis_utils.py index 6d52296..d9dfcae 100644 --- a/rks/redis_utils.py +++ b/rks/redis_utils.py @@ -62,7 +62,7 @@ def get_redis_keys(r, batch_size, db_num, use_pretty): return analyze_redis_keys(min_heap, prefix_statistics_map, total_key_size, total_key_count, key_count_by_type, db_num, use_pretty) -def get_redis_cluster_keys(rc, batch_size, replica_only, use_pretty): +def get_redis_cluster_keys(rc, batch_size, replica_only, use_pretty, password, ssl, socket_timeout): slave_flag = False min_heap = [] @@ -96,7 +96,13 @@ def get_redis_cluster_keys(rc, batch_size, replica_only, use_pretty): node = next((node for node in nodes if node['id'] == master['master']), None) else: node = next((node for node in nodes if node['id'] == master['slaves'][0]), None) - r = redis.Redis(host=node['host'], port=node['port']) + r = redis.Redis( + host=node['host'], + port=node['port'], + ssl=ssl, + password=password, + socket_timeout=socket_timeout + ) r.execute_command('READONLY') script = """