n1qlback file generator

gen_stmt_file.py

#!/usr/bin/env python
from argparse import ArgumentParser
import hashlib
import json
import sys
import os

from couchbase.bucket import Bucket


class Prepared(object):
    def __init__(self, qstr, connstr):
        self._qs = qstr
        self._connstr = connstr
        self._prepared = None
        self._name = None

    def _do_prepare(self):
        if self._prepared:
            return

        cb = Bucket(self._connstr)
        self._prepared = cb.n1ql_query('PREPARE ' + self._qs).get_single_result()

    @property
    def name(self):
        self._do_prepare()
        return self._prepared['name']

    @property
    def encoded_plan(self):
        self._do_prepare()
        return self._prepared['encoded_plan']

    @property
    def full_plan(self):
        self._do_prepare()
        return self._prepared


ap = ArgumentParser()

ap.add_argument('-U', '--connspec', help='Connection string',
                default='couchbase://localhost/default')
ap.add_argument('-s', '--statement', help='Path to statement string',
                required=True)
ap.add_argument('-a', '--arguments', help='Path to argument array',
                required=True)

ap.add_argument('-p', '--prefix', help='Base string to use for output files',
                default='query-')

ap.add_argument('-o', '--output-dir', help='Output directory')
ap.add_argument('-f', '--format', help='Output format',
                choices=('string', 'name', 'plan', 'encoded', 'all'),
                default='string')

options = ap.parse_args()

with open(options.statement, 'r') as sfp:
    query_string = sfp.read()

with open(options.arguments, 'r') as afp:
    args = json.load(afp)

# Strip whitespace from query
query_string = query_string.replace("\n", " ").rstrip().lstrip()
formats = []
if options.format == 'all':
    formats = ['string', 'name', 'plan', 'encoded']
else:
    formats = [options.format]

def make_string_dict(qs, args):
    return {'args': args, 'statement': qs}

OUTPUTS = {}
PREPARED = Prepared(query_string, options.connspec)

if 'string' in formats:
    OUTPUTS['string'] = make_string_dict(query_string, args)

if 'plan' in formats:
    OUTPUTS['plan'] = {'args': args, 'prepared': PREPARED.full_plan}

if 'name' in formats:
    OUTPUTS['named'] = {'args': args, 'prepared': PREPARED.name}

if 'encoded' in formats:
    OUTPUTS['encoded'] = {
        'args': args,
        'encoded_plan': PREPARED.encoded_plan,
        'prepared': PREPARED.name
    }

if len(formats) == 1: # Not all
    for v in OUTPUTS.values():
        json.dump(v, sys.stdout)
        print ""
        sys.exit(0)

if options.output_dir:
    if not os.path.exists(options.output_dir):
        os.makedirs(options.output_dir)
    output_dir = options.output_dir
else:
    output_dir = '.'

for name, body in OUTPUTS.items():
    name = options.prefix + name
    pth = os.path.join(output_dir, name + '.json')
    with open(pth, 'w') as fp:
        json.dump(body, fp)