Insertion script for blog post example

insert.py

"""Simple Data loader.

Usage:
  insert.py [-c] [-n=<x>] [-p=<y>]

Options:
  -c  enable compression (disabled by default)
  -n=<x> total number of rows to insert (default: 1M)
  -p=<y> number of partitions to insert the row into (default: 100)
"""
from __future__ import print_function
from uuid import uuid4, uuid1
from docopt import docopt
from cassandra.cluster import Cluster
from random import randint, choice, random
from string import ascii_letters

max_async_query = 1000

if __name__ == '__main__':
    args = docopt(__doc__)
    use_compression = args['-c']

    rows = 1000000 if args['-n'] is None else args['-n']
    partitions = 100 if args['-p'] is None else args['-p']

    assert rows > partitions

    cluster = Cluster()
    session = cluster.connect()

    session.execute("CREATE KEYSPACE ks WITH replication = { 'class': 'SimpleStrategy', 'replication_factor': '1' };")

    session.execute("""
        CREATE TABLE ks.events (
            id uuid,
            received_at timeuuid,
            property_1 int,
            property_2 text,
            property_3 float,
            PRIMARY KEY (id, received_at)
        ) WITH compression = {'sstable_compression': '%s'}
        """ % ('LZ4Compressor' if use_compression else ''))

    insert = session.prepare('INSERT INTO ks.events (id, received_at, property_1, property_2, property_3) VALUES (?, ?, ?, ?, ?)')

    row_per_partition = rows // partitions
    futures = []
    for _ in range(0, partitions):
        key = uuid4()
        for _ in range(0, row_per_partition):
            random_text = ''.join(choice(ascii_letters) for _ in range(randint(3, 10)))
            futures.append(session.execute_async(insert, [key, uuid1(), randint(0, 2^31), random_text, random()]))
            if (len(futures) > max_async_query):
                for f in futures:
                    f.result()
                futures = []

    for f in futures:
        f.result()

    assert len(list(session.execute('SELECT * from ks.events;'))) == rows