riak_key_count_example.py

#!/usr/bin/env python

import httplib

import simplejson

# Based on Christian Dowski's iresponse http://blog.dowski.com/2008/04/02/reading-chunked-http11-responses/
class IterableResponse(httplib.HTTPResponse, object):
    def __iter__(self):
        assert self.chunked != httplib._UNKNOWN, "Response is not chunked!"
        chunk_left = self.chunk_left

        while 1:
            if chunk_left is None:
                line = self.fp.readline()
                i = line.find(';')
                if i >= 0:
                    line = line[:i]
                chunk_left = int(line, 16)
                if chunk_left == 0:
                    break
            yield self._safe_read(chunk_left)

            self._safe_read(2)
            chunk_left = None

        while 1:
            line = self.fp.readline()
            if not line:
                break
            if line == '\r\n':
                break

        self.close()

if __name__ == '__main__':
    loads = simplejson.loads
    conn = httplib.HTTPConnection('localhost:8098')
    conn.response_class = IterableResponse
    conn.request('GET', '/riak/sessions?keys=stream&props=false')
    r = conn.getresponse()
    num_keys = 0
    for chunk in r:
        try:
            decoded = loads(chunk)
        except ValueError:
            continue
        keys = (decoded or {}).get('keys', [])
        num_keys += len(keys)
        print num_keys
    print 'TOTAL KEYS:', num_keys