Last active
September 4, 2015 21:16
-
-
Save mnunberg/ea58edc4303492bd0985 to your computer and use it in GitHub Desktop.
n1qlback file generator
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment