attach_large.py

#!/usr/bin/env python

import requests
import sys, socket, hashlib, base64, itertools, time
import couchdb


DEBUG = True
PORT = 15984
DATABASE = 'db'

def lg(*args):
    if DEBUG:
        msg = ' '.join([str(a) for a in args])
        sys.stderr.write('> ' + msg + '\n')

def le(*args):
    msg = ' '.join([str(a) for a in args])
    sys.stderr.write('!ERROR: ' + msg + '\n')

def setup_db(dbname, user, password, port):
    url = 'http://%s:%s@127.0.0.1:%d'% (user, password, port)
    lg("db url:", url)
    srv=couchdb.Server(url)
    if dbname in srv:
        lg("Deleting previous database:", dbname)
        srv.delete(dbname)
    db = srv.create(dbname)
    lg("Created DB:", dbname)
    return db

def get_mp_body_bak(att_size):
    att = 'x' * att_size
    att_md5  = hashlib.md5(att).digest().encode('base64').strip()
    return '\r\n'.join([
        '--195044bdf7abd8b23cce7b30d39a19dd',
        'Content-Type: application/json',
        '',
        '''{"_id":"doc2","_rev":"1-2df9eed63e6f4df24c6a7b593adfc195","_revisions":{"start":1,"ids":["2df9eed63e6f4df24c6a7b593adfc195"]},
           "_attachments":{"att2":{"content_type":"app/binary","revpos":1, "digest":"md5-%s", "length":%s, "follows":true}}}
        ''' % (att_md5, att_size),
        '--195044bdf7abd8b23cce7b30d39a19dd',
        'Content-Disposition: attachment; filename="att2"',
        'Content-Type: app/binary',
        '',
        att,
        '--195044bdf7abd8b23cce7b30d39a19dd--',
        ])

def get_mp_body(att_size):
    att = 'x' * att_size
    return '\r\n'.join([
        '--bound',
        'Content-Type: application/json',
        '',
        '''{"_id":"doc2","_rev":"1-2df9eed63e6f4df24c6a7b593adfc195","_revisions":{"start":1,"ids":["2df9eed63e6f4df24c6a7b593adfc195"]},
           "_attachments":{"att2":{"content_type":"app/binary","revpos":1, "length":%s, "follows":true}}}
        ''' %  att_size,
        '--bound',
        'Content-Disposition: attachment; filename="att2"',
        'Content-Type: app/binary',
        '',
        att,
        '--bound--',
        ])

def get_basic_auth(uname, upass):
    return 'Basic %s' % base64.b64encode("%s:%s" % (uname, upass))

def get_mp_request(dbname, port, att_size, user, pwd):
    body = get_mp_body(att_size)
    return '\r\n'.join([
        'PUT /%s/doc2?new_edits=false HTTP/1.1' % dbname,
        'Authorization : %s' % get_basic_auth(user, pwd),
        'Content-Type: multipart/related; boundary="bound"',
        'Content-Length: %d' % (len(body)),
        'Accept: application/json',
        #'Host: 127.0.0.1:%d' % port,
        '',
        body,
        ])

def _batchit(it, batchsize):
    while True:
        batch = [x for (_, x) in itertools.izip(xrange(batchsize), it)]
        if not batch:
            raise StopIteration
        yield ''.join(batch)


def send(s, data, mintime=0):
    if mintime == 0:
        lg("sending all data", len(data), "bytes using sendall()")
        return s.sendall(data)
    else:
        return send_paced(s, data, mintime)

def get_chunk_size(dlen):
    if dlen > 2**20:
        return 4029
    if dlen > 2**16:
        return 256
    return 128

def send_paced(s, data, mintime):
    dlen = len(data)
    chunk_size = get_chunk_size(dlen)
    chunk_count = dlen / chunk_size
    dt = mintime / chunk_count
    lg("chunk size:", chunk_size, "chunk count:", chunk_count)
    for chunk in _batchit(iter(data), chunk_size):
        lg("sending data chunk:",len(chunk), "then sleep %0.3f sec" % dt)
        send_t0 = time.time()
        res = s.sendall(chunk)
        send_dt = time.time() - send_t0
        lg("sent data in %0.3f sec" % send_dt, " res:", res)
        time.sleep(dt)

def send_recv(s, data, mintime):
    err = None
    try:
        send(s, data, mintime)
    except socket.error, e:
        err = e
        le(" socket Error:", e)
    lg("Sent data:",len(data),"bytes, receiving")
    rdata = s.recv(8192)
    lg("Received:",len(rdata),"bytes")
    lg("")
    sys.stdout.write(rdata)
    return (rdata,err)

def main(dbname, user, password, port, att_size, mintime, atype):
    lg("Settings DB:", dbname, "Port:", port, "User:", user, "Pass:", password, "Size:",att_size, "Type:", atype)
    db = setup_db(dbname, user, password, port)
    if atype == 'mpexact':
        req = get_mp_request(dbname, port, att_size, user, password)
        s = socket.socket()
        s.connect(('127.0.0.1',port))
        lg("Connected to localhost:",port)
        raw_input("continue")
        rdata,err = send_recv(s, req, mintime)
        if len(rdata)==0:
            lg("received 0 bytes, socket closed")
            return 0
        if err:
            lg("error sending data", err)
            return 1
        if 'Connection: close' in rdata:
            return 1
    elif atype == 'single':
        doc2 = {'_id':'doc2'}
        db.save(doc2)
        res = db.put_attachment(doc2, att_size * 'x', filename="att2", content_type="app/binary")
        #lg("PUT attachment res:", dict(db['doc2']))
    elif atype == 'singlechunked':
        '''PUT /db/doc2/att2?rev=1-967a00dff5e02add41819138abb3284d HTTP/1.1
Host: 127.0.0.1:15984
Transfer-Encoding: chunked
Content-Type: app/binary
Authorization: Basic YWRtOnBhc3M=
Accept: application/json

xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx...
'''
        doc2 = {'_id':'doc2'}
        db.save(doc2)
        with open('/tmp/att2.bin', 'w') as fh:
            fh.write(att_size * 'z')
        with open('/tmp/att2.bin') as fh:
            res = db.put_attachment(doc2, fh, filename="att2", content_type="app/binary")
            #lg("PUT attachment from file:", dict(db['doc2']))
    elif atype == 'inline':
        data = att_size * 'z'
        doc2 = {
            '_id':'doc2',
            '_attachments': {
                'att2':{
                    'content_type':'app/binary',
                    'data': data.encode('base64').strip()
                }
            }
        }
        db.save(doc2)
    else:
        raise Exception("Invalid type of request %s" % atype)


if __name__=='__main__':
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument('--db', default=DATABASE)
    p.add_argument('--port',type=int, default=PORT)
    p.add_argument('--user', default='adm')
    p.add_argument('--password', default='pass')
    p.add_argument('--size', type=int, default=64)
    p.add_argument('--debug', action="store_true")
    p.add_argument('--mintime', type=float, default=0)
    p.add_argument('--type', default="mpexact")
    args = p.parse_args()
    if args.debug:
        DEBUG = True
    sys.exit(main(args.db, args.user, args.password, args.port, args.size, args.mintime, args.type))