Permalink
Switch branches/tags
Nothing to show
Find file
Fetching contributors…
Cannot retrieve contributors at this time
executable file 133 lines (113 sloc) 3.76 KB
#!/usr/bin/env python
# -*- python -*-
import sys, zipfile, os, os.path
import getopt
import json
import tempfile
import re
from couchbase import client
def usage(err=0):
print >> sys.stderr, """
Usage: %s [-u user [-p password]] [-n host:port] [-b bucket] <directory>|zipfile
Example:
%s -u user -p secret9876 -n localhost:8091 -b mybucket datadir
""" % (os.path.basename(sys.argv[0]), (os.path.basename(sys.argv[0])))
sys.exit(err)
def parse_args(args):
user = None
pswd = None
host = 'localhost:8091'
bucket = 'default'
try:
opts, args = getopt.getopt(args, 'hu:p:n:b:', ['help'])
except getopt.GetoptError, e:
usage("ERROR: " + e.msg)
for (o, a) in opts:
if o == '--help' or o == '-h':
usage()
elif o == '-u':
user = a
elif o == '-p':
pswd = a
elif o == '-n':
host = a
elif o == '-b':
bucket = a
else:
usage("ERROR: unknown option - " + o)
if not args or len(args) < 1:
usage("ERROR: missing upload directory or zipped file.")
return user, pswd, host, bucket, args
def save_doc(bucket, dockey, fp, views):
buf = fp.read()
result = json.loads(buf)
if isinstance(result, dict):
if '_id' not in result:
bucket.set(dockey, 0, 0, json.dumps(result))
else:
try:
doc_id = bucket.save(result)
print "just now saving", doc_id
except:
doc_id = "_design/testing"
if result['_id'] and 'views' in result:
for key in result['views'].iterkeys():
viewpath = result['_id'] + '/_view/' + key
views.append(viewpath)
def gen_dockey(filename):
fileslashed = re.split(".json", filename)
fileinlist = re.split("/", fileslashed[0])
file_for_key = fileinlist.pop()
return file_for_key
def listFiles(bucket, dir, views):
basedir = dir
#print "Files in ", os.path.abspath(dir), ": "
subdirlist = []
for item in os.listdir(dir):
if os.path.isfile(os.path.join(basedir, item)):
with open(os.path.join(basedir, item), 'r') as fp:
print "working with ", item
dockey = gen_dockey(item)
save_doc(bucket, dockey, fp, views)
else:
subdirlist.append(os.path.join(basedir, item))
for subdir in subdirlist:
listFiles(bucket, subdir, views)
def unzip_file_and_upload(bucket,file, views):
zfobj = zipfile.ZipFile(file)
for name in zfobj.namelist():
if not name.endswith('/'):
print 'working with ', name
dockey = gen_dockey(name)
temp = tempfile.NamedTemporaryFile(delete=False)
fname = temp.name
temp.write(zfobj.read(name))
temp.close()
with open(fname, 'r') as fp:
save_doc(bucket, dockey, fp, views)
os.remove(fname)
def populate_docs(bucket, dir, views):
if dir.endswith('.zip'):
unzip_file_and_upload(bucket, dir, views)
else:
listFiles(bucket, dir, views)
def main():
user, pswd, host, bucket, args = parse_args(sys.argv[1:])
print user,pswd,host,bucket, args
cb = client.Server(host, user, pswd)
try:
newbucket = cb.create(bucket, ram_quota_mb=100, replica=1)
except:
newbucket = cb[bucket]
#upload documents
dir = args[0]
views = []
populate_docs(newbucket, dir, views)
# execute views at least once
try:
for viewpath in views:
rows = newbucket.view(viewpath)
except:
#do nothing, maybe view result is not ready yet
print "View result will be ready after a while"
if __name__ == '__main__': main()