rpc cli client

rpc.py

#!/usr/bin/env python

import argparse
import base64
import code
import json
import pprint
import requests


def request(sql, method='run', request_id=1, timeout=None):
    url = 'http://localhost:8580/jsonrpc'
    headers = {'content-type': 'application/json'}
    encoded = base64.standard_b64encode(sql.encode('utf-8')).decode('utf-8')
    data = {
        'jsonrpc': '2.0',
        'method': method,
        'params': {
            'sql': encoded,
            'name': 'foo',
            'timeout': timeout
        },
        'id': request_id
    }
    print('request={}'.format(data))

    result = requests.post(url, headers=headers, data=json.dumps(data))
    if result.ok:
        return result.json()
    else:
        return result


def run(sql, request_id=1, timeout=None):
    return request(sql=sql, method='run', request_id=request_id,
                   timeout=timeout)


def compile(sql, request_id=1, timeout=None):
    return request(sql=sql, method='compile', request_id=request_id,
                   timeout=timeout)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--method', default='run')
    parser.add_argument('--sql', default=None)
    parser.add_argument('--id', default=1, type=int, dest='request_id')
    parser.add_argument('--timeout', default=None, type=lambda x: int(x) if x is not None else None)
    return parser.parse_args()


def interactive(defaults):
    del defaults.sql

    def run(sql, request_id=defaults.request_id, timeout=defaults.timeout):
        return request(sql=sql, method='run', request_id=request_id,
                       timeout=timeout)

    def compile(sql, request_id=defaults.request_id, timeout=defaults.timeout):
        return request(sql=sql, method='compile', request_id=request_id,
                       timeout=timeout)

    try:
        import IPython
        IPython.embed()
        return
    except ImportError:
        pass

    ctx = {'defaults': defaults, 'run': run, 'compile': compile}

    return code.interact(local=ctx)


def main():
    parsed = parse_args()
    if parsed.sql is None:
        interactive(parsed)
        return

    result = request(
        method=parsed.method,
        sql=parsed.sql,
        request_id=parsed.request_id,
        timeout=parsed.timeout)
    pprint.pprint(result)

if __name__ == '__main__':
    main()