pump.py

#!/usr/bin/env python

import sys
from optparse import OptionParser
import random
import sha
import base64
import httplib

parser = OptionParser()
parser.add_option(
    "--bucket_type", 
    dest="bucket_type",
    help="Riak Bucket Type", 
    metavar="BUCKET_TYPE"
)
parser.add_option(
    "--bucket_name", 
    dest="bucket_name",
    help="Riak Bucket Name", 
    metavar="BUCKET_NAME"
)
parser.add_option(
    "--key", 
    dest="key",
    help="Riak Key", 
    metavar="KEY"
)

parser.add_option(
    "--host", 
    dest="host",
    help="Solr endpoint host", 
    metavar="HOST",
    type="string",
    default="localhost"
)
parser.add_option(
    "--port", 
    dest="port",
    help="Solr endpoint port", 
    metavar="PORT",
    type="int",
    default=8093
)
parser.add_option(
    "--index", 
    dest="index",
    help="Solr core", 
    metavar="INDEX",
    type="string"
)
parser.add_option(
    "--batch_size", 
    dest="batch_size",
    help="Batch size", 
    metavar="BATCH_SIZE",
    type="int",
    default=1
)
parser.add_option(
    "--num_puts", 
    dest="num_puts",
    help="Number of puts (<0 for infinity)", 
    metavar="NUM_PUTS",
    type="int",
    default=1
)

(options, args) = parser.parse_args()

if not options.index :
    parser.error("The index option is required")
if not options.bucket_type :
    parser.error("The bucket_type option is required")

def file_contents(path) :
    file = open(path,  "r")
    return file.read()

class Connection :

    def __init__(self, host, port, context) :
        self.connection = None
        self.host = host
        self.port = int(port)
        self.context = context

    def get(self) :
        try :
            if self.connection == None :
                self.connection = httplib.HTTPConnection(self.host, self.port)
            self.connection.request("GET", self.context)
            response = self.connection.getresponse()
            body = response.read()
            return body
        except Exception as e :
            self.connection = None
            raise Exception("Connection error: %s to %s: %s" % (e, self.host, self.port, e))
    
    def post(self, body) :
        try :
            if self.connection == None :
                self.connection = httplib.HTTPConnection(self.host, self.port)
            self.connection.request("POST", self.context, body, {'Content-Type': "application/json"})
            response = self.connection.getresponse()
            return response.status
        except Exception as e :
            self.connection = None
            raise Exception("Connection error: %s to %s" % (str(e), self.str()))
    
    def str(self) :
        return "http://" + self.host + ":" + str(self.port) + self.context

def do_hash(string) :
    return sha.new(string).hexdigest()



docstr = file_contents("./doc.json")

distribution = [
    (23, 34),
    (56, 78),
    (101, 109),
    (114, 130),
    (143, 149),
    (153, 167),
    (191, 207),
    (234, 242)
]

random.seed()

for iteration in range(0, options.num_puts) :
    connection = Connection(options.host, options.port, "/internal_solr/%s/update/" % options.index)
    # print connection.str()
    body = "{"
    bucket_type = options.bucket_type
    bucket = options.index
    for i in range(0, options.batch_size) :
        key = random.randint(1, 99999999)
        lb, ub = distribution[random.randint(0, 7)]
        pn = random.randint(lb, ub)
        fpn = pn - 1
        hash_ = do_hash("%s %s %s" % (bucket_type, bucket, key))
        body += docstr % (
            bucket_type, bucket, key, hash_, options.bucket_type, bucket, key, pn, fpn, pn,
            base64.b64encode(bucket_type), base64.b64encode(bucket), base64.b64encode(str(key)), base64.b64encode(hash_)
        )
        if i < options.batch_size - 1 :
            body += ","
    body += "}"
    # print body
    status = connection.post(body)
    if status != 200 :
        print "Received bad status from %s: %s.  Exiting..." % (connection.str(), status)
        sys.exit(1)