Creating Deterministic Shard Sentinel Value keys

shard_key_creator.py

from subprocess import check_output
from uuid import uuid4
from redis import Redis

PORT = 19475
SHARD_IDS = [1,2,3,4]
R = Redis(port=PORT)

## keys is unsafe for production as it can block for a long time - use scan instead if testing there
_shard_keys_command = lambda shard_id, pattern: ['/opt/redislabs/bin/shard-cli', str(shard_id), 'keys', str(pattern)]

def keys_from_shard_id(shard_id):
    return check_output(_shard_keys_command(shard_id, '*')).split('\n')[0:-1] # trim last item as it's an empty string

def check_key_on_shard(shard_id, key):
    return bool(check_output(_shard_keys_command(shard_id, key)).replace('\n', ''))

if __name__ == '__main__':
    ## Create Enough Keys that you'll get a few per shard
    shard_keys = dict(((str(uuid4()), None) for i in range(100)))

    ## Put them into redis
    for key in shard_keys.keys():
        R.set(key, 'data')

    ## go through each shard and figure out which key ended up on which shard
    for shard_id in SHARD_IDS:
        for key in keys_from_shard_id(shard_id):
            shard_keys[key] = shard_id

    ## invert the index so we now have one key per shard
    keys_for_shard_id = dict(((v,k) for k,v in shard_keys.items()))

    ## Create new keys with the sentinel key values in there
    test_keys = ['{' + keys_for_shard_id[shard_id] + '}:shard:' + str(shard_id) for shard_id in SHARD_IDS]

    ## put test keys onto shards
    for key in test_keys:
        R.set(key, 'worf')

    ## check to make sure that it lined up correctly
    for shard_id, key in zip(SHARD_IDS, test_keys):
        assert check_key_on_shard(shard_id, key)