multiprocess_execute_python3.py

#!/usr/bin/env python

from multiprocessing import Pool
import sys
import time

from cassandra.cluster import Cluster
from cassandra.query import tuple_factory

def query_gen(n):
    for _ in range(n):
        yield ('local', )


class QueryManager(object):
    session = None
    prepared = None
    batch_size = 10

    def __init__(self, process_count=None):
        self.pool = Pool(processes=process_count, initializer=self._setup)

    @classmethod
    def _setup(cls):
        cluster = Cluster()
        cls.session = cluster.connect()
        cls.session.row_factory = tuple_factory
        cls.prepared = cls.session.prepare('SELECT * FROM system.local WHERE key=?')

    def close_pool(self):
        self.pool.close()
        self.pool.join()

    def get_results(self, params):
        results = self.pool.map(_get_multiproc, params, self.batch_size)
        return results

    @classmethod
    def _execute_request(cls, params):
        result = cls.session.execute(cls.prepared, params)
        if result:
            return result[0]
        return None


def _get_multiproc(params):
    return QueryManager._execute_request(params)


if __name__ == '__main__':
    try:
        iterations = int(sys.argv[1])
        processes = int(sys.argv[2]) if len(sys.argv) > 2 else None
    except (IndexError, ValueError):
        print("Usage: %s <num iterations> [<num processes>]" % sys.argv[0])
        sys.exit(1)

    qm = QueryManager(processes)

    start = time.time()
    rows = qm.get_results(query_gen(iterations))
    delta = time.time() - start
    for row in rows:
        print(row)
    print("%d queries in %s seconds (%s/s)" % (iterations, delta, iterations / delta))