This is a small script to get all keys from a Riak bucket using only standard Python libraries. It can be used as-is to list all keys from a bucket line-by-line, or can be used with other code to do things like purge an entire bucket, or perform some action over all keys in a bucket, etc. It takes advantage of Riak's streaming to grab a list of …

get-all-bucket-keys.py

#!/usr/bin/env python

import sys
from optparse import OptionParser
from urlparse import urljoin
import urllib2
import json

DEFAULT_RIAK_URL = "http://localhost:8098/"

def generate_all_keys(riak_url, bucket_name):
    """Generate all keys from a bucket"""
    
    url = urljoin(riak_url, "/buckets/%s/keys?keys=stream" % bucket_name)
    
    r = urllib2.Request(url=url)
    f = urllib2.urlopen(r)
    
    decoder = json.JSONDecoder()
    
    json_str = ''
    
    while True:
        chunk = f.read(1024)
        
        json_str += chunk
        
        while True:
            try:
                obj, pos = decoder.raw_decode(json_str)
            except ValueError:
                break
            else:
                if 'keys' not in obj:
                    if obj['error']:
                        raise Exception("Error: %s" % (obj['error']))
                    else:
                        raise Exception("Unknown error")
            
                for key in obj['keys']:
                    yield key
            
                json_str = json_str[pos:]
                
        if not chunk:
            break
                
    
def get_options():
    usage = "usage: %prog [options] bucket-name"
    parser = OptionParser(usage=usage)
    
    (options, args) = parser.parse_args()
    
    parser.add_option("-u", "--riak-url",
                      dest="riak_url",
                      help="Turn on debug output",
                      default=DEFAULT_RIAK_URL)
    
    if len(args) != 1:
        parser.error("Please specifiy exactly one arg - the bucket name")
    
    (options, args) = parser.parse_args()
    
    return (options, args)

def main(args):
    (options, args) = get_options()
    
    riak_url = options.riak_url
    bucket_name = args[0]

    for url in generate_all_keys(riak_url, bucket_name):
        print url

if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))