332 lines
14 KiB
Python
332 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Redis Queue Management CLI Tool for yt-ops-client.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
from typing import Optional
|
|
|
|
import redis
|
|
|
|
try:
|
|
from dotenv import load_dotenv
|
|
except ImportError:
|
|
load_dotenv = None
|
|
|
|
try:
|
|
from tabulate import tabulate
|
|
except ImportError:
|
|
print("'tabulate' library not found. Please install it with: pip install tabulate", file=sys.stderr)
|
|
tabulate = None
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class QueueManager:
|
|
"""Manages Redis lists (queues)."""
|
|
|
|
def __init__(self, redis_host='localhost', redis_port=6379, redis_password=None):
|
|
"""Initialize Redis connection."""
|
|
logger.info(f"Attempting to connect to Redis at {redis_host}:{redis_port}...")
|
|
try:
|
|
self.redis = redis.Redis(
|
|
host=redis_host,
|
|
port=redis_port,
|
|
password=redis_password,
|
|
decode_responses=True,
|
|
socket_connect_timeout=5,
|
|
socket_timeout=5
|
|
)
|
|
self.redis.ping()
|
|
logger.info(f"Successfully connected to Redis.")
|
|
except redis.exceptions.ConnectionError as e:
|
|
logger.error(f"Failed to connect to Redis at {redis_host}:{redis_port}: {e}")
|
|
sys.exit(1)
|
|
|
|
def list_queues(self, pattern: str):
|
|
"""Lists queues matching a pattern and their sizes."""
|
|
queues = []
|
|
for key in self.redis.scan_iter(match=pattern):
|
|
key_type = self.redis.type(key)
|
|
if key_type == 'list':
|
|
size = self.redis.llen(key)
|
|
queues.append({'name': key, 'size': size})
|
|
return queues
|
|
|
|
def peek(self, queue_name: str, count: int):
|
|
"""Returns the top `count` items from a queue without removing them."""
|
|
return self.redis.lrange(queue_name, 0, count - 1)
|
|
|
|
def count(self, queue_name: str) -> int:
|
|
"""Returns the number of items in a queue."""
|
|
return self.redis.llen(queue_name)
|
|
|
|
def push_from_file(self, queue_name: str, file_path: str, wrap_key: Optional[str] = None) -> int:
|
|
"""Populates a queue from a file (text with one item per line, or JSON with an array of items)."""
|
|
count = 0
|
|
|
|
if file_path.lower().endswith('.json'):
|
|
if wrap_key:
|
|
logger.warning("--wrap-file-line-in-json is ignored for JSON files, as they are expected to contain complete items.")
|
|
logger.info("Detected JSON file. Attempting to parse as an array of items.")
|
|
try:
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
data = json.load(f)
|
|
if not isinstance(data, list):
|
|
logger.error("JSON file must contain a list/array.")
|
|
return 0
|
|
|
|
# Items can be strings or objects. If objects, they should be converted to JSON strings.
|
|
items_to_add = []
|
|
for item in data:
|
|
if isinstance(item, str):
|
|
items_to_add.append(item.strip())
|
|
else:
|
|
items_to_add.append(json.dumps(item))
|
|
|
|
items_to_add = [item for item in items_to_add if item]
|
|
|
|
pipe = self.redis.pipeline()
|
|
for item in items_to_add:
|
|
pipe.rpush(queue_name, item)
|
|
count += 1
|
|
if count > 0 and count % 1000 == 0:
|
|
pipe.execute()
|
|
logger.info(f"Pushed {count} items...")
|
|
pipe.execute()
|
|
|
|
except (IOError, json.JSONDecodeError) as e:
|
|
logger.error(f"Failed to read or parse JSON file '{file_path}': {e}")
|
|
return 0
|
|
else:
|
|
logger.info("Reading items from text file (one per line).")
|
|
try:
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
pipe = self.redis.pipeline()
|
|
for line in f:
|
|
item = line.strip()
|
|
if item:
|
|
if wrap_key:
|
|
payload = json.dumps({wrap_key: item})
|
|
else:
|
|
payload = item
|
|
pipe.rpush(queue_name, payload)
|
|
count += 1
|
|
if count > 0 and count % 1000 == 0:
|
|
pipe.execute()
|
|
logger.info(f"Pushed {count} items...")
|
|
pipe.execute()
|
|
except IOError as e:
|
|
logger.error(f"Failed to read file '{file_path}': {e}")
|
|
return 0
|
|
|
|
logger.info(f"Finished. Pushed a total of {count} items to '{queue_name}'.")
|
|
return count
|
|
|
|
def push_generated(self, queue_name: str, prefix: str, count: int) -> int:
|
|
"""Pushes generated payloads to a queue."""
|
|
from datetime import datetime
|
|
timestamp = datetime.now().strftime('%Y%m%dt%H%M')
|
|
|
|
pipe = self.redis.pipeline()
|
|
pushed_count = 0
|
|
for i in range(count):
|
|
generated_value = f"{prefix}_{timestamp}_{i:04d}"
|
|
payload = json.dumps({"url": generated_value})
|
|
pipe.rpush(queue_name, payload)
|
|
pushed_count += 1
|
|
if pushed_count > 0 and pushed_count % 1000 == 0:
|
|
pipe.execute()
|
|
logger.info(f"Pushed {pushed_count} of {count} items...")
|
|
pipe.execute()
|
|
logger.info(f"Finished. Pushed a total of {pushed_count} items to '{queue_name}'.")
|
|
return pushed_count
|
|
|
|
def push_static(self, queue_name: str, payload: str, count: int) -> int:
|
|
"""Pushes a static payload multiple times to a queue."""
|
|
try:
|
|
json.loads(payload)
|
|
except json.JSONDecodeError:
|
|
logger.error(f"Invalid JSON in --payload-json: {payload}")
|
|
return 0
|
|
|
|
pipe = self.redis.pipeline()
|
|
pushed_count = 0
|
|
for _ in range(count):
|
|
pipe.rpush(queue_name, payload)
|
|
pushed_count += 1
|
|
if pushed_count > 0 and pushed_count % 1000 == 0:
|
|
pipe.execute()
|
|
logger.info(f"Pushed {pushed_count} of {count} items...")
|
|
pipe.execute()
|
|
logger.info(f"Finished. Pushed a total of {pushed_count} items to '{queue_name}'.")
|
|
return pushed_count
|
|
|
|
def clear(self, queue_name: str, dump_path: Optional[str] = None) -> int:
|
|
"""Clears a queue, optionally dumping its contents to a file."""
|
|
size = self.redis.llen(queue_name)
|
|
if size == 0:
|
|
logger.info(f"Queue '{queue_name}' is already empty.")
|
|
return 0
|
|
|
|
if dump_path:
|
|
logger.info(f"Dumping {size} items from '{queue_name}' to '{dump_path}'...")
|
|
with open(dump_path, 'w') as f:
|
|
# Use lpop to be memory efficient for very large queues
|
|
while True:
|
|
item = self.redis.lpop(queue_name)
|
|
if item is None:
|
|
break
|
|
f.write(item + '\n')
|
|
logger.info("Dump complete.")
|
|
# After lpop, the queue is already empty.
|
|
return size
|
|
|
|
deleted_count = self.redis.delete(queue_name)
|
|
if deleted_count > 0:
|
|
logger.info(f"Cleared queue '{queue_name}' ({size} items).")
|
|
|
|
return size
|
|
|
|
|
|
def add_queue_manager_parser(subparsers):
|
|
"""Adds the parser for the 'queue' command."""
|
|
parser = subparsers.add_parser(
|
|
'queue',
|
|
description='Manage Redis queues.',
|
|
formatter_class=argparse.RawTextHelpFormatter,
|
|
help='Manage Redis queues.'
|
|
)
|
|
|
|
# Common arguments for all queue manager subcommands
|
|
common_parser = argparse.ArgumentParser(add_help=False)
|
|
common_parser.add_argument('--env-file', help='Path to a .env file to load environment variables from.')
|
|
common_parser.add_argument('--env', default='dev', help="Environment name for queue prefixes (e.g., 'stg', 'prod'). Defaults to 'dev'.")
|
|
common_parser.add_argument('--redis-host', default=None, help='Redis host. Defaults to REDIS_HOST or MASTER_HOST_IP env var, or localhost.')
|
|
common_parser.add_argument('--redis-port', type=int, default=None, help='Redis port. Defaults to REDIS_PORT env var, or 6379.')
|
|
common_parser.add_argument('--redis-password', default=None, help='Redis password. Defaults to REDIS_PASSWORD env var.')
|
|
common_parser.add_argument('--verbose', action='store_true', help='Enable verbose logging')
|
|
|
|
subparsers = parser.add_subparsers(dest='queue_command', help='Command to execute', required=True)
|
|
|
|
# List command
|
|
list_parser = subparsers.add_parser('list', help='List queues and their sizes.', parents=[common_parser])
|
|
list_parser.add_argument('--pattern', default='*queue*', help="Pattern to search for queue keys (default: '*queue*')")
|
|
|
|
# Peek command
|
|
peek_parser = subparsers.add_parser('peek', help='View items in a queue without removing them.', parents=[common_parser])
|
|
peek_parser.add_argument('queue_name', nargs='?', help="Name of the queue. Defaults to '<env>_stress_inbox'.")
|
|
peek_parser.add_argument('--count', type=int, default=10, help='Number of items to show (default: 10)')
|
|
|
|
# Push command
|
|
push_parser = subparsers.add_parser('push', help='Push items to a queue from a file, a generator, or a static payload.', parents=[common_parser])
|
|
push_parser.add_argument('queue_name', nargs='?', help="Name of the queue. Defaults to '<env>_stress_inbox'.")
|
|
push_parser.add_argument('--count', type=int, default=1, help='Number of items to push (for --payload-json or --generate-payload-prefix).')
|
|
|
|
source_group = push_parser.add_mutually_exclusive_group(required=True)
|
|
source_group.add_argument('--from-file', dest='file_path', help='Path to a file containing items to add (one per line, or a JSON array).')
|
|
source_group.add_argument('--payload-json', help='A static JSON payload to push. Use with --count to push multiple times.')
|
|
source_group.add_argument('--generate-payload-prefix', help='Generate JSON payloads with a timestamp and counter. Example: {"url": "PREFIX_yyyymmddthhmm_0001"}. Use with --count.')
|
|
|
|
push_parser.add_argument('--wrap-file-line-in-json', metavar='KEY', help="For text files (--from-file), wrap each line in a JSON object with the specified key (e.g., 'url' -> {\"url\": \"line_content\"}).")
|
|
|
|
# Clear command
|
|
clear_parser = subparsers.add_parser('clear', help='Clear a queue, optionally dumping its contents.', parents=[common_parser])
|
|
clear_parser.add_argument('queue_name', nargs='?', help="Name of the queue to clear. Defaults to '<env>_stress_inbox'.")
|
|
clear_parser.add_argument('--dump-to', help='File path to dump queue contents before clearing.')
|
|
clear_parser.add_argument('--confirm', action='store_true', help='Confirm this destructive action (required).')
|
|
|
|
return parser
|
|
|
|
|
|
def main_queue_manager(args):
|
|
"""Main dispatcher for 'queue' command."""
|
|
if load_dotenv:
|
|
was_loaded = load_dotenv(args.env_file)
|
|
if was_loaded:
|
|
print(f"Loaded environment variables from {args.env_file or '.env file'}", file=sys.stderr)
|
|
elif args.env_file:
|
|
print(f"ERROR: The specified --env-file was not found: {args.env_file}", file=sys.stderr)
|
|
return 1
|
|
|
|
if args.redis_host is None:
|
|
args.redis_host = os.getenv('REDIS_HOST', os.getenv('MASTER_HOST_IP', 'localhost'))
|
|
if args.redis_port is None:
|
|
args.redis_port = int(os.getenv('REDIS_PORT', 6379))
|
|
if args.redis_password is None:
|
|
args.redis_password = os.getenv('REDIS_PASSWORD')
|
|
|
|
if args.verbose:
|
|
logging.getLogger().setLevel(logging.DEBUG)
|
|
|
|
manager = QueueManager(
|
|
redis_host=args.redis_host,
|
|
redis_port=args.redis_port,
|
|
redis_password=args.redis_password
|
|
)
|
|
|
|
# For commands that operate on a single queue, set a default name based on the environment if not provided.
|
|
is_single_queue_command = args.queue_command in ['peek', 'push', 'clear']
|
|
if is_single_queue_command:
|
|
# `push`, `peek` and `clear` use a positional argument for queue_name.
|
|
# We check for `queue_name` attribute and if it's falsy (None or empty string).
|
|
if not getattr(args, 'queue_name', None):
|
|
default_queue_name = f"{args.env}_stress_inbox"
|
|
args.queue_name = default_queue_name
|
|
print(f"INFO: No queue name specified, defaulting to '{default_queue_name}' based on --env='{args.env}'.", file=sys.stderr)
|
|
|
|
if args.queue_command == 'list':
|
|
queues = manager.list_queues(args.pattern)
|
|
if not queues:
|
|
print(f"No queues found matching pattern '{args.pattern}'.")
|
|
return 0
|
|
if tabulate:
|
|
print(tabulate(queues, headers='keys', tablefmt='grid'))
|
|
else:
|
|
for q in queues:
|
|
print(f"{q['name']}: {q['size']}")
|
|
return 0
|
|
|
|
elif args.queue_command == 'peek':
|
|
size = manager.count(args.queue_name)
|
|
items = manager.peek(args.queue_name, args.count)
|
|
print(f"Queue '{args.queue_name}' has {size} items. Showing top {len(items)}:")
|
|
for i, item in enumerate(items):
|
|
print(f"{i+1: >3}: {item}")
|
|
return 0
|
|
|
|
elif args.queue_command == 'push':
|
|
if args.file_path:
|
|
if not os.path.exists(args.file_path):
|
|
print(f"Error: File not found at '{args.file_path}'", file=sys.stderr)
|
|
return 1
|
|
if args.count > 1:
|
|
logger.warning("--count is ignored when using --from-file.")
|
|
manager.push_from_file(args.queue_name, args.file_path, args.wrap_file_line_in_json)
|
|
elif args.payload_json:
|
|
manager.push_static(args.queue_name, args.payload_json, args.count)
|
|
elif args.generate_payload_prefix:
|
|
if args.count <= 0:
|
|
print("Error: --count must be 1 or greater for --generate-payload-prefix.", file=sys.stderr)
|
|
return 1
|
|
manager.push_generated(args.queue_name, args.generate_payload_prefix, args.count)
|
|
return 0
|
|
|
|
elif args.queue_command == 'clear':
|
|
if not args.confirm:
|
|
print("Error: --confirm flag is required for this destructive action.", file=sys.stderr)
|
|
return 1
|
|
manager.clear(args.queue_name, args.dump_to)
|
|
return 0
|
|
|
|
return 1 # Should not be reached
|