redis space usage report grouped by prefixes on keys

redis_usage.py

#!/usr/bin/env python

# production friendly
# slow if there are too many keys but the script will print partial outputs
# usage is an approximation for complex objects

import math
import redis
import urlparse
import os

redis_url = os.getenv("redis_url")

redis_info = urlparse.urlparse(redis_url)
REDIS_HOST = redis_info.hostname
REDIS_PORT = redis_info.port or 6379
REDIS_PASSWORD = redis_info.password
REDIS_DATABASE = int(redis_info.path.replace("/", "")) if redis_info.path else 0

r = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_DATABASE)

cursor = 0
total_size = 0
sizes = {}
num_keys = {}
dtype_keys = {}
example_keys = {}

def convert_size(size_bytes):
   if size_bytes == 0:
       return "0B"
   size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
   i = int(math.floor(math.log(size_bytes, 1024)))
   p = math.pow(1024, i)
   s = round(size_bytes / p, 2)
   return "%s %s" % (s, size_name[i])

def display():
    print '#' * 40
    print "Job Status - %s size parsed (%s bytes)" % (convert_size(total_size), total_size)
    print '#' * 40
    print

    top_keys_by_space()
    top_keys_by_count()

def _pretty_print(result, title, show_helpful_info=True):
    print title
    print '=' * 40
    if not result:
        print 'n/a\n'
        return

    max_key_len = max((len(x[0]) for x in result))
    max_val_len = max((len(str(x[1])) for x in result))
    for key, val in result:
        key_padding = max(max_key_len - len(key), 0) * ' '

        if show_helpful_info:
            val_padding = max(max_val_len - len(str(val)), 0) * ' '
            info = "%s (%s)" % (dtype_keys[key], example_keys[key])
            print key,key_padding,'\t',val,val_padding,'\t',info
        else:
            print key,key_padding,'\t',val

    print

def top_keys_by_space():
    sorted_space_rank = sorted(sizes.iteritems(), key=lambda (k,v): (v,k), reverse=True)
    sorted_space_rank = sorted_space_rank[:20]
    sorted_space_rank = [(k, convert_size(v)) for (k,v) in sorted_space_rank]
    _pretty_print(sorted_space_rank, "Key Prefixes with Top Space Usage")


def top_keys_by_count():
    sorted_space_rank = sorted(num_keys.iteritems(), key=lambda (k,v): (v,k), reverse=True)
    sorted_space_rank = sorted_space_rank[:20]
    _pretty_print(sorted_space_rank, "Most Common Key Prefixes")


def length(key, dtype):
    if dtype == "string":
        l = len(r.get(key))
    elif dtype == "list":
        l = len(str(r.lrange(key, 0, -1)))
    elif dtype == "set":
        l = len(str(r.smembers(key)))
    elif dtype == "hash":
        l = len(str(r.hgetall(key)))
    elif dtype == "zset":
        l = len(str(r.zrange(key, 0, -1)))
    else:
        raise "Could not recognize type #{dtype}"
    # print "length of key %s with dtype %s - %f" % (key, dtype, l)
    return l

count = 0
while True:
    count += 1
    result = r.scan(cursor)
    cursor = int(result[0])
    keys = result[1]
    # print " cursor is at %d with %d length of keys in this loop" % (cursor, len(keys))
    for k in keys:
        dtype = r.type(k)
        key_prefix = k.partition(':')[0]
        if key_prefix not in sizes:
            sizes[key_prefix] = 0
        if key_prefix not in num_keys:
            num_keys[key_prefix] = 0
        if key_prefix not in example_keys:
            example_keys[key_prefix] = k
        if key_prefix not in dtype_keys:
            dtype_keys[key_prefix] = dtype

        num_keys[key_prefix] += 1
        try:
            size = length(k, dtype)
            sizes[key_prefix] += size
            total_size += size
        except:
            print "[warning] could not get size of key %s" % k
            
    display()

    if cursor == 0:
        break