hyperdex-json-bridge-python

#!/usr/bin/env python

# Copyright (c) 2012, Cornell University
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#     * Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of HyperDex nor the names of its contributors may be
#       used to endorse or promote products derived from this software without
#       specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

from __future__ import absolute_import
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import with_statement


'''Interpret JSON objects as a script for testing HyperDex.'''


import simplejson
import sys

import hyperclient


def _frozen_dict(x):
    return tuple(sorted(x.items()))


class HyperDexJSONBridge(object):

    VALUE_IDENTITY = set(['string', 'int64', 'float',
                          'list(string)', 'list(int64)', 'list(float)',
                          'map(string,string)', 'map(string,int64)', 'map(string,float)',
                          'map(int64,string)', 'map(int64,int64)', 'map(int64,float)',
                          'map(float,string)', 'map(float,int64)', 'map(float,float)'])
    VALUE_CONVERT = {'set(string)': lambda x: set(x),
                     'set(int64)': lambda x: set(x),
                     'set(float)': lambda x: set(x)}

    def __init__(self, host, port):
        self._client = hyperclient.Client(host, port)

    def perform_action(self, action):
        if 'action' not in action:
            raise RuntimeError('missing \'action\' field')
        if not hasattr(self, '_action_' + action['action']):
            raise RuntimeError('unknown action %r' % action['action'])
        getattr(self, '_action_' + action['action'])(action)

    def _kve(self, action, call):
        self._check_fields(action, 'space', 'key', 'value', 'expected')
        expected = action['expected']
        try:
            returned = call(action['space'], action['key'], self._to_value(action['value']))
            self._compare_success(action['expected'], returned)
        except hyperclient.HyperClientException as e:
            self._compare_exception(expected, e)

    def _action_get(self, action):
        self._check_fields(action, 'space', 'key', 'expected')
        expected = action['expected']
        try:
            returned = self._client.get(action['space'], action['key'])
            self._compare_success(action['expected'], returned)
        except hyperclient.HyperClientException as e:
            self._compare_exception(expected, e)

    def _action_put(self, action):
        self._kve(action, self._client.put)

    def _action_put_if_not_exist(self, action):
        self._kve(action, self._client.put_if_not_exist)

    def _action_condput(self, action):
        self._check_fields(action, 'space', 'key', 'condition', 'value', 'expected')
        expected = action['expected']
        try:
            returned = self._client.condput(action['space'], action['key'],
                                            self._to_value(action['condition']),
                                            self._to_value(action['value']))
            self._compare_success(action['expected'], returned)
        except hyperclient.HyperClientException as e:
            self._compare_exception(expected, e)

    def _action_del(self, action):
        self._check_fields(action, 'space', 'key', 'expected')
        expected = action['expected']
        try:
            returned = self._client.delete(action['space'], action['key'])
            self._compare_success(action['expected'], returned)
        except hyperclient.HyperClientException as e:
            self._compare_exception(expected, e)

    def _action_atomic_add(self, action):
        self._kve(action, self._client.atomic_add)

    def _action_atomic_sub(self, action):
        self._kve(action, self._client.atomic_sub)

    def _action_atomic_mul(self, action):
        self._kve(action, self._client.atomic_mul)

    def _action_atomic_div(self, action):
        self._kve(action, self._client.atomic_div)

    def _action_atomic_mod(self, action):
        self._kve(action, self._client.atomic_mod)

    def _action_atomic_and(self, action):
        self._kve(action, self._client.atomic_and)

    def _action_atomic_or(self, action):
        self._kve(action, self._client.atomic_or)

    def _action_atomic_xor(self, action):
        self._kve(action, self._client.atomic_xor)

    def _action_string_prepend(self, action):
        self._kve(action, self._client.string_prepend)

    def _action_string_append(self, action):
        self._kve(action, self._client.string_append)

    def _action_list_lpush(self, action):
        self._kve(action, self._client.list_lpush)

    def _action_list_rpush(self, action):
        self._kve(action, self._client.list_rpush)

    def _action_set_add(self, action):
        self._kve(action, self._client.set_add)

    def _action_set_remove(self, action):
        self._kve(action, self._client.set_remove)

    def _action_set_intersect(self, action):
        self._kve(action, self._client.set_intersect)

    def _action_set_union(self, action):
        self._kve(action, self._client.set_union)

    def _action_map_add(self, action):
        self._kve(action, self._client.map_add)

    def _action_map_remove(self, action):
        self._kve(action, self._client.map_remove)

    def _action_map_atomic_add(self, action):
        self._kve(action, self._client.map_atomic_add)

    def _action_map_atomic_sub(self, action):
        self._kve(action, self._client.map_atomic_sub)

    def _action_map_atomic_mul(self, action):
        self._kve(action, self._client.map_atomic_mul)

    def _action_map_atomic_div(self, action):
        self._kve(action, self._client.map_atomic_div)

    def _action_map_atomic_mod(self, action):
        self._kve(action, self._client.map_atomic_mod)

    def _action_map_atomic_and(self, action):
        self._kve(action, self._client.map_atomic_and)

    def _action_map_atomic_or(self, action):
        self._kve(action, self._client.map_atomic_or)

    def _action_map_atomic_xor(self, action):
        self._kve(action, self._client.map_atomic_xor)

    def _action_map_string_prepend(self, action):
        self._kve(action, self._client.map_string_prepend)

    def _action_map_string_append(self, action):
        self._kve(action, self._client.map_string_append)

    def _action_search(self, action):
        self._check_fields(action, 'space', 'predicate', 'expected')
        expected = action['expected']
        try:
            returned = [x for x in self._client.search(action['space'], self._to_predicate(action['predicate']))]
            self._compare_search_results(action['expected'], returned)
        except hyperclient.HyperClientException as e:
            self._compare_exception(expected, e)

    '''
    def group_del(self, bytes space, dict predicate):
        async = self.async_group_del(space, predicate)
        return async.wait()

    def count(self, bytes space, dict predicate, bool unsafe=False):
        async = self.async_count(space, predicate, unsafe)
        return async.wait()

    def search(self, bytes space, dict predicate):
        return Search(self, space, predicate)

    def sorted_search(self, bytes space, dict predicate, bytes sort_by, long limit, bytes compare):
        return SortedSearch(self, space, predicate, sort_by, limit, compare)
        '''

    def _check_fields(self, action, *args):
        for arg in args:
            if arg not in action:
                raise RuntimeError('missing field \'%s\'' % arg)

    def _compare_success(self, expected, returned):
        if isinstance(expected, bytes):
            error  = 'should have raised an exception but returned a value instead\n'
            error += '    expected:  %r\n' % expected
            error += '    returned:  %r' % returned
            raise RuntimeError(error)
        if expected in (True, False, None):
            if expected != returned:
                error  = 'unexpected value returned\n'
                error += '    expected:  %r\n' % expected
                error += '    returned:  %r' % returned
                raise RuntimeError(error)
            return
        if not isinstance(expected, dict):
            raise RuntimeError('\'expected\' should be True, False, None or a dict')
        if not isinstance(returned, dict):
            error  = 'should have returned a dict\n'
            error += '    expected:  %r\n' % expected
            error += '    returned:  %r' % returned
            raise RuntimeError(error)
        if self._to_value(expected) != returned:
            error  = 'returned unexpected value\n'
            error += '    expected:  %r\n' % expected
            error += '    returned:  %r' % returned
            raise RuntimeError(error)

    def _compare_search_results(self, expected, returned):
        if isinstance(expected, bytes):
            error  = 'should have raised an exception but returned a value instead\n'
            error += '    expected:  %r\n' % expected
            error += '    returned:  %r' % returned
            raise RuntimeError(error)
        if not isinstance(expected, list):
            raise RuntimeError('\'expected\' should be a list')
        if not isinstance(returned, list):
            error  = 'should have returned a list\n'
            error += '    expected:  %r\n' % expected
            error += '    returned:  %r' % returned
            raise RuntimeError(error)
        if set([_frozen_dict(self._to_value(e)) for e in expected]) != set([_frozen_dict(x) for x in returned]):
            error  = 'returned unexpected value\n'
            error += '    expected:  %r\n' % expected
            error += '    returned:  %r' % returned
            raise RuntimeError(error)

    def _compare_exception(self, expected, exception):
        if not isinstance(expected, bytes):
            raise RuntimeError('caused an exception:  %r' % exception.symbol())
        elif expected != exception.symbol():
            raise RuntimeError('should have raised %r but raised %r instead' % (expected, exception.symbol()))

    def _to_value(self, valuedict):
        ret = {}
        if not isinstance(valuedict, dict):
            raise RuntimeError('%r should be a value dict' % valuedict)
        for attrname, value in valuedict.iteritems():
            if not isinstance(attrname, bytes):
                raise RuntimeError('value dict\'s keys should be bytes')
            if not isinstance(value, dict):
                raise RuntimeError('value dict\'s values should be dictionaries')
            if set(value.keys()) != set(['type', 'value']):
                raise RuntimeError('value dict\'s inner dictionaries must have \'type\' and \'value\' as keys')
            if value['type'] in HyperDexJSONBridge.VALUE_IDENTITY:
                ret[attrname] = value['value']
            elif value['type'] in HyperDexJSONBridge.VALUE_CONVERT:
                ret[attrname] = HyperDexJSONBridge.VALUE_CONVERT[value['type']](value['value'])
            else:
                raise RuntimeError('value dict\'s \'type\' is invalid')
        return ret

    def _to_predicate(self, predicatedict):
        ret = {}
        if not isinstance(predicatedict, dict):
            raise RuntimeError('%r should be a predicate dict' % predicatedict)
        for attrname, predicate in predicatedict.iteritems():
            if not isinstance(attrname, bytes):
                raise RuntimeError('predicate dict\'s keys should be bytes')
            if not isinstance(predicate, dict):
                raise RuntimeError('predicate dict\'s values should be dictionaries')
            if set(predicate.keys()) not in (set(['equality']),):
                raise RuntimeError('predicate dict\'s inner dictionaries must have \'equality\' key')
            if 'equality' in predicate:
                if set(predicate['equality'].keys()) != set(['type', 'value']):
                    raise RuntimeError('predicate dict\'s equality predicate must have \'type\' and \'value\' as keys')
                if predicate['equality']['type'] in HyperDexJSONBridge.VALUE_IDENTITY:
                    ret[attrname] = predicate['equality']['value']
                elif predicate['equality']['type'] in HyperDexJSONBridge.VALUE_CONVERT:
                    ret[attrname] = HyperDexJSONBridge.VALUE_CONVERT[predicate['equality']['type']](predicate['equality']['value'])
                else:
                    raise RuntimeError('predicate dict\'s \'type\' is invalid')
        return ret


def main(argv):
    import argparse
    import time
    parser = argparse.ArgumentParser()
    parser.add_argument('--host', default='127.0.0.1', metavar='HOST',
                        help='Address of the coordinator')
    parser.add_argument('--port', type=int, default=1234, metavar='PORT',
                        help='Port for the coordinator')
    parser.add_argument('--read', default=None, metavar='FILE',
                        help='The script to read from')
    parser.add_argument('--no-time', default=False, action='store_false',
                        help='Don\'t print statistics about running time')
    parser.add_argument('--time', default=False, action='store_true',
                        help='Print statistics about running time')
    args = parser.parse_args(argv)
    hdjb = HyperDexJSONBridge(args.host, args.port)
    json_source = None
    if args.read is None:
        json_source = sys.stdin
    else:
        json_source = open(args.read)
    lineno = 1
    comments = 0
    try:
        start = time.time()
        for line in json_source:
            line = line.strip()
            if not line.startswith('#') and len(line) > 0:
                action = simplejson.loads(line)
                hdjb.perform_action(action)
            else:
                comments += 1
            lineno += 1
        end = time.time()
        if args.time:
            ops = lineno - comments - 1
            secs = end - start
            opssec = ops / secs
            print('%i ops in %f seconds (%f ops/second)' % (ops, secs, opssec))
    except simplejson.decoder.JSONDecodeError as e:
        print('error on line %i:  %s' % (lineno, e), file=sys.stderr)
        sys.exit(-1)
    except RuntimeError as e:
        print('error on line %i:  %s' % (lineno, e), file=sys.stderr)
        sys.exit(-1)


if __name__ == '__main__':
    try:
        sys.exit(main(sys.argv[1:]))
    except Exception as e:
        import traceback
        print('uncaught exception: ', e, file=sys.stderr)
        traceback.print_exc(file=sys.stderr)
        sys.exit(-1)