Skip to content
This repository

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Fetching contributors…

Cannot retrieve contributors at this time

executable file 237 lines (196 sloc) 8.91 kb
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237
#!/usr/bin/env python

import logging
import simplejson as json
import time
import urllib

import pump
import pump_mc


class CBSink(pump_mc.MCSink):
    DDOC_HEAD = "_design/"

    """Smart client sink to couchbase cluster."""
    def __init__(self, opts, spec, source_bucket, source_node,
                 source_map, sink_map, ctl, cur):
        super(CBSink, self).__init__(opts, spec, source_bucket, source_node,
                                     source_map, sink_map, ctl, cur)

        self.rehash = opts.extra.get("rehash", 0)

    def scatter_gather(self, mconns, batch):
        sink_map_buckets = self.sink_map['buckets']
        if len(sink_map_buckets) != 1:
            return "error: CBSink.run() expected 1 bucket in sink_map", None, None

        vbuckets_num = len(sink_map_buckets[0]['vBucketServerMap']['vBucketMap'])
        vbuckets = batch.group_by_vbucket_id(vbuckets_num, self.rehash)

        # Scatter or send phase.
        for vbucket_id, msgs in vbuckets.iteritems():
            rv, conn = self.find_conn(mconns, vbucket_id, msgs)
            if rv != 0:
                return rv, None, None
            rv = self.send_msgs(conn, msgs, self.operation(),
                                vbucket_id=vbucket_id)
            if rv != 0:
                return rv, None, None

        # Yield to let other threads do stuff while server's processing.
        time.sleep(0.01)

        retry_batch = None
        need_refresh = False

        # Gather or recv phase.
        for vbucket_id, msgs in vbuckets.iteritems():
            rv, conn = self.find_conn(mconns, vbucket_id, msgs)
            if rv != 0:
                return rv, None, None
            rv, retry, refresh = self.recv_msgs(conn, msgs, vbucket_id=vbucket_id)
            if rv != 0:
                return rv, None, None
            if retry:
                retry_batch = batch
            if refresh:
                need_refresh = True

        if need_refresh:
            self.refresh_sink_map()

        return 0, retry_batch, retry_batch and not need_refresh

    @staticmethod
    def can_handle(opts, spec):
        return (spec.startswith("http://") or
                spec.startswith("couchbase://"))

    @staticmethod
    def check_source(opts, source_class, source_spec, sink_class, sink_spec):
        if (source_spec.startswith("http://") or
            source_spec.startswith("couchbase://")):
            return None
        return pump.Sink.check_source(opts, source_class, source_spec,
                                      sink_class, sink_spec)

    @staticmethod
    def check(opts, spec, source_map):
        rv, sink_map = pump.rest_couchbase(opts, spec)
        if rv != 0:
            return rv, None

        rv, source_bucket_name = pump.find_source_bucket_name(opts, source_map)
        if rv != 0:
            return rv, None
        rv, sink_bucket_name = pump.find_sink_bucket_name(opts, source_bucket_name)
        if rv != 0:
            return rv, None

        # Adjust sink_map['buckets'] to have only our sink_bucket.
        sink_buckets = [bucket for bucket in sink_map['buckets']
                        if bucket['name'] == sink_bucket_name]
        if not sink_buckets:
            return "error: missing bucket-destination: " + sink_bucket_name + \
                " at destination: " + spec + \
                "; perhaps your username/password is missing or incorrect", None
        if len(sink_buckets) != 1:
            return "error: multiple buckets with name: " + sink_bucket_name + \
                " at destination: " + spec, None
        sink_map['buckets'] = sink_buckets

        return 0, sink_map

    def refresh_sink_map(self):
        """Grab a new vbucket-server-map."""
        logging.warn("refreshing sink map: %s" % (self.spec))
        rv, new_sink_map = CBSink.check(self.opts, self.spec, self.source_map)
        if rv == 0:
            self.sink_map = new_sink_map
        return rv

    @staticmethod
    def consume_design(opts, sink_spec, sink_map,
                       source_bucket, source_map, source_design):
        if not source_design:
            return 0

        try:
            sd = json.loads(source_design)
        except ValueError, e:
            return "error: could not parse source design; exception: %s" % (e)
        if not sd:
            return 0

        if (not sink_map['buckets'] or
            len(sink_map['buckets']) != 1 or
            not sink_map['buckets'][0] or
            not sink_map['buckets'][0]['name']):
            return "error: design sink incorrect sink_map bucket"

        spec_parts = pump.parse_spec(opts, sink_spec, 8091)
        if not spec_parts:
            return "error: design sink no spec_parts: " + sink_spec

        sink_bucket = sink_map['buckets'][0]
        sink_nodes = pump.filter_bucket_nodes(sink_bucket, spec_parts) or \
            sink_bucket['nodes']
        if not sink_nodes:
            return "error: design sink nodes missing"

        couch_api_base = sink_nodes[0].get('couchApiBase')
        if not couch_api_base:
            return "error: cannot restore bucket design" \
                " on a couchbase cluster that does not support couch API;" \
                " the couchbase cluster may be an older, pre-2.0 version;" \
                " please check your cluster URL: " + sink_spec

        host, port, user, pswd, path = \
            pump.parse_spec(opts, couch_api_base, 8092)

        if user is None:
            user = spec_parts[2] # Default to the main REST user/pwsd.
            pswd = spec_parts[3]

        for row in sd:
            logging.debug("design_doc row: " + str(row))

            doc = row.get('doc', None)
            if not doc:
                return "error: missing design doc in row: %s" % (row)

            if 'json' in doc and 'meta' in doc:
                js = doc['json']
                id = doc['meta'].get('id', None)
                if not id:
                    return "error: missing id for design doc: %s" % (row)
            else:
                # Handle design-doc from 2.0DP4.
                js = doc
                if '_rev' in js:
                    del js['_rev']
                id = row.get('id', None)
                if not id:
                    return "error: missing id for row: %s" % (row)

            js_doc = json.dumps(js)
            if id.startswith(CBSink.DDOC_HEAD):
                id = CBSink.DDOC_HEAD + urllib.quote(id[len(CBSink.DDOC_HEAD):], '')
            else:
                id = urllib.quote(id, '')
            logging.debug("design_doc: " + js_doc)
            logging.debug("design_doc id: " + id + " at: " + path + "/" + id)

            try:
                err, conn, response = \
                    pump.rest_request(host, int(port), user, pswd,
                                      path + "/" + id, method='PUT', body=js_doc,
                                      reason="consume_design")
                if conn:
                    conn.close()
                if err:
                    return ("error: could not restore design doc id: %s" +
                            "; response: %s; err: %s") % (id, response, err)
            except Exception, e:
                return ("error: design sink exception: %s" +
                        "; couch_api_base: %s") % (e, couch_api_base)

            logging.debug("design_doc created at: " + path + "/" + id)

        return 0

    def find_conn(self, mconns, vbucket_id, msgs):
        bucket = self.sink_map['buckets'][0]

        vBucketMap = bucket['vBucketServerMap']['vBucketMap']
        serverList = bucket['vBucketServerMap']['serverList']

        if vbucket_id > len(vBucketMap):
            return "error: map missing vbucket_id: " + str(vbucket_id) + \
                "; perhaps your source does not have vbuckets" + \
                "; if so, try using moxi (HOST:11211) as a destination", None

        # Primary server for a vbucket_id is the 0'th entry.
        host_port = serverList[vBucketMap[vbucket_id][0]]

        conn = mconns.get(host_port, None)
        if not conn:
            host, port = host_port.split(':')
            user = bucket['name']
            pswd = bucket['saslPassword']
            rv, conn = CBSink.connect_mc(host, port, user, pswd)
            if rv != 0:
                logging.error("error: CBSink.connect() for send: " + rv)
                return rv, None
            mconns[host_port] = conn
            #check if we need to calll hello command
            for i, msg in enumerate(msgs):
                msg_format_length = len(msg)
                if msg_format_length > 8:
                    try:
                        conn.hello()
                    except Exception, e:
                        logging.warn("fail to call hello command, maybe it is not supported")
                        pass
                break
        return 0, conn
Something went wrong with that request. Please try again.