class MemcachedClient(object):
    """Simple memcached client."""

    vbucket_id = 0

    def __init__(self, host='127.0.0.1', port=11211, family=socket.AF_UNSPEC, use_ssl=False, verify=True, cacert=None):
        self.host = host
        self.port = port
        sock_error = None
        for info in socket.getaddrinfo(self.host, self.port, family,
                                       socket.SOCK_STREAM):
            _family, socktype, proto, _, sockaddr = info
            try:
                sock = socket.socket(_family, socktype, proto)
                sock.settimeout(10)
                if use_ssl:
                    cert_req = ssl.CERT_REQUIRED
                    if not verify:
                        cert_req = ssl.CERT_NONE
                    sock = ssl.wrap_socket(sock, server_side=False, do_handshake_on_connect=True, cert_reqs=cert_req,
                                           ca_certs=cacert)

                self.s = sock
                self.s.connect_ex(sockaddr)
                break
            except socket.error as e:
                sock_error = e
                # If we get here socket objects will be close()d via
                # garbage collection.
                pass
        else:
            # Didn't break from the loop, re-raise the last error
            raise sock_error

        self.r = random.Random()

    def close(self):
        self.s.close()

    def __del__(self):
        self.close()

    def _send_cmd(self, cmd: int, key: bytes, val: bytes, opaque: int, extra_header: bytes = b'', cas: int = 0,
                  extra_meta: bytes = b''):
        self._send_msg(cmd, key, val, opaque, extra_header=extra_header, cas=cas,
                       vbucket_id=self.vbucket_id, extra_meta=extra_meta)

    def _send_msg(self, cmd: int, key: bytes, val: bytes, opaque: int, extra_header: bytes = b'', cas: int = 0,
                  dtype: int = 0, vbucket_id: int = 0, extra_meta: bytes = b'',
                  fmt: str = REQ_PKT_FMT, magic: int = REQ_MAGIC_BYTE):
        msg: bytes = struct.pack(fmt, magic, cmd, len(key), len(extra_header), dtype, vbucket_id,
                                 len(key) + len(extra_header) + len(val) + len(extra_meta), opaque, cas)
        self.s.sendall(msg + extra_header + key + val + extra_meta)

    def _recv_msg(self) -> Tuple[int, int, int, int, int, int, bytes]:
        response = b''
        while len(response) < MIN_RECV_PACKET:
            data = self.s.recv(MIN_RECV_PACKET - len(response))
            if data == b'':
                raise EOFError("Got empty data (remote died?).")
            response += data
        assert len(response) == MIN_RECV_PACKET
        magic, cmd, keylen, extralen, dtype, errcode, remaining, opaque, cas = \
            struct.unpack(RES_PKT_FMT, response)

        rv = b''
        while remaining > 0:
            data = self.s.recv(remaining)
            if data == b'':
                raise EOFError("Got empty data (remote died?).")
            rv += data
            remaining -= len(data)

        assert (magic in (RES_MAGIC_BYTE, REQ_MAGIC_BYTE)), f'Got magic: {magic}'
        return cmd, errcode, opaque, cas, keylen, extralen, rv

    def _handle_keyed_response(self, myopaque: Union[int, None]) -> Tuple[int, int, int, int, int, bytes]:
        cmd, errcode, opaque, cas, keylen, extralen, rv = self._recv_msg()
        assert myopaque is None or opaque == myopaque, f'expected opaque {myopaque:x}, got {opaque:x}'
        if errcode != 0:
            raise MemcachedError(errcode, str(rv))
        return cmd, opaque, cas, keylen, extralen, rv

    def _handle_single_response(self, myopaque: Union[int, None]) -> Tuple[int, int, bytes]:
        cmd, opaque, cas, keylen, extralen, data = self._handle_keyed_response(myopaque)
        return opaque, cas, data

    def _do_cmd(self, cmd: int, key: bytes, val: bytes, extra_header: bytes = b'',
                cas: int = 0) -> Tuple[int, int, bytes]:
        """Send a command and await its response."""
        opaque = self.r.randint(0, 2**32)
        self._send_cmd(cmd, key, val, opaque, extra_header, cas)
        return self._handle_single_response(opaque)

    def _mutate(self, cmd: int, key: bytes, exp: int, flags: int, cas: int, val: bytes) -> Tuple[int, int, bytes]:
        return self._do_cmd(cmd, key, val, struct.pack(SET_PKT_FMT, flags, exp), cas)

    def _cat(self, cmd: int, key: bytes, cas: int, val: bytes) -> Tuple[int, int, bytes]:
        return self._do_cmd(cmd, key, val, b'', cas)

    def append(self, key: bytes, value: bytes, cas: int = 0) -> Tuple[int, int, bytes]:
        return self._cat(couchbaseConstants.CMD_APPEND, key, cas, value)

    def prepend(self, key: bytes, value: bytes, cas: int = 0) -> Tuple[int, int, bytes]:
        return self._cat(couchbaseConstants.CMD_PREPEND, key, cas, value)

    def __incrdecr(self, cmd, key, amt, init, exp):
        something, cas, val = self._do_cmd(cmd, key, b'',
                                           struct.pack(couchbaseConstants.INCRDECR_PKT_FMT, amt, init, exp))
        return struct.unpack(INCRDECR_RES_FMT, val)[0], cas

    def incr(self, key, amt=1, init=0, exp=0):
        """Increment or create the named counter."""
        return self.__incrdecr(couchbaseConstants.CMD_INCR, key, amt, init, exp)

    def decr(self, key, amt=1, init=0, exp=0):
        """Decrement or create the named counter."""
        return self.__incrdecr(couchbaseConstants.CMD_DECR, key, amt, init, exp)

    def _do_meta_cmd(self, cmd: int, key: bytes, value: bytes, cas: int, exp: int, flags: int, seqno: int,
                     remote_cas: int, options: Optional[int] = None, meta_len: Optional[int] = None) -> Tuple[int, int,
                                                                                                              bytes]:
        extra = b''
        if options is None and meta_len is None:
            extra = struct.pack('>IIQQ', flags, exp, seqno, remote_cas)
        if options is not None and meta_len is not None:
            extra = struct.pack('>IIQQIH', flags, exp, seqno, remote_cas, options, meta_len)
        if options is not None and meta_len is None:
            extra = struct.pack('>IIQQI', flags, exp, seqno, remote_cas, options)
        if options is None and meta_len is not None:
            extra = struct.pack('>IIQQH', flags, exp, seqno, remote_cas, meta_len)
        return self._do_cmd(cmd, key, value, extra, cas)

    def _do_rev_cmd(self, cmd: int, key: bytes, exp: int, flags: int,
                    value: bytes, rev: Tuple[int, bytes], cas: int = 0):
        seqno, revid = rev
        meta_data = struct.pack('>I', seqno) + revid
        meta_type = couchbaseConstants.META_REVID
        meta = (meta_type, meta_data)
        return self._do_meta_cmd(cmd, key, value, cas, exp, flags, seqno, 0)  # pylint: disable=no-value-for-parameter

    def set(self, key: bytes, exp: int, flags: int, val: bytes) -> Tuple[int, int, bytes]:
        """Set a value in the memcached server."""
        return self._mutate(couchbaseConstants.CMD_SET, key, exp, flags, 0, val)

    def set_with_meta(self, key: bytes, value: bytes, exp: int, flags: int, seqno: int, remote_cas: int,
                      options: Optional[int] = None, meta_len: Optional[int] = None) -> Tuple[int, int, bytes]:
        """Set a value and its meta data in the memcached server."""
        return self._do_meta_cmd(couchbaseConstants.CMD_SET_WITH_META,
                                 key, value, 0, exp, flags, seqno, remote_cas, options, meta_len)

    def set_with_rev(self, key: bytes, exp: int, flags: int, value: bytes, rev: Tuple[int, bytes]):
        """Set a value and its revision in the memcached server."""
        return self._do_rev_cmd(couchbaseConstants.CMD_SET_WITH_META,
                                key, exp, flags, value, rev)

    def add(self, key: bytes, exp: int, flags: int, val: bytes):
        """Add a value in the memcached server iff it doesn't already exist."""
        return self._mutate(couchbaseConstants.CMD_ADD, key, exp, flags, 0, val)

    def add_with_meta(self, key: bytes, value: bytes, exp: int, flags: int, seqno: int, remote_cas: int):
        return self._do_meta_cmd(couchbaseConstants.CMD_ADD_WITH_META,
                                 key, value, 0, exp, flags, seqno, remote_cas)

    def add_with_rev(self, key: bytes, exp: int, flags: int, value: bytes, rev: Tuple[int, bytes]):
        return self._do_rev_cmd(couchbaseConstants.CMD_ADD_WITH_META,
                                key, exp, flags, value, rev)

    def replace(self, key: bytes, exp: int, flags: int, val: bytes):
        """Replace a value in the memcached server iff it already exists."""
        return self._mutate(couchbaseConstants.CMD_REPLACE, key, exp, flags, 0,
                            val)

    def observe(self, key: bytes, vbucket: int) -> Tuple[int, int, int, bytes]:
        """Observe a key for persistence and replication."""
        value = struct.pack('>HH', vbucket, len(key)) + key
        opaque, cas, data = self._do_cmd(couchbaseConstants.CMD_OBSERVE, b'', value)
        rep_time = (cas & 0xFFFFFFFF)
        persist_time = (cas >> 32) & 0xFFFFFFFF
        persisted = struct.unpack('>B', bytes(data[4 + len(key)]))[0]
        return opaque, rep_time, persist_time, persisted

    def __parse_get(self, data: bytes, klen: int = 0) -> Tuple[int, int, bytes]:
        flags = struct.unpack(couchbaseConstants.GET_RES_FMT, data[:4])[0]
        return flags, data[1], data[4 + klen:]

    def get(self, key: bytes) -> Tuple[int, int, bytes]:
        """Get the value for a given key within the memcached server."""
        _, _, data = self._do_cmd(couchbaseConstants.CMD_GET, key, b'')
        return self.__parse_get(data)

    def get_meta(self, key: bytes) -> Tuple[int, int, int, int, int]:
        """Get the metadata for a given key within the memcached server."""
        opaque, cas, data = self._do_cmd(couchbaseConstants.CMD_GET_META, key, b'')
        deleted = struct.unpack('>I', data[0:4])[0]
        flags = struct.unpack('>I', data[4:8])[0]
        exp = struct.unpack('>I', data[8:12])[0]
        seqno = struct.unpack('>Q', data[12:20])[0]
        return (deleted, flags, exp, seqno, cas)

    def getl(self, key: bytes, exp: int = 15):
        """Get the value for a given key within the memcached server."""
        _, _, data = self._do_cmd(couchbaseConstants.CMD_GET_LOCKED, key, b'',
                                  struct.pack(couchbaseConstants.GETL_PKT_FMT, exp))
        return self.__parse_get(data)

    def cas(self, key: bytes, exp: int, flags: int, old_val: int, val: bytes):
        """CAS in a new value for the given key and comparison value."""
        self._mutate(couchbaseConstants.CMD_SET, key, exp, flags, old_val, val)

    def touch(self, key: bytes, exp: int):
        """Touch a key in the memcached server."""
        return self._do_cmd(couchbaseConstants.CMD_TOUCH, key, b'', struct.pack(couchbaseConstants.TOUCH_PKT_FMT, exp))

    def gat(self, key: bytes, exp: int):
        """Get the value for a given key and touch it within the memcached server."""
        _, _, data = self._do_cmd(couchbaseConstants.CMD_GAT, key, b'',
                                  struct.pack(couchbaseConstants.GAT_PKT_FMT, exp))
        return self.__parse_get(data)

    def getr(self, key: bytes):
        """Get the value for a given key in a replica vbucket within the memcached server."""
        _, _, data = self._do_cmd(couchbaseConstants.CMD_GET_REPLICA, key, b'')
        return self.__parse_get(data, len(key))

    def version(self):
        """Get the value for a given key within the memcached server."""
        return self._do_cmd(couchbaseConstants.CMD_VERSION, b'', b'')

    def verbose(self, level: int):
        """Set the verbosity level."""
        return self._do_cmd(couchbaseConstants.CMD_VERBOSE, b'', b'',
                            extra_header=struct.pack(">I", level))

    def sasl_mechanisms(self):
        """Get the supported SASL methods."""
        return set(self._do_cmd(couchbaseConstants.CMD_SASL_LIST_MECHS, b'', b'')[2].split(b' '))

    def sasl_auth_start(self, mech: bytes, data: bytes) -> Tuple[int, int, bytes]:
        """Start a sasl auth session."""
        return self._do_cmd(couchbaseConstants.CMD_SASL_AUTH, mech, data)

    def sasl_auth_plain(self, user: str, password: str, foruser: str = '') -> Tuple[int, int, bytes]:
        """Perform plain auth."""
        return self.sasl_auth_start(b'PLAIN', b'\0'.join([foruser.encode(), user.encode(), password.encode()]))

    def stop_persistence(self) -> Tuple[int, int, bytes]:
        return self._do_cmd(couchbaseConstants.CMD_STOP_PERSISTENCE, b'', b'')

    def start_persistence(self) -> Tuple[int, int, bytes]:
        return self._do_cmd(couchbaseConstants.CMD_START_PERSISTENCE, b'', b'')

    def set_param(self, key: bytes, val: bytes, type: int) -> Tuple[int, int, bytes]:
        print("setting param:", key, val)
        type_bytes: bytes = struct.pack(couchbaseConstants.SET_PARAM_FMT, type)
        return self._do_cmd(couchbaseConstants.CMD_SET_PARAM, key, val, type_bytes)

    def set_vbucket_state(self, vbucket: int, state_name: str):
        assert isinstance(vbucket, int)
        self.vbucket_id = vbucket
        state = struct.pack(couchbaseConstants.VB_SET_PKT_FMT,
                            couchbaseConstants.VB_STATE_NAMES[state_name])
        return self._do_cmd(couchbaseConstants.CMD_SET_VBUCKET_STATE, b'', b'', state)

    def get_vbucket_state(self, vbucket: int):
        assert isinstance(vbucket, int)
        self.vbucket_id = vbucket
        return self._do_cmd(couchbaseConstants.CMD_GET_VBUCKET_STATE, b'', b'')

    def delete_vbucket(self, vbucket: int):
        assert isinstance(vbucket, int)
        self.vbucket_id = vbucket
        return self._do_cmd(couchbaseConstants.CMD_DELETE_VBUCKET, b'', b'')

    def evict_key(self, key: bytes):
        return self._do_cmd(couchbaseConstants.CMD_EVICT_KEY, key, b'')

    def get_multi(self, keys: Iterable[bytes]):
        """Get values for any available keys in the given iterable.

        Returns a dict of matched keys to their values."""
        opaqued = dict(enumerate(keys))
        terminal = len(opaqued) + 10
        # Send all of the keys in quiet
        for k, v in opaqued.items():
            self._send_cmd(couchbaseConstants.CMD_GETQ, v, b'', k)

        self._send_cmd(couchbaseConstants.CMD_NOOP, b'', b'', terminal)

        # Handle the response
        rv = {}
        done = False
        while not done:
            opaque, cas, data = self._handle_single_response(None)  # type: ignore
            if opaque != terminal:
                rv[opaqued[opaque]] = self.__parse_get((opaque, cas, data))  # type: ignore
            else:
                done = True

        return rv

    def set_multi(self, exp: int, flags: int, items):
        """Multi-set (using setq).

        Give me (key, value) pairs."""

        # If this is a dict, convert it to a pair generator
        if hasattr(items, 'iteritems'):
            items = items.items()

        opaqued = dict(enumerate(items))
        terminal = len(opaqued) + 10
        extra = struct.pack(SET_PKT_FMT, flags, exp)

        # Send all of the keys in quiet
        for opaque, kv in opaqued.items():
            self._send_cmd(couchbaseConstants.CMD_SETQ, kv[0], kv[1], opaque, extra)

        self._send_cmd(couchbaseConstants.CMD_NOOP, b'', b'', terminal)

        # Handle the response
        failed = []
        done = False
        while not done:
            try:
                opaque, cas, data = self._handle_single_response(None)  # type: ignore
                done = opaque == terminal
            except MemcachedError as e:
                failed.append(e)

        return failed

    def del_multi(self, items):
        """Multi-delete (using delq).

        Give me a collection of keys."""

        opaqued = dict(enumerate(items))
        terminal = len(opaqued) + 10
        extra = b''

        # Send all of the keys in quiet
        for opaque, k in opaqued.items():
            self._send_cmd(couchbaseConstants.CMD_DELETEQ, k, b'', opaque, extra)

        self._send_cmd(couchbaseConstants.CMD_NOOP, b'', b'', terminal)

        # Handle the response
        failed = []
        done = False
        while not done:
            try:
                opaque, cas, data = self._handle_single_response(None)
                done = opaque == terminal
            except MemcachedError as e:
                failed.append(e)

        return failed

    def stats(self, sub: bytes = b''):
        """Get stats."""
        opaque = self.r.randint(0, 2**32)
        self._send_cmd(couchbaseConstants.CMD_STAT, sub, b'', opaque)
        done = False
        rv = {}
        while not done:
            cmd, opaque, cas, klen, extralen, data = self._handle_keyed_response(None)
            if klen:
                rv[data[0:klen]] = data[klen:]
            else:
                done = True
        return rv

    def noop(self) -> Tuple[int, int, bytes]:
        """Send a noop command."""
        return self._do_cmd(couchbaseConstants.CMD_NOOP, b'', b'')

    def helo(self, features: List[int]) -> Tuple[int, int, Tuple[int, ...]]:
        """Send a hello command for feature checking"""
        r1, r2, value = self._do_cmd(couchbaseConstants.CMD_HELLO, f'couchbase-cli/{VERSION}'.encode(),
                                     struct.pack('>' + ('H' * len(features)), *features))
        # divide length of value by 2 to find out how many 2-byte elements exist
        return r1, r2, struct.unpack('>' + ('H' * int(len(value) / 2)), value)

    def delete(self, key: bytes, cas: int = 0) -> Tuple[int, int, bytes]:
        """Delete the value for a given key within the memcached server."""
        return self._do_cmd(couchbaseConstants.CMD_DELETE, key, b'', b'', cas)

    def flush(self, timebomb: int = 0) -> Tuple[int, int, bytes]:
        """Flush all storage in a memcached instance."""
        return self._do_cmd(couchbaseConstants.CMD_FLUSH, b'', b'',
                            struct.pack(couchbaseConstants.FLUSH_PKT_FMT, timebomb))

    def bucket_select(self, name: bytes) -> Tuple[int, int, bytes]:
        return self._do_cmd(couchbaseConstants.CMD_SELECT_BUCKET, name, b'')

    def restore_file(self, filename):
        """Initiate restore of a given file."""
        return self._do_cmd(couchbaseConstants.CMD_RESTORE_FILE, filename, b'', b'', 0)

    def restore_complete(self):
        """Notify the server that we're done restoring."""
        return self._do_cmd(couchbaseConstants.CMD_RESTORE_COMPLETE, b'', b'', b'', 0)

    def deregister_tap_client(self, tap_name):
        """Deregister the TAP client with a given name."""
        return self._do_cmd(couchbaseConstants.CMD_DEREGISTER_TAP_CLIENT, tap_name, b'', b'', 0)

    def reset_replication_chain(self):
        """Reset the replication chain."""
        return self._do_cmd(couchbaseConstants.CMD_RESET_REPLICATION_CHAIN, b'', b'', b'', 0)

    def audit(self, auditid, event):
        # print couchbaseConstants.CMD_AUDIT_PUT, auditid, event

        # return self._do_cmd(couchbaseConstants.CMD_AUDIT_PUT, '', event, \
        #                   struct.pack(AUDIT_PKT_FMT, auditid))
        return 0 class Batch(object):
    """Holds a batch of data being transfered from source to sink."""

    def __init__(self, source):
        self.source = source
        self.msgs: List[couchbaseConstants.BATCH_MSG] = []
        self.bytes: int = 0
        self.adjust_size: int = 0

    def append(self, msg: couchbaseConstants.BATCH_MSG, num_bytes: int):
        self.msgs.append(msg)
        self.bytes = self.bytes + num_bytes

    def size(self) -> int:
        return len(self.msgs)

    def msg(self, i: int) -> couchbaseConstants.BATCH_MSG:
        return self.msgs[i]

    def group_by_vbucket_id(self, vbuckets_num, rehash=0) -> Dict[int, List[couchbaseConstants.BATCH_MSG]]:
        """Returns dict of vbucket_id->[msgs] grouped by msg's vbucket_id."""
        g: Dict[int, List[couchbaseConstants.BATCH_MSG]] = defaultdict(list)
        for msg in self.msgs:
            cmd, vbucket_id, key = msg[:3]
            if vbucket_id == 0x0000ffff or rehash == 1:
                if self.source.opts.collection:
                    # Collections embeds the ID into the key field, but does not
                    # hash the ID as part of VB hashing
                    key = cb_bin_client.skip_collection_id(key)

                if isinstance(key, str):
                    key = key.encode()
                # Special case when the source did not supply a vbucket_id
                # (such as stdin source), so we calculate it.
                vbucket_id = ((zlib.crc32(key) >> 16) & 0x7FFF) % vbuckets_num
                msg = (cmd, vbucket_id) + msg[2:]  # type: ignore
            g[vbucket_id].append(msg)
        return g def filter_bucket_nodes(bucket: Dict[str, Any], spec_parts: Sequence[Any]) -> List[Any]:
    host, port = spec_parts[:2]
    if host in ['localhost', '127.0.0.1']:
        host = get_ip()
    # Convert from raw IPv6
    if ':' in host:
        host_port = f'[{host}]:{port!s}'
    else:
        host_port = f'{host}:{port!s}'

    # Hosts are case-insensitive but the url parser used by the CLI normalizes the input which turns it into lower
    # case (this is the spec_parts). To avoid any issues we will convert the hostnames we get from the REST API to lower
    # case. See MB-38453.
    return [n for n in bucket['nodes'] if n.get('hostname').lower() == host_port] class BFD:
    """Mixin for backup-file/directory EndPoint helper methods."""
    NUM_VBUCKET = 1024

    def bucket_name(self) -> str:
        return self.source_bucket['name']  # type: ignore # pylint: disable=no-member

    def node_name(self) -> str:
        return self.source_node['hostname']  # type: ignore # pylint: disable=no-member

    @staticmethod
    def get_file_path(spec: str, bucket_name: str, file: str) -> str:
        bucket_name = urllib.parse.quote_plus(bucket_name)
        bucket_path = os.path.join(os.path.normpath(spec), "bucket-" + bucket_name)
        if os.path.isdir(bucket_path):
            return os.path.join(bucket_path, file)

        path, _ = BFD.find_latest_dir(spec, None)
        if path:
            path, _ = BFD.find_latest_dir(path, None)
            if path:
                return os.path.join(path, "bucket-" + bucket_name, file)
        return os.path.join(bucket_path, file)

    @staticmethod
    def construct_dir(parent: str, bucket_name: str, node_name: str) -> str:
        return os.path.join(parent,
                            "bucket-" + urllib.parse.quote_plus(bucket_name),
                            "node-" + urllib.parse.quote_plus(node_name))

    @staticmethod
    def check_full_dbfiles(parent_dir: str) -> List[str]:
        return glob.glob(os.path.join(parent_dir, "data-*.cbb"))

    @staticmethod
    def get_predecessors(parent_dir: str):
        try:
            json_file = open(os.path.join(parent_dir, "meta.json"), "r")
            json_data = json.load(json_file)
            json_file.close()
            return json_data["pred"]
        except IOError:
            return []

    @staticmethod
    def get_server_version(parent_dir: str) -> str:
        try:
            json_file = open(os.path.join(parent_dir, "meta.json"), "r")
            json_data = json.load(json_file)
            json_file.close()
            if "version" in json_data:
                return json_data["version"]
            return "0.0.0"
        except IOError:
            return "0.0.0"

    @staticmethod
    def get_failover_log(parent_dir: str) -> Dict[str, Any]:
        try:
            filepath = os.path.join(parent_dir, "failover.json")
            if os.path.isfile(filepath):
                json_file = open(filepath, "r")
                json_data = json.load(json_file)
                json_file.close()
                return json_data
            else:
                return {}
        except IOError:
            return {}

    @staticmethod
    def list_files(opts, spec, bucket, node, pattern):
        pass

    @staticmethod
    def write_json_file(parent_dir, filename, output_data):
        filepath = os.path.join(parent_dir, filename)
        json_data = {}
        if os.path.isfile(filepath):
            # load into existed meta data generated from previous batch run
            try:
                json_file = open(filepath, "r")
                json_data = json.load(json_file)
                json_file.close()
            except IOError:
                pass

        for i in range(BFD.NUM_VBUCKET):
            if output_data.get(i):
                str_index = str(i)
                historic_data_exist = json_data.get(str_index)
                if historic_data_exist:
                    # Historic data will share same data type as incoming ones.
                    # It will be sufficient to only check type for historic data
                    if isinstance(json_data[str_index], int):
                        # For seqno, we want to keep the highest seqno
                        if json_data[str_index] < output_data[i]:
                            json_data[str_index] = output_data[i]
                    elif isinstance(json_data[str_index], list):
                        # For each vbucket, we want to get the superset of its references
                        if len(json_data[str_index]) <= len(output_data[i]):
                            json_data[str_index] = output_data[i]
                else:
                    # Bookkeeping the incoming one.
                    json_data[str_index] = output_data[i]

        json_file = open(filepath, "w")
        json.dump(json_data, json_file, ensure_ascii=False)
        json_file.close()

    @staticmethod
    def db_dir(spec: str, bucket_name: str, node_name: str, tmstamp: Optional[str] = None, mode: Optional[str] = None,
               new_session: bool = False) -> str:
        parent_dir = f'{os.path.normpath(spec)}/bucket-{urllib.parse.quote_plus(bucket_name)}/node-' \
            f'{urllib.parse.quote_plus(node_name)}'
        if os.path.isdir(parent_dir):
            return parent_dir

        # check 3.0 directory structure
        if not tmstamp:
            tmstamp = time.strftime("%Y-%m-%dT%H%M%SZ", time.gmtime())
        parent_dir = os.path.normpath(spec)
        rootpath, _ = BFD.find_latest_dir(parent_dir, None)
        if not rootpath or not mode or mode == "full":
            # no any backup roots exists
            path = os.path.join(parent_dir, tmstamp, tmstamp + "-full")
            return BFD.construct_dir(path, bucket_name, node_name)

        # check if any full backup exists
        path, _ = BFD.find_latest_dir(rootpath, "full")  # type: ignore
        if not path:
            path = os.path.join(rootpath, tmstamp + "-full")
            return BFD.construct_dir(path, bucket_name, node_name)

        # further check full backup for this bucket and node
        path = BFD.construct_dir(path, bucket_name, node_name)
        if not os.path.isdir(path):
            return path

        if mode.find("diff") >= 0:
            path, _ = BFD.find_latest_dir(rootpath, "diff")  # type: ignore
            if not path or new_session:
                path = os.path.join(rootpath, tmstamp + "-diff")
                return BFD.construct_dir(path, bucket_name, node_name)

            path = BFD.construct_dir(path, bucket_name, node_name)
            if not os.path.isdir(path):
                return path

            path = os.path.join(rootpath, tmstamp + "-diff")
            return BFD.construct_dir(path, bucket_name, node_name)

        if mode.find("accu") >= 0:
            path, _ = BFD.find_latest_dir(rootpath, "accu")  # type: ignore
            if not path or new_session:
                path = os.path.join(rootpath, tmstamp + "-accu")
                return BFD.construct_dir(path, bucket_name, node_name)

            path = BFD.construct_dir(path, bucket_name, node_name)
            if not os.path.isdir(path):
                return path

            path = os.path.join(rootpath, tmstamp + "-accu")
            return BFD.construct_dir(path, bucket_name, node_name)

        return parent_dir

    @staticmethod
    def find_latest_dir(parent_dir: str, mode: Union[str, None]) -> Tuple[Optional[str], List[str]]:
        all_subdirs: List[str] = []
        latest_dir: Optional[str] = None
        if not os.path.isdir(parent_dir):
            return latest_dir, all_subdirs
        for d in os.listdir(parent_dir):
            if not mode or d.find(mode) >= 0:
                bd = os.path.join(parent_dir, d)
                if os.path.isdir(bd):
                    all_subdirs.append(bd)
        if all_subdirs:
            all_subdirs = sorted(all_subdirs, key=os.path.getmtime, reverse=True)
            latest_dir = all_subdirs[0]
        return latest_dir, all_subdirs

    @staticmethod
    def find_seqno(opts: Any, spec: str, bucket_name: str, node_name: str, mode: str) -> Tuple[Dict[Any, Any],
                                                                                               List[Any],
                                                                                               Dict[Any, Any],
                                                                                               Dict[Any, Any]]:
        seqno: Dict[Any, Any] = {}
        dep: Dict[Any, Any] = {}
        dep_list: List[Any] = []
        failover_log: Dict[Any, Any] = {}
        snapshot_markers: Dict[Any, Any] = {}
        for i in range(BFD.NUM_VBUCKET):
            seqno[str(i)] = 0
            dep[i] = None
            failover_log[i] = None
            snapshot_markers[i] = None

        file_list: List[Any] = []
        failoverlog_list: List[Any] = []
        snapshot_list: List[Any] = []
        seqno_list: List[Any] = []
        parent_dir: str = os.path.normpath(spec)

        if mode == "full":
            return seqno, dep_list, failover_log, snapshot_markers
        timedir, latest_dirs = BFD.find_latest_dir(parent_dir, None)
        if not timedir:
            return seqno, dep_list, failover_log, snapshot_markers
        fulldir, latest_dirs = BFD.find_latest_dir(timedir, "full")
        if not fulldir:
            return seqno, dep_list, failover_log, snapshot_markers

        path = BFD.construct_dir(fulldir, bucket_name, node_name)
        if not os.path.isdir(path):
            return seqno, dep_list, failover_log, snapshot_markers

        last_backup = BFD.construct_dir(fulldir, bucket_name, node_name)
        file_list.extend(recursive_glob(path, 'data-*.cbb'))
        failoverlog_list.extend(recursive_glob(path, 'failover.json'))
        snapshot_list.extend(recursive_glob(path, 'snapshot_markers.json'))
        seqno_list.extend(recursive_glob(path, 'seqno.json'))

        if mode.find("accu") < 0:
            accudir, accu_dirs = BFD.find_latest_dir(timedir, "accu")
            if accudir:
                path = BFD.construct_dir(accudir, bucket_name, node_name)
                last_backup = path
                if os.path.isdir(path):
                    file_list.extend(recursive_glob(path, 'data-*.cbb'))
                    failoverlog_list.extend(recursive_glob(path, 'failover.json'))
                    snapshot_list.extend(recursive_glob(path, 'snapshot_markers.json'))
                    seqno_list.extend(recursive_glob(path, 'seqno.json'))
        if mode.find("diff") >= 0:
            diffdir, diff_dirs = BFD.find_latest_dir(timedir, "diff")
            if diff_dirs:
                last_backup = BFD.construct_dir(diffdir, bucket_name, node_name)  # type: ignore
                for dir in diff_dirs:
                    path = BFD.construct_dir(dir, bucket_name, node_name)
                    if os.path.isdir(path):
                        file_list.extend(recursive_glob(path, 'data-*.cbb'))
                        failoverlog_list.extend(recursive_glob(path, 'failover.json'))
                        snapshot_list.extend(recursive_glob(path, 'snapshot_markers.json'))
                        seqno_list.extend(recursive_glob(path, 'seqno.json'))

        if BFD.get_server_version(last_backup) < "5.0.0":
            return dict(), list(), dict(), dict()

        for x in sorted(seqno_list):
            try:
                json_file = open(x, "r")
                json_data = json.load(json_file)
                json_file.close()

                for vbid, seq in json_data.items():
                    if not seq:
                        continue
                    if seqno.get(vbid) < seq:
                        seqno[vbid] = seq
            except IOError:
                pass

        for log_file in sorted(failoverlog_list):
            try:
                json_file = open(log_file, "r")
                json_data = json.load(json_file)
                json_file.close()

                for vbid, flogs in json_data.items():
                    if not flogs:
                        continue
                    elif vbid not in failover_log:
                        failover_log[vbid] = flogs
                    else:
                        for logpair in flogs:
                            if not failover_log[vbid]:
                                failover_log[vbid] = [logpair]
                            elif logpair not in failover_log[vbid]:
                                failover_log[vbid].append(logpair)
            except IOError:
                pass

        for snapshot in sorted(snapshot_list):
            try:
                json_file = open(snapshot, "r")
                json_data = json.load(json_file)
                json_file.close()

                for vbid, markers in json_data.items():
                    snapshot_markers[vbid] = markers
            except IOError:
                pass

        for i in range(BFD.NUM_VBUCKET):
            if dep[i] and dep[i] not in dep_list:
                dep_list.append(dep[i])
        return seqno, dep_list, failover_log, snapshot_markers CBB_VERSION = [2004, 2014, 2015]  # sqlite pragma user version. DDOC_FILE_NAME = "design.json" FTS_FILE_NAME = "fts_index.json" INDEX_FILE_NAME = "index.json" class BFDSource(BFD, pump.Source):
    """Can read from backup-file/directory layout."""

    def __init__(self, opts, spec, source_bucket, source_node,
                 source_map, sink_map, ctl, cur):
        super(BFDSource, self).__init__(opts, spec, source_bucket, source_node,
                                        source_map, sink_map, ctl, cur)
        self.done = False
        self.files = None
        self.cursor_db = None

    @staticmethod
    def can_handle(opts, spec: str) -> Union[bool, List[str]]:
        if os.path.isdir(spec):
            dirs = glob.glob(f'{spec}/bucket-*/node-*/data-*.cbb')
            if dirs:
                return True
            path, dirs = BFD.find_latest_dir(spec, None)
            if path:
                path, dirs = BFD.find_latest_dir(path, "full")
                if path:
                    return glob.glob(f'{path}*/bucket-*/node-*/data-*.cbb')
        return False

    @staticmethod
    def check(opts, spec: str) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[Dict[str, Any]]]:
        spec = os.path.normpath(spec)
        if not os.path.isdir(spec):
            return f'error: backup_dir is not a directory: {spec}', None

        buckets = []

        bucket_dirs = glob.glob(f'{spec}/bucket-*')
        if not bucket_dirs:
            # check 3.0 directory structure
            path, dirs = BFD.find_latest_dir(spec, None)
            if not path:
                return f'error: no backup directory found: {spec}', None
            latest_dir, dir = BFD.find_latest_dir(path, "full")
            if not latest_dir:
                return f'error: no valid backup directory found: {path}', None
            bucket_dirs = glob.glob(f'{latest_dir}/bucket-*')

        for bucket_dir in sorted(bucket_dirs):
            if not os.path.isdir(bucket_dir):
                return f'error: not a bucket directory: {bucket_dir}', None

            bucket_name = urllib.parse.unquote_plus(os.path.basename(bucket_dir)[len("bucket-"):].strip()).encode()
            if not bucket_name:
                return f'error: bucket_name too short: {bucket_dir}', None

            bucket = {'name': bucket_name, 'nodes': []}
            buckets.append(bucket)

            node_dirs = glob.glob(f'{bucket_dir}/node-*')
            for node_dir in sorted(node_dirs):
                if not os.path.isdir(node_dir):
                    return f'error: not a node directory: {node_dir}', None

                node_name = urllib.parse.unquote_plus(os.path.basename(node_dir)[len("node-"):].strip()).encode()
                if not node_name:
                    return f'error: node_name too short: {node_dir}', None

                bucket['nodes'].append({'hostname': node_name})  # type: ignore

        return 0, {'spec': spec, 'buckets': buckets}

    @staticmethod
    def provide_design(opts, source_spec: str, source_bucket: Dict[str, Any],
                       source_map: Any) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[str]]:
        return BFDSource.read_index_file(source_spec, source_bucket, DDOC_FILE_NAME)

    @staticmethod
    def provide_index(opts, source_spec: str, source_bucket: Dict[str, Any],
                      source_map: Any) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[str]]:
        return BFDSource.read_index_file(source_spec, source_bucket, INDEX_FILE_NAME)

    @staticmethod
    def provide_fts_index(opts, source_spec: str, source_bucket: Dict[str, Any],
                          source_map: Any) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[str]]:
        return BFDSource.read_index_file(source_spec, source_bucket, FTS_FILE_NAME)

    @staticmethod
    def provide_fts_alias(opts, source_spec: str, source_bucket: Dict[str, Any],
                          source_map: Any) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[str]]:
        base_dir = source_spec
        path, dirs = BFD.find_latest_dir(source_spec, None)
        if path:
            path, dirs = BFD.find_latest_dir(path, None)
            if path:
                base_dir = path

        if os.path.exists(os.path.join(base_dir, FTS_ALIAS_FILE_NAME)):
            try:
                f = open(os.path.join(base_dir, FTS_ALIAS_FILE_NAME), 'r')
                d = f.read()
                f.close()
                return 0, d
            except IOError as e:
                return f'error: could not write {path}; exception: {e}', None

        return 0, None

    @staticmethod
    def read_index_file(source_spec: str, source_bucket: Dict[str, Any],
                        file_name: str) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[str]]:
        path = BFD.get_file_path(source_spec, source_bucket['name'], file_name)
        if os.path.exists(path):
            try:
                f = open(path, 'r')
                d = f.read()
                f.close()
                return 0, d
            except IOError as e:
                return f'error: could not read {path}; exception: {e}', None
        return 0, None

    def get_conflict_resolution_type(self) -> str:
        dir = BFD.construct_dir(self.source_map['spec'],
                                self.source_bucket['name'],
                                self.source_node['hostname'])
        metafile = os.path.join(dir, "meta.json")
        try:
            json_file = open(metafile, "r")
            json_data = json.load(json_file)
            json_file.close()
            if "conflict_resolution_type" in json_data:
                return json_data["conflict_resolution_type"]
            return "seqno"
        except IOError:
            return "seqno"

    def provide_batch(self) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[pump.Batch]]:
        if self.done:
            return 0, None

        batch = pump.Batch(self)

        batch_max_size = self.opts.extra['batch_max_size']
        batch_max_bytes = self.opts.extra['batch_max_bytes']

        s = ["SELECT cmd, vbucket_id, key, flg, exp, cas, meta, val FROM cbb_msg",
             "SELECT cmd, vbucket_id, key, flg, exp, cas, meta, val, seqno, dtype, meta_size FROM cbb_msg",
             "SELECT cmd, vbucket_id, key, flg, exp, cas, meta, val, seqno, dtype, meta_size, conf_res FROM cbb_msg"]

        if self.files is None:  # None != [], as self.files will shrink to [].
            g = glob.glob(BFD.db_dir(self.spec, self.bucket_name(), self.node_name()) + "/data-*.cbb")
            if not g:
                # check 3.0 file structure
                rv, file_list = BFDSource.list_files(self.opts,
                                                     self.spec,
                                                     self.bucket_name(),
                                                     self.node_name(),
                                                     "data-*.cbb")
                if rv != 0:
                    return rv, None
                from_date = getattr(self.opts, "from_date", None)
                if from_date:
                    from_date = datetime.datetime.strptime(from_date, "%Y-%m-%d")

                to_date = getattr(self.opts, "to_date", None)
                if to_date:
                    to_date = datetime.datetime.strptime(to_date, "%Y-%m-%d")
                g = []
                for f in file_list:
                    mtime = datetime.datetime.fromtimestamp(os.path.getmtime(f))
                    if (not from_date or mtime >= from_date) and (not to_date or mtime <= to_date):
                        g.append(f)
            self.files = sorted(g)
        try:
            ver = 0
            while (not self.done and
                   batch.size() < batch_max_size and
                   batch.bytes < batch_max_bytes):
                if self.cursor_db is None:
                    if not self.files:
                        self.done = True
                        return 0, batch

                    rv, db, ver = connect_db(self.files[0], self.opts, CBB_VERSION)
                    if rv != 0:
                        return rv, None
                    self.files = self.files[1:]

                    cursor = db.cursor()
                    cursor.execute(s[ver])

                    self.cursor_db = (cursor, db, ver)

                cursor, db, ver = self.cursor_db

                row = cursor.fetchone()
                if row:
                    vbucket_id = row[1]
                    key = row[2]
                    val = row[7]

                    if self.skip(key, vbucket_id):
                        continue

                    if ver == 2:
                        msg = (row[0], row[1], row[2], row[3], row[4],
                               int(row[5]),  # CAS as 64-bit integer not string.
                               row[6],  # revid as 64-bit integer too
                               row[7], row[8], row[9], row[10], row[11])
                    elif ver == 1:
                        msg = (row[0], row[1], row[2], row[3], row[4],
                               int(row[5]),  # CAS as 64-bit integer not string.
                               row[6],  # revid as 64-bit integer too
                               row[7], row[8], row[9], row[10], 0)
                    else:
                        msg = (row[0], row[1], row[2], row[3], row[4],
                               int(row[5]),  # CAS as 64-bit integer not string.
                               row[6],  # revid as 64-bit integer too
                               row[7], 0, 0, 0, 0)

                    batch.append(msg, len(val))
                else:
                    if self.cursor_db:
                        self.cursor_db[0].close()
                        self.cursor_db[1].close()
                    self.cursor_db = None

            return 0, batch

        except Exception as e:
            self.done = True
            if self.cursor_db:
                self.cursor_db[0].close()
                self.cursor_db[1].close()
            self.cursor_db = None

            return f'error: exception reading backup file: {e!s}', None

    @staticmethod
    def total_msgs(opts, source_bucket, source_node, source_map) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[int]]:
        t = 0
        file_list = glob.glob(BFD.db_dir(
            source_map['spec'],
            source_bucket['name'],
            source_node['hostname']) + "/data-*.cbb")
        if not file_list:
            # check 3.0 directory structure
            rv, file_list = BFDSource.list_files(opts,
                                                 source_map['spec'],
                                                 source_bucket['name'],
                                                 source_node['hostname'],
                                                 "data-*.cbb")
            if rv != 0:
                return rv, None

        for x in sorted(file_list):
            rv, db, ver = connect_db(x, opts, CBB_VERSION)
            if rv != 0:
                return rv, None

            cur = db.cursor()
            # TODO: (1) BFDSource - COUNT(*) is not indexed.
            cur.execute("SELECT COUNT(*) FROM cbb_msg;")
            t = t + cur.fetchone()[0]
            cur.close()
            db.close()

        return 0, t

    @staticmethod
    def list_files(opts, spec: str, bucket: str, node: str, pattern: str) -> Tuple[couchbaseConstants.PUMP_ERROR,
                                                                                   List[str]]:
        file_list: List[str] = []
        prec_list: List[str] = []
        path, dirs = BFD.find_latest_dir(spec, None)
        if not path:
            return f'error: No valid data in path: {spec}', file_list

        path, dirs = BFD.find_latest_dir(path, None)
        if not path:
            return 0, file_list

        for dir in dirs:
            latest_dir = BFD.construct_dir(dir, bucket, node)
            file_list.extend(glob.glob(os.path.join(latest_dir, pattern)))
            for p in BFD.get_predecessors(latest_dir):
                prec_list.append(os.path.dirname(p))
            if len(prec_list) > 0:
                break

        while len(prec_list) > 0:
            for d in glob.glob(os.path.join(prec_list[0], pattern)):
                if d not in file_list:
                    file_list.append(d)
            for p in BFD.get_predecessors(prec_list[0]):
                dirname = os.path.dirname(p)
                if dirname not in prec_list:
                    prec_list.append(dirname)
            prec_list = prec_list[1:]

        return 0, file_list class BFDSinkEx(pump_bfd.BFDSink):

    def __init__(self, opts, spec, source_bucket, source_node,
                 source_map, sink_map, ctl, cur):
        # pylint: disable=bad-super-call
        super(pump_bfd.BFDSink, self).__init__(opts, spec, source_bucket, source_node, source_map,
                                               sink_map, ctl, cur)
        self.mode = getattr(opts, "mode", "diff")
        self.init_worker(pump_bfd.BFDSink.run)

    @staticmethod
    def check_spec(source_bucket, source_node, opts, spec, cur):
        pump.Sink.check_spec(source_bucket, source_node, opts, spec, cur)

        seqno, dep, faillover_log, snapshot_markers = pump_bfd.BFD.find_seqno(opts, spec,
                                                                              source_bucket['name'],
                                                                              source_node['hostname'],
                                                                              getattr(opts, "mode", "diff"))
        if 'seqno' in cur:
            cur['seqno'][(source_bucket['name'], source_node['hostname'])] = seqno
        else:
            cur['seqno'] = {(source_bucket['name'], source_node['hostname']): seqno}

        if 'failoverlog' in cur:
            cur['failoverlog'][(source_bucket['name'], source_node['hostname'])] = faillover_log
        else:
            cur['failoverlog'] = {(source_bucket['name'], source_node['hostname']): faillover_log}

        if 'snapshot' in cur:
            cur['snapshot'][(source_bucket['name'], source_node['hostname'])] = snapshot_markers
        else:
            cur['snapshot'] = {(source_bucket['name'], source_node['hostname']): snapshot_markers} class CSVSink(pump.Sink):
    """Emits batches to stdout in CSV format."""
    CSV_SCHEME = "csv:"
    CSV_JSON_SCHEME = "csv-json:"

    def __init__(self, opts, spec, source_bucket, source_node,
                 source_map, sink_map, ctl, cur):
        super(CSVSink, self).__init__(opts, spec, source_bucket, source_node,
                                      source_map, sink_map, ctl, cur)
        self.csvfile = None
        self.writer = None
        self.fields = None

    def bucket_name(self) -> str:
        if 'name' in self.source_bucket:
            return self.source_bucket['name']
        return ""

    def node_name(self) -> str:
        if 'hostname' in self.source_node:
            return self.source_node['hostname']
        return ""

    @staticmethod
    def can_handle(opts, spec: str) -> bool:
        if spec.startswith(CSVSink.CSV_SCHEME) or spec.startswith(CSVSink.CSV_JSON_SCHEME):
            opts.threads = 1  # Force 1 thread to not overlap stdout.
            return True
        return False

    @staticmethod
    def check(opts, spec, source_map) -> Tuple[couchbaseConstants.PUMP_ERROR, Any]:
        rv: couchbaseConstants.PUMP_ERROR = 0
        if spec.endswith(".csv"):
            if spec.startswith(CSVSink.CSV_JSON_SCHEME):
                targetpath = spec[len(CSVSink.CSV_JSON_SCHEME):]
            else:
                targetpath = spec[len(CSVSink.CSV_SCHEME):]
            targetpath = os.path.normpath(targetpath)
            rv = pump.mkdirs(targetpath)

        return rv, None

    @staticmethod
    def consume_design(opts, sink_spec, sink_map,
                       source_bucket, source_map, source_design) -> couchbaseConstants.PUMP_ERROR:
        if source_design:
            logging.warn("warning: cannot save bucket design"
                         " on a CSV destination")
        return 0

    def consume_batch_async(self, batch: pump.Batch) -> Tuple[couchbaseConstants.PUMP_ERROR,
                                                              Optional[pump.SinkBatchFuture]]:
        if not self.writer:
            self.csvfile = sys.stdout
            if self.spec.startswith(CSVSink.CSV_JSON_SCHEME):
                if len(batch.msgs) <= 0:
                    future = pump.SinkBatchFuture(self, batch)
                    self.future_done(future, 0)
                    return 0, future

                cmd, vbucket_id, key, flg, exp, cas, meta, val_bytes = batch.msgs[0][:8]
                doc = json.loads(val_bytes)
                self.fields = sorted(doc.keys())
                if 'id' not in self.fields:
                    self.fields = ['id'] + self.fields
                if self.spec.endswith(".csv"):
                    filename = self.get_csvfile(self.spec[len(CSVSink.CSV_JSON_SCHEME):])
                    try:
                        self.csvfile = open(filename, "w", encoding='utf-8')
                    except IOError:
                        return f'error: could not write csv to file: {filename}', None
                self.writer = csv.writer(self.csvfile)
                self.writer.writerow(self.fields)
            else:
                if self.spec.endswith(".csv"):
                    filename = self.get_csvfile(self.spec[len(CSVSink.CSV_SCHEME):])
                    try:
                        self.csvfile = open(filename, "w", encoding='utf-8')
                    except IOError:
                        return f'error: could not write csv to file: {filename}', None
                self.writer = csv.writer(self.csvfile)
                self.writer.writerow(['id', 'flags', 'expiration', 'cas', 'value', 'rev', 'vbid', 'dtype'])
        msg_tuple_format = 0
        for msg in batch.msgs:
            cmd, vbucket_id, key, flg, exp, cas, meta, val_bytes = msg[:8]
            if self.skip(key, vbucket_id):
                continue
            if not msg_tuple_format:
                msg_tuple_format = len(msg)
            seqno = dtype = nmeta = 0
            if msg_tuple_format > 8:
                seqno, dtype, nmeta, conf_res = msg[8:12]
            if dtype & couchbaseConstants.DATATYPE_COMPRESSED != 0:
                try:
                    val_bytes = snappy.uncompress(val_bytes)
                except Exception:
                    pass
            try:
                if cmd in [couchbaseConstants.CMD_TAP_MUTATION,
                           couchbaseConstants.CMD_DCP_MUTATION]:
                    if self.fields:
                        if val_bytes and len(val_bytes) > 0:
                            try:
                                row = []
                                doc = json.loads(val_bytes)
                                if isinstance(doc, dict):
                                    for field in self.fields:
                                        if field == 'id':
                                            row.append(pump.return_string(key))
                                        else:
                                            row.append(doc[field])
                                    self.writer.writerow(row)
                            except ValueError:
                                pass
                    else:
                        self.writer.writerow([pump.return_string(key), flg, exp, cas,
                                             val_bytes, meta, vbucket_id, dtype])
                elif cmd in [couchbaseConstants.CMD_TAP_DELETE, couchbaseConstants.CMD_DCP_DELETE]:
                    pass
                elif cmd == couchbaseConstants.CMD_GET:
                    pass
                else:
                    return f'error: CSVSink - unknown cmd: {cmd!s}', None
            except IOError:
                return "error: could not write csv to stdout", None

        future = pump.SinkBatchFuture(self, batch)
        self.future_done(future, 0)
        return 0, future

    def close(self):
        if self.csvfile is not None and self.csvfile != sys.stdout:
            self.csvfile.close()
            self.csvfile = None

    def get_csvfile(self, base: str) -> str:
        extension = os.path.splitext(base)
        filename = extension[0]
        if self.bucket_name():
            filename = f'{filename}_{urllib.parse.quote_plus(self.bucket_name())}'
        if self.node_name():
            filename = f'{filename}_{urllib.parse.quote_plus(self.node_name())}'
        return filename + extension[1]

    def convert_meta(self, meta) -> int:
        seq_no = str(meta)
        if len(seq_no) > 8:
            seq_no = seq_no[0:8]
        if len(seq_no) < 8:
            seq_no = ('\x00\x00\x00\x00\x00\x00\x00\x00' + seq_no)[-8:]
        check_seqno, = struct.unpack(">Q", seq_no.encode())
        if not check_seqno:
            check_seqno = 1

        return check_seqno class CSVSource(pump.Source):
    """Reads csv file, where first line is field names and one field
       should be 'id'."""

    def __init__(self, opts, spec, source_bucket, source_node,
                 source_map, sink_map, ctl, cur):
        super(CSVSource, self).__init__(opts, spec, source_bucket, source_node,
                                        source_map, sink_map, ctl, cur)
        self.done = False
        self.file = None
        self.r = None  # An iterator of csv.reader()

    @staticmethod
    def can_handle(opts, spec: str) -> bool:
        return spec.endswith(".csv") and os.path.isfile(spec)

    @staticmethod
    def check(opts, spec: str) -> Tuple[couchbaseConstants.PUMP_ERROR, Dict[str, Any]]:
        return 0, {'spec': spec,
                   'buckets': [{'name': os.path.basename(spec),
                                'nodes': [{'hostname': 'N/A'}]}]}

    @staticmethod
    def provide_design(opts, source_spec, source_bucket, source_map) -> Tuple[couchbaseConstants.PUMP_ERROR,
                                                                              Optional[str]]:
        return 0, None

    def close_file(self):
        self.file.close()
        self.file = None

    def provide_batch(self) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[pump.Batch]]:
        if self.done:
            return 0, None

        if not self.r:
            try:
                self.file = open(self.spec, 'r', encoding='utf-8')
                self.r = csv.reader(self.file)
                self.fields = next(self.r)
                if 'id' not in self.fields:
                    self.close_file()
                    return f'error: no \'id\' field in 1st line of csv: {self.spec}', None
            except StopIteration:
                self.close_file()
                return f'error: could not read 1st line of csv: {self.spec}', None
            except IOError as e:
                return f'error: could not open csv: {self.spec}; exception: {e!s}', None

        batch = pump.Batch(self)

        batch_max_size = self.opts.extra['batch_max_size']
        batch_max_bytes = self.opts.extra['batch_max_bytes']

        cmd = couchbaseConstants.CMD_TAP_MUTATION
        vbucket_id = 0x0000ffff

        while (self.r
                and batch.size() < batch_max_size
                and batch.bytes < batch_max_bytes):
            try:
                vals = next(self.r)
                doc = {}
                for i, field in enumerate(self.fields):
                    if i >= len(vals):
                        continue
                    if field == 'id':
                        doc[field] = vals[i]
                    else:
                        doc[field] = number_try_parse(vals[i])
                if doc['id']:
                    msg: couchbaseConstants.BATCH_MSG = (cmd, vbucket_id, doc['id'].encode(), 0, 0, 0, b'',
                                                         literal_eval(doc['value']), 0, 0, 0, 0)
                    batch.append(msg, len(doc))
            except StopIteration:
                self.done = True
                self.close_file()
                self.r = None
            except Exception as e:
                logging.error(f'error: fails to read from csv file {e}')
                continue

        if batch.size() <= 0:
            return 0, None
        return 0, batch class DCPStreamSource(pump.Source, threading.Thread):
    """Can read from cluster/server/bucket via DCP streaming."""
    HIGH_SEQNO = "high_seqno"
    VB_UUID = "uuid"
    ABS_HIGH_SEQNO = "abs_high_seqno"
    PURGE_SEQNO = "purge_seqno"

    HIGH_SEQNO_BYTE = 8
    UUID_BYTE = 8

    def __init__(self, opts, spec: str, source_bucket, source_node: Dict[str, Any],
                 source_map, sink_map, ctl: Dict[str, Any], cur: Dict[str, Any]):
        if spec.startswith("https://"):
            setattr(opts, "ssl", True)
        super(DCPStreamSource, self).__init__(opts, spec, source_bucket, source_node,
                                              source_map, sink_map, ctl, cur)
        threading.Thread.__init__(self)

        self.dcp_done: bool = False
        self.dcp_conn: Optional[cb_bin_client.MemcachedClient] = None
        self.mem_conn: Optional[cb_bin_client.MemcachedClient] = None
        self.dcp_name: str = opts.process_name
        self.ack_last: bool = False
        self.cmd_last: Optional[int] = None
        self.num_msg: int = 0
        # change made to support development version
        self.version_supported: bool = self.source_node['version'].split(".") >= ["3", "0", "0"] or \
            (self.source_node['version'].split(".") == ['0', '0', '0-0000-enterprise'])
        self.recv_min_bytes: int = int(opts.extra.get("recv_min_bytes", 4096))
        self.batch_max_bytes: int = int(opts.extra.get("batch_max_bytes", 400000))
        self.flow_control = int(opts.extra.get("flow_control", 1))
        self.vbucket_list = getattr(opts, "vbucket_list", None)
        self.r = random.Random()
        self.queue_size = int(opts.extra.get("dcp_consumer_queue_length", 1000))
        self.response: queue.Queue = queue.Queue(self.queue_size)
        self.running = False
        self.stream_list: Dict[Any, Any] = {}
        self.unack_size = 0
        self.node_vbucket_map: Optional[List[int]] = None
        self.uncompress = opts.extra.get("uncompress", 0)

    @staticmethod
    def can_handle(opts, spec: str) -> bool:
        return (spec.startswith("http://") or
                spec.startswith("couchbase://") or
                spec.startswith("https://")
                )

    @staticmethod
    def check(opts, spec: str) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[Dict[str, Any]]]:
        err, map = pump.rest_couchbase(opts, spec)
        if err:
            return err, map
        if not map or not map.get('buckets'):
            return f'error: no buckets supporting DCP at source: {spec};' \
                f' please check your username/password to the cluster', None
        return 0, map

    @staticmethod
    def provide_design(opts, source_spec: str, source_bucket: Dict[str, Any], source_map) -> \
            Tuple[couchbaseConstants.PUMP_ERROR, Optional[str]]:
        # Ephemeral buckets do not have design docs
        if source_bucket['bucketType'] == 'ephemeral':
            return 0, None
        spec_parts = source_map.get('spec_parts')
        if not spec_parts:
            return "error: no design spec_parts", None
        host, port, user, pswd, path = spec_parts

        source_nodes = pump.filter_bucket_nodes(source_bucket, spec_parts)
        if not source_nodes:
            source_nodes = source_bucket['nodes']
            if not source_nodes:
                return f'error: no design source node; spec_parts: {spec_parts}', None

        couch_api_base = source_nodes[0].get('couchApiBase')
        if not couch_api_base:
            return 0, None  # No couchApiBase; probably not 2.0.

        err, _, ddocs = \
            pump.rest_request_json(host, int(port), user, pswd, opts.ssl,
                                   "/pools/default/buckets/%s/ddocs" %
                                   (source_bucket['name']),
                                   reason="provide_design", verify=not opts.no_ssl_verify, cacert=opts.cacert)
        if err and "response: 404" in err:  # A 404/not-found likely means 2.0-DP4.
            ddocs_url = f'{couch_api_base}/_all_docs'
            ddocs_qry = "?startkey=\"_design/\"&endkey=\"_design0\"&include_docs=true"

            host, port, user, pswd, path = \
                pump.parse_spec(opts, ddocs_url, 8092)
            # Not using user/pwd as 2.0-DP4 CAPI did not support auth.
            err, _, ddocs = pump.rest_request_json(host, int(port),
                                                   None, None, opts.ssl, path + ddocs_qry,
                                                   reason="provide_design-2.0DP4", verify=not opts.no_ssl_verify,
                                                   cacert=opts.cacert)
        if err is not None:
            return err, None

        if ddocs is None:
            return 0, None

        if not ddocs.get('rows', None):
            return 0, None

        return 0, json.dumps(ddocs.get('rows', []))

    @staticmethod
    def provide_index(opts, source_spec: str, source_bucket: Dict[str, Any], source_map) -> \
            Tuple[couchbaseConstants.PUMP_ERROR, Optional[str]]:
        try:
            rest = ClusterManager(source_spec, opts.username, opts.password, opts.ssl, not opts.no_ssl_verify,
                                  opts.cacert, False)
            result, errors = rest.get_index_metadata(source_bucket['name'])
            if errors:
                return errors, None
            return 0, json.dumps(result["result"])
        except ServiceNotAvailableException:
            return 0, None

    @staticmethod
    def provide_fts_index(opts, source_spec: str, source_bucket: Dict[str, Any], source_map) -> \
            Tuple[couchbaseConstants.PUMP_ERROR, Optional[str]]:
        try:
            rest = ClusterManager(source_spec, opts.username, opts.password, opts.ssl, not opts.no_ssl_verify,
                                  opts.cacert, False)
            result, errors = rest.get_fts_index_metadata(source_bucket['name'])
            if errors:
                return errors, None
            return 0, json.dumps(result)
        except ServiceNotAvailableException:
            return 0, None

    @staticmethod
    def provide_fts_alias(opts, source_spec: str, source_bucket: Dict[str, Any], source_map) -> \
            Tuple[couchbaseConstants.PUMP_ERROR, Optional[str]]:
        try:
            rest = ClusterManager(source_spec, opts.username, opts.password, opts.ssl, not opts.no_ssl_verify,
                                  opts.cacert, False)
            result, errors = rest.get_fts_index_alias()
            if errors:
                return errors, None
            return 0, json.dumps(result)
        except ServiceNotAvailableException:
            return 0, None

    def get_conflict_resolution_type(self) -> str:
        conf_res_type = "seqno"
        if "conflictResolutionType" in self.source_bucket:
            conf_res_type = self.source_bucket["conflictResolutionType"]

        return conf_res_type

    def add_start_event(self, conn: Optional[cb_bin_client.MemcachedClient]) -> couchbaseConstants.PUMP_ERROR:
        sasl_user = str(self.source_bucket.get("name"))
        event = {"timestamp": self.get_timestamp(),
                 "real_userid": {"source": "internal",
                                 "user": pump.return_string(sasl_user),
                                 },
                 "mode": pump.return_string(getattr(self.opts, "mode", "diff")),
                 "source_bucket": pump.return_string(self.source_bucket['name']),
                 "source_node": pump.return_string(self.source_node['hostname'])
                 }
        if conn:
            try:
                conn.audit(couchbaseConstants.AUDIT_EVENT_BACKUP_START, json.dumps(event))
            except Exception as e:
                logging.warning(f'auditing error: {e}')
        return 0

    def add_stop_event(self, conn: Optional[cb_bin_client.MemcachedClient]) -> couchbaseConstants.PUMP_ERROR:
        sasl_user = str(self.source_bucket.get("name"))
        event = {"timestamp": self.get_timestamp(),
                 "real_userid": {"source": "internal",
                                 "user": pump.return_string(sasl_user)
                                 },
                 "source_bucket": pump.return_string(self.source_bucket['name']),
                 "source_node": pump.return_string(self.source_node['hostname'])
                 }
        if conn:
            try:
                conn.audit(couchbaseConstants.AUDIT_EVENT_BACKUP_STOP, json.dumps(event))
            except Exception as e:
                logging.warning(f'auditing error: {e}')
        return 0

    def build_node_vbucket_map(self) -> Optional[List[int]]:
        if "vBucketServerMap" in self.source_bucket:
            server_list = self.source_bucket["vBucketServerMap"]["serverList"]
            vbucket_map = self.source_bucket["vBucketServerMap"]["vBucketMap"]
        else:
            return None

        node_vbucket_map = []

        host, _ = couchbaseConstants.parse_host_port(self.source_node.get('hostname', 'N/A'))
        nodename = f'{host}:{self.source_node["ports"]["direct"]!s}'
        nodeindex = -1
        for index, node in enumerate(server_list):
            if nodename == node:
                nodeindex = index
                break
        for vbucket_id, vblist in enumerate(vbucket_map):
            if vblist[0] >= 0 and vblist[0] == nodeindex:
                node_vbucket_map.append(vbucket_id)
        return node_vbucket_map

    def provide_batch(self) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[pump.Batch]]:
        if not self.version_supported:
            return "error: cannot back up 2.x or older clusters with 3.x tools", None

        cur_sleep = 0.2
        cur_retry = 0
        max_retry = self.opts.extra['max_retry']
        if not self.node_vbucket_map:
            self.node_vbucket_map = self.build_node_vbucket_map()

        while True:
            if self.dcp_done:
                if self.dcp_conn:
                    self.add_stop_event(self.dcp_conn)
                    self.dcp_conn.close()
                    self.dcp_conn = None
                return 0, None

            rv = self.get_dcp_conn()
            if rv != 0:
                self.dcp_done = True
                return rv, None

            rv, batch = self.provide_dcp_batch_actual()
            if rv == 0:
                return 0, batch

            if self.dcp_conn:
                self.dcp_conn.close()
                self.dcp_conn = None

            if cur_retry > max_retry:
                self.dcp_done = True
                return rv, batch

            logging.warning(f'backoff: {cur_retry}, sleeping: {cur_sleep}, on error: {rv}')
            time.sleep(cur_sleep)
            cur_sleep = min(cur_sleep * 2, 20)  # Max backoff sleep 20 seconds.
            cur_retry = cur_retry + 1

    def maybe_uncompress_value(self, data_type: int, value: bytes) -> Tuple[int, bytes]:
        if self.uncompress and data_type & couchbaseConstants.DATATYPE_COMPRESSED:
            try:
                value = snappy.uncompress(value)
                data_type = data_type ^ couchbaseConstants.DATATYPE_COMPRESSED
            except snappy.UncompressError as e:
                logging.warning(f'Could not uncompress value: {e}')

        return data_type, value

    def provide_dcp_batch_actual(self) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[pump.Batch]]:
        batch = pump.Batch(self)

        batch_max_size = self.opts.extra['batch_max_size']
        batch_max_bytes = self.opts.extra['batch_max_bytes']
        delta_ack_size = batch_max_bytes * 10 / 4  # ack every 25% of buffer size
        last_processed = 0
        total_bytes_read = 0

        vbid: int = 0
        start_seqno: int = 0
        end_seqno: int = 0
        vb_uuid: int = 0
        ss_start_seqno: int = 0
        ss_end_seqno: int = 0
        no_response_count: int = 0
        try:
            while (not self.dcp_done and
                   batch.size() < batch_max_size and
                   batch.bytes < batch_max_bytes):

                if self.response.empty():
                    if len(self.stream_list) > 0:
                        logging.debug(f'no response while there {len(self.stream_list)} active streams')
                        time.sleep(.25)
                        no_response_count = no_response_count + 1
                        # if not had a response after a minimum of 30 seconds then state we are done
                        if no_response_count == 120:
                            logging.warning(f'no response for 30 seconds while there {len(self.stream_list)}'
                                            ' active streams'
                                            )
                            self.dcp_done = True
                    else:
                        self.dcp_done = True
                    continue

                no_response_count = 0
                unprocessed_size = total_bytes_read - last_processed
                if unprocessed_size > delta_ack_size:
                    rv = self.ack_buffer_size(unprocessed_size)
                    if rv:
                        logging.error(rv)
                    else:
                        last_processed = total_bytes_read

                cmd, errcode, opaque, cas, keylen, extlen, data, datalen, dtype, bytes_read = \
                    self.response.get()  # type: int, int, int, int, int, int, bytes, int, int, int
                total_bytes_read += bytes_read
                rv = 0
                metalen = flags = flg = exp = 0
                key = val = ext = b''
                need_ack = False
                seqno = 0
                if cmd == couchbaseConstants.CMD_DCP_REQUEST_STREAM:
                    total_bytes_read -= bytes_read
                    if errcode == couchbaseConstants.ERR_SUCCESS:
                        pair_index = (self.source_bucket['name'], self.source_node['hostname'])
                        start = 0
                        step = DCPStreamSource.HIGH_SEQNO_BYTE + DCPStreamSource.UUID_BYTE
                        while start + step <= datalen:
                            uuid, seqno = struct.unpack(
                                couchbaseConstants.DCP_VB_UUID_SEQNO_PKT_FMT,
                                data[start:start + step])
                            if pair_index not in self.cur['failoverlog']:
                                self.cur['failoverlog'][pair_index] = {}
                            if opaque not in self.cur['failoverlog'][pair_index] or \
                               not self.cur['failoverlog'][pair_index][opaque]:
                                self.cur['failoverlog'][pair_index][opaque] = [(uuid, seqno)]
                            else:
                                self.cur['failoverlog'][pair_index][opaque].append((uuid, seqno))
                            start = start + step
                    elif errcode == couchbaseConstants.ERR_KEY_ENOENT:
                        logging.warn("producer doesn't know about the vbucket uuid, rollback to 0")
                        vbid, flags, start_seqno, end_seqno, vb_uuid, ss_start_seqno, ss_end_seqno = \
                            self.stream_list[opaque]
                        del self.stream_list[opaque]
                    elif errcode == couchbaseConstants.ERR_KEY_EEXISTS:
                        logging.warning(f'a stream exists on the connection for vbucket: {opaque}')
                    elif errcode == couchbaseConstants.ERR_NOT_MY_VBUCKET:
                        logging.warning(f'Vbucket is not active anymore, skip it:{vbid!s}')
                        del self.stream_list[opaque]
                    elif errcode == couchbaseConstants.ERR_ERANGE:
                        logging.warning(f'Start or end sequence numbers specified incorrectly,({start_seqno},'
                                        f' {end_seqno})')
                        del self.stream_list[opaque]
                    elif errcode == couchbaseConstants.ERR_ROLLBACK:
                        vbid, flags, start_seqno, end_seqno, vb_uuid, ss_start_seqno, ss_stop_seqno = \
                            self.stream_list[opaque]
                        start_seqno, = struct.unpack(couchbaseConstants.DCP_VB_SEQNO_PKT_FMT, data)
                        # find the most latest uuid, hi_seqno that fit start_seqno
                        if self.cur['failoverlog']:
                            pair_index = (self.source_bucket['name'], self.source_node['hostname'])
                            if self.cur['failoverlog'][pair_index].get("vbid"):
                                for uuid, seqno in self.cur['failoverlog'][pair_index][vbid]:
                                    if start_seqno >= seqno:
                                        vb_uuid = uuid
                                        break
                        ss_start_seqno = start_seqno
                        ss_end_seqno = start_seqno
                        self.request_dcp_stream(vbid, flags, start_seqno, end_seqno, vb_uuid, ss_start_seqno,
                                                ss_end_seqno)

                        del self.stream_list[opaque]
                        self.stream_list[opaque] = \
                            (vbid, flags, start_seqno, end_seqno, vb_uuid, ss_start_seqno, ss_end_seqno)
                    else:
                        logging.error(f'unprocessed errcode: {errcode}')
                        del self.stream_list[opaque]
                elif cmd == couchbaseConstants.CMD_DCP_MUTATION:
                    vbucket_id = errcode
                    seqno, rev_seqno, flg, exp, locktime, metalen, nru = \
                        struct.unpack(couchbaseConstants.DCP_MUTATION_PKT_FMT, data[0:extlen])
                    key_start = extlen
                    val_start = key_start + keylen
                    val_len = datalen - keylen - metalen - extlen
                    meta_start = val_start + val_len
                    key = data[extlen:val_start]
                    val = data[val_start:meta_start]
                    conf_res = 0
                    if meta_start < datalen:
                        # handle extra conflict resolution fields
                        extra_meta = data[meta_start:]
                        extra_index = 0
                        version = extra_meta[extra_index]
                        extra_index += 1
                        while extra_index < metalen:
                            id, extlen = struct.unpack(couchbaseConstants.DCP_EXTRA_META_PKG_FMT,
                                                       extra_meta[extra_index: extra_index + 3])
                            extra_index += 3
                            if id == couchbaseConstants.DCP_EXTRA_META_CONFLICT_RESOLUTION:
                                if extlen == 1:
                                    conf_res, = struct.unpack(">B", extra_meta[extra_index:extra_index + 1])
                                elif extlen == 2:
                                    conf_res, = struct.unpack(">H", extra_meta[extra_index:extra_index + 2])
                                elif extlen == 4:
                                    conf_res, = struct.unpack(">I", extra_meta[extra_index:extra_index + 4])
                                elif extlen == 8:
                                    conf_res, = struct.unpack(">Q", extra_meta[extra_index:extra_index + 8])
                                else:
                                    logging.error(f'unsupported extra meta data format: {extlen:d}')
                                    conf_res = 0
                            extra_index += extlen

                    if not self.skip(key, vbucket_id):
                        dtype, val = self.maybe_uncompress_value(dtype, val)
                        msg = (cmd, vbucket_id, key, flg, exp, cas, rev_seqno.to_bytes(8, 'big'), val, seqno, dtype,
                               metalen, conf_res)
                        batch.append(msg, len(val))
                        self.num_msg += 1
                elif cmd in [couchbaseConstants.CMD_DCP_DELETE, couchbaseConstants.CMD_DCP_EXPIRATION]:
                    vbucket_id = errcode
                    seqno, rev_seqno, metalen = struct.unpack(couchbaseConstants.DCP_DELETE_PKT_FMT, data[0:extlen])
                    key_start = extlen
                    val_start = key_start + keylen
                    key = data[extlen:val_start]
                    # If the delete has the Xattr data type get the Xattrs from the body
                    if dtype & couchbaseConstants.DATATYPE_HAS_XATTR:
                        val = data[val_start:]
                    if not self.skip(key, vbucket_id):
                        dtype, val = self.maybe_uncompress_value(dtype, val)
                        msg = (cmd, vbucket_id, key, flg, exp, cas, rev_seqno.to_bytes(8, 'big'), val, seqno, dtype,
                               metalen, 0)
                        batch.append(msg, len(val))
                        self.num_msg += 1
                    if cmd == couchbaseConstants.CMD_DCP_DELETE:
                        batch.adjust_size += 1
                elif cmd == couchbaseConstants.CMD_DCP_FLUSH:
                    total_bytes_read -= bytes_read
                    logging.warning("stopping: saw CMD_DCP_FLUSH")
                    self.dcp_done = True
                    break
                elif cmd == couchbaseConstants.CMD_DCP_END_STREAM:
                    del self.stream_list[opaque]
                    if not len(self.stream_list):
                        self.dcp_done = True
                elif cmd == couchbaseConstants.CMD_DCP_SNAPSHOT_MARKER:
                    ss_start_seqno, ss_end_seqno, _ = struct.unpack(couchbaseConstants.DCP_SNAPSHOT_PKT_FMT,
                                                                    data[0:extlen])
                    pair_index = (self.source_bucket['name'], self.source_node['hostname'])
                    if not self.cur['snapshot']:
                        self.cur['snapshot'] = {}
                    if pair_index not in self.cur['snapshot']:
                        self.cur['snapshot'][pair_index] = {}
                    self.cur['snapshot'][pair_index][opaque] = (ss_start_seqno, ss_end_seqno)
                elif cmd == couchbaseConstants.CMD_DCP_NOOP:
                    total_bytes_read -= bytes_read
                    need_ack = True
                elif cmd == couchbaseConstants.CMD_DCP_BUFFER_ACK:
                    total_bytes_read -= bytes_read
                    if errcode != couchbaseConstants.ERR_SUCCESS:
                        logging.warning(f'buffer ack response errcode: {errcode}')
                    continue
                else:
                    total_bytes_read -= bytes_read
                    logging.warning(f'warning: unexpected DCP message: {cmd}')
                    return f'unexpected DCP message: {cmd}', batch

                if need_ack:
                    self.ack_last = True
                    try:
                        if self.dcp_conn is not None:
                            self.dcp_conn._send_msg(cmd, b'', b'', opaque, vbucket_id=0,
                                                    fmt=couchbaseConstants.RES_PKT_FMT,
                                                    magic=couchbaseConstants.RES_MAGIC_BYTE)
                    except socket.error:
                        return f'error: socket.error on sendall(); perhaps the source server:' \
                            f' {self.source_node["hostname"]} was rebalancing or had' \
                            f' connectivity/server problems', batch
                    except EOFError:
                        self.dcp_done = True
                        return f'error: EOFError on socket sendall(); perhaps the source server:' \
                            f' {self.source_node["hostname"]} was rebalancing or had ' \
                            f'connectivity/server problems', batch

                    # Close the batch when there's an ACK handshake, so
                    # the server can concurrently send us the next batch.
                    # If we are slow, our slow ACK's will naturally slow
                    # down the server.
                    self.ack_buffer_size(total_bytes_read - last_processed)
                    return 0, batch

                self.ack_last = False
                self.cmd_last = cmd

        except EOFError:
            if batch.size() <= 0 and self.ack_last:
                # A closed conn after an ACK means clean end of TAP dump.
                self.dcp_done = True

        if batch.size() <= 0:
            return 0, None
        self.ack_buffer_size(total_bytes_read - last_processed)
        return 0, batch

    def get_dcp_conn(self):
        """Return previously connected dcp conn."""

        if not self.dcp_conn:
            host, _ = pump.hostport(self.source_node['hostname'])
            port = self.source_node['ports']['direct']
            username = self.opts.username
            password = self.opts.password
            bucket = str(self.source_bucket.get("name"))
            if self.opts.ssl:
                port = couchbaseConstants.SSL_PORT

            if getattr(self, 'alt_add', None):
                if 'alternateAddresses' in self.source_node:
                    host = self.source_node['alternateAddresses']['external']['hostname']
                    if 'ports' not in self.source_node['alternateAddresses']['external']:
                        return ['no data port available in external address: {}'. format(host)], None

                    alterante_ports = self.source_node['alternateAddresses']['external']['ports']
                    if self.opts.ssl:
                        if 'kvSSL' not in alterante_ports:
                            return ['Secure data port not available in external host: {}'.format(host)]
                        port = alterante_ports['kvSSL']
                    elif 'kv' not in alterante_ports:
                        return ['Data port not available in external host: {}'.format(host)]
                    else:
                        port = alterante_ports['kv']
                else:
                    return ['alterante address not available in node: {}'.format(self.source_node['otpNode'])], None

            logging.debug(f'  DCPStreamSource connecting mc: {host}:{port!s}')

            err, self.dcp_conn = pump.get_mcd_conn(host, port, username, password, bucket, self.opts.ssl,
                                                   not self.opts.no_ssl_verify, self.opts.cacert)
            if err:
                return err, None

            err, self.mem_conn = pump.get_mcd_conn(host, port, username, password, bucket, self.opts.ssl,
                                                   not self.opts.no_ssl_verify, self.opts.cacert)
            if err:
                return err, None

            flags = couchbaseConstants.FLAG_DCP_PRODUCER | couchbaseConstants.FLAG_DCP_XATTRS
            extra = struct.pack(couchbaseConstants.DCP_CONNECT_PKT_FMT, 0, flags)
            try:
                opaque = self.r.randint(0, 2**32)
                self.dcp_conn._send_cmd(couchbaseConstants.CMD_DCP_CONNECT, self.dcp_name.encode(), b'', opaque, extra)
                self.dcp_conn._handle_single_response(opaque)

                opaque = self.r.randint(0, 2**32)
                self.dcp_conn._send_cmd(couchbaseConstants.CMD_DCP_CONTROL,
                                        couchbaseConstants.KEY_DCP_CONNECTION_BUFFER_SIZE,
                                        str(self.batch_max_bytes * 10).encode(), opaque)
                self.dcp_conn._handle_single_response(opaque)

                opaque = self.r.randint(0, 2**32)
                self.dcp_conn._send_cmd(couchbaseConstants.CMD_DCP_CONTROL,
                                        couchbaseConstants.KEY_DCP_NOOP, b'true', opaque)
                self.dcp_conn._handle_single_response(opaque)

                opaque = self.r.randint(0, 2**32)
                self.dcp_conn._send_cmd(couchbaseConstants.CMD_DCP_CONTROL,
                                        couchbaseConstants.KEY_DCP_NOOP_INTERVAL,
                                        b'180', opaque)
                self.dcp_conn._handle_single_response(opaque)
            except EOFError:
                return "error: Fail to set up DCP connection"
            except cb_bin_client.MemcachedError:
                return "error: DCP connect memcached error"
            except socket.error:
                return "error: DCP connection error"
            try:
                opaque = self.r.randint(0, 2**32)
                self.dcp_conn._send_cmd(couchbaseConstants.CMD_DCP_CONTROL,
                                        couchbaseConstants.KEY_DCP_EXT_METADATA,
                                        bool_to_str(True).encode(), opaque)
                self.dcp_conn._handle_single_response(opaque)
            except EOFError:
                return "error: Fail to set up DCP connection"
            except cb_bin_client.MemcachedError:
                pass
            except socket.error:
                return "error: DCP connection error"

            self.running = True
            self.start()

            self.add_start_event(self.dcp_conn)
            self.setup_dcp_streams()
        return 0

    def ack_buffer_size(self, buf_size: int) -> couchbaseConstants.PUMP_ERROR:
        if self.flow_control:
            try:
                opaque = 0
                if self.dcp_conn is not None:
                    self.dcp_conn._send_cmd(couchbaseConstants.CMD_DCP_BUFFER_ACK, b'', b'', opaque,
                                            struct.pack(">I", int(buf_size)))
                    logging.debug(f'Send buffer size: {buf_size:d}')
            except socket.error:
                return "error: socket error during sending buffer ack msg"
            except EOFError:
                return "error: send buffer ack msg"

        return 0

    def run(self):
        if not self.dcp_conn:
            logging.error("socket to memcached server is not created yet.")
            return

        bytes_read = b''
        rd_timeout = 1
        desc = [self.dcp_conn.s]
        while self.running:
            try:
                if self.dcp_done:
                    time.sleep(1)
                    continue
                readers, _, _ = select.select(desc, [], [], rd_timeout)
                rd_timeout = .25
                if len(self.stream_list) == 0:
                    time.sleep(1)
                    continue

                for reader in readers:
                    data = reader.recv(self.recv_min_bytes)
                    logging.debug(f'Read {len(data)} bytes off the wire')
                    if len(data) == 0:
                        raise EOFError("Got empty data (remote died?).")
                    bytes_read += data
                while len(bytes_read) >= couchbaseConstants.MIN_RECV_PACKET:
                    magic, opcode, keylen, extlen, datatype, status, bodylen, opaque, cas = \
                        struct.unpack(couchbaseConstants.RES_PKT_FMT,
                                      bytes_read[0:couchbaseConstants.MIN_RECV_PACKET])

                    if len(bytes_read) < (couchbaseConstants.MIN_RECV_PACKET + bodylen):
                        break

                    rd_timeout = 0
                    body = bytes_read[couchbaseConstants.MIN_RECV_PACKET:couchbaseConstants.MIN_RECV_PACKET + bodylen]
                    bytes_read = bytes_read[couchbaseConstants.MIN_RECV_PACKET + bodylen:]
                    self.response.put((opcode, status, opaque, cas, keylen, extlen, body, bodylen, datatype,
                                       couchbaseConstants.MIN_RECV_PACKET + bodylen))
            except socket.error:
                break
            except Exception:
                pass

    def setup_dcp_streams(self):
        # send request to retrieve vblist and uuid for the node
        stats = self.mem_conn.stats(b'vbucket-seqno')
        if not stats:
            return "error: fail to retrive vbucket seqno", None
        self.mem_conn.close()

        vb_list = {}
        for key, val in stats.items():
            vb, counter = key.decode().split(":")
            if counter in [DCPStreamSource.VB_UUID,
                           DCPStreamSource.HIGH_SEQNO,
                           DCPStreamSource.ABS_HIGH_SEQNO,
                           DCPStreamSource.PURGE_SEQNO]:
                if vb[3:] not in vb_list:
                    vb_list[vb[3:]] = {}
                vb_list[vb[3:]][counter] = int(val)
        flags = 0
        pair_index = (self.source_bucket['name'], self.source_node['hostname'])
        vbuckets_or_dict = None
        vbuckets = None
        if self.vbucket_list:
            vbuckets_or_dict = json.loads(self.vbucket_list)
            if isinstance(vbuckets_or_dict, dict):
                vbuckets = []
                for key, value in vbuckets_or_dict.items():
                    vbuckets += value
            else:
                vbuckets = vbuckets_or_dict
        for vbid in vb_list:
            if int(vbid) not in self.node_vbucket_map:
                # skip nonactive vbucket
                continue
            if vbuckets and int(vbid) not in vbuckets:
                # skip vbuckets that are not in this run
                continue

            if self.cur['seqno'] and self.cur['seqno'][pair_index]:
                start_seqno = self.cur['seqno'][pair_index][vbid]
            else:
                start_seqno = 0
            uuid = 0
            if self.cur['failoverlog'] and self.cur['failoverlog'][pair_index]:
                if vbid in self.cur['failoverlog'][pair_index] and \
                   self.cur['failoverlog'][pair_index][vbid]:
                    # Use the latest failover log
                    self.cur['failoverlog'][pair_index][vbid] = \
                        sorted(self.cur['failoverlog'][pair_index][vbid], key=lambda tup: tup[1], reverse=True)
                    uuid, _ = self.cur['failoverlog'][pair_index][vbid][0]
            ss_start_seqno = start_seqno
            ss_end_seqno = start_seqno
            if (self.cur['snapshot'] and self.cur['snapshot'][pair_index] and vbid in self.cur['snapshot'][pair_index]
                    and self.cur['snapshot'][pair_index][vbid]):
                ss_start_seqno, ss_end_seqno = self.cur['snapshot'][pair_index][vbid]
                if start_seqno == ss_end_seqno:
                    ss_start_seqno = start_seqno

            self.request_dcp_stream(int(vbid), flags, start_seqno, vb_list[vbid][DCPStreamSource.HIGH_SEQNO],
                                    uuid, ss_start_seqno, ss_end_seqno)

    def request_dcp_stream(self,
                           vbid: int,
                           flags: int,
                           start_seqno: int,
                           end_seqno: int,
                           vb_uuid: int,
                           ss_start_seqno: int,
                           ss_end_seqno: int):
        if not self.dcp_conn:
            return "error: no dcp connection setup yet.", None
        extra = struct.pack(couchbaseConstants.DCP_STREAM_REQ_PKT_FMT,
                            int(flags), 0,
                            int(start_seqno),
                            int(end_seqno),
                            int(vb_uuid),
                            int(ss_start_seqno),
                            int(ss_end_seqno))
        self.dcp_conn._send_msg(couchbaseConstants.CMD_DCP_REQUEST_STREAM, b'', b'', vbid, extra, 0, 0, vbid)
        self.stream_list[vbid] = (vbid, flags, start_seqno, end_seqno, vb_uuid, ss_start_seqno, ss_end_seqno)

    @staticmethod
    def total_msgs(opts, source_bucket: Dict[str, Any], source_node, source_map: Dict[str, Any]) ->\
            Tuple[couchbaseConstants.PUMP_ERROR, Optional[int]]:
        source_name = source_node.get("hostname", None)
        if not source_name:
            return 0, None

        spec = source_map['spec']
        name = source_bucket['name']
        vbuckets_num = len(source_bucket['vBucketServerMap']['vBucketMap'])
        if not vbuckets_num:
            return 0, None

        stats_vals = {}
        host, port, user, pswd, _ = pump.parse_spec(opts, spec, 8091)
        for stats in ["curr_items", "vb_active_resident_items_ratio"]:
            path = f'/pools/default/buckets/{name}/stats/{stats}'
            err, json, data = pump.rest_request_json(host, int(port),
                                                     user, pswd, opts.ssl, path, reason="total_msgs",
                                                     verify=not opts.no_ssl_verify, cacert=opts.cacert)
            if err:
                return 0, None

            if data is not None:
                node_stats = data.get("nodeStats", None)

            if not node_stats:
                return 0, None
            vals = node_stats.get(source_name, None)
            if not vals:
                return 0, None
            stats_vals[stats] = vals[-1]

        total_msgs = stats_vals["curr_items"]
        resident_ratio = stats_vals["vb_active_resident_items_ratio"]
        if 0 < resident_ratio < 100:
            # for DGM case, server will transfer both in-memory items and
            # backfill all items on disk
            total_msgs += (resident_ratio / 100.0) * stats_vals["curr_items"]
        return 0, int(total_msgs) class GenSource(pump.Source):
    """Generates simple SET/GET workload, useful for basic testing.
       Examples:
         ./cbtransfer gen: http://10.3.121.192:8091
         ./cbtransfer gen:max-items=50000,min-value-size=10,exit-after-creates=1,\
                          prefix=steve1-,ratio-sets=1.0 \
                      http://10.3.121.192:8091 -B my-other-bucket --threads=10
    """

    def __init__(self, opts, spec, source_bucket, source_node,
                 source_map, sink_map, ctl, cur):
        super(GenSource, self).__init__(opts, spec, source_bucket, source_node,
                                        source_map, sink_map, ctl, cur)
        self.done = False
        self.body = None
        self.cur_ops = source_map['cfg']['cur-ops']
        self.cur_gets = source_map['cfg']['cur-gets']
        self.cur_sets = source_map['cfg']['cur-sets']
        self.cur_items = source_map['cfg']['cur-items']

    @staticmethod
    def can_handle(opts, spec: str) -> bool:
        """The gen spec follows gen:[key=value[,key=value]] format."""
        return spec.startswith("gen:")

    @staticmethod
    def check(opts, spec: str) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[Dict[str, Any]]]:
        rv, cfg = GenSource.parse_spec(opts, spec)
        if rv != 0:
            return rv, None
        return 0, {'cfg': cfg,
                   'spec': spec,
                   'buckets': [{'name': 'default',
                                'nodes': [{'hostname': f'N/A-{i!s}'}
                                          for i in range(opts.threads)]}]}

    @staticmethod
    def parse_spec(opts, spec: str) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[Dict[str, Any]]]:
        """Parse the comma-separated key=value configuration from the gen spec.
           Names and semantics were inspired from subset of mcsoda parameters."""
        cfg = {'cur-ops': 0,
               'cur-gets': 0,
               'cur-sets': 0,
               'cur-items': 0,
               'exit-after-creates': 0,
               'max-items': 10000,
               'min-value-size': 10,
               'prefix': "",
               'ratio-sets': 0.05,
               'json': 0,
               'low-compression': False,
               'xattr': False}

        for kv in spec[len("gen:"):].split(','):
            if kv:
                k = kv.split('=')[0].strip()
                v = kv.split('=')[1].strip()
                try:
                    if k in cfg:
                        # this type casting does not work for bools so this code has to be added
                        if v == 'True':
                            cfg[k] = True
                        elif v == 'False':
                            cfg[k] = False
                        else:
                            cfg[k] = type(cfg[k])(v)  # type: ignore
                    else:
                        return f'error: unknown workload gen parameter: {k}', None
                except ValueError:
                    return f'error: could not parse value from: {kv}', None
        return 0, cfg

    @staticmethod
    def provide_design(opts, source_spec, source_bucket, source_map) -> Tuple[couchbaseConstants.PUMP_ERROR,
                                                                              Optional[str]]:
        """No design from a GenSource."""
        return 0, None

    def provide_batch(self) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[pump.Batch]]:
        """Provides a batch of messages, with GET/SET ratios and keys
           controlled by a mcsoda-inspired approach, but simpler."""
        if self.done:
            return 0, None

        cfg: Dict[str, Any] = self.source_map['cfg']
        prefix: str = cfg['prefix']
        max_items: int = cfg['max-items']
        ratio_sets: float = cfg['ratio-sets']
        exit_after_creates: bool = cfg['exit-after-creates']
        low_compression: bool = cfg['low-compression']
        xattrs: bool = cfg['xattr']
        itr = None
        collections = self.opts.collection
        if collections:
            itr = iter(collections)

        json_body: bool = cfg['json']
        if not self.body:

            if low_compression:
                # Generate a document which snappy will struggle to compress.
                # Useful if your creating data-sets which utilise disk.
                random.seed(0)  # Seed to a fixed value so we always have the same document pattern.
                document = ''.join(random.choice(string.ascii_uppercase) for _ in range(cfg['min-value-size']))
            else:
                # else a string of 0 is fine, but will compress very well.
                document = "0" * cfg['min-value-size']

            self.body = document

        batch = pump.Batch(self)

        batch_max_size = self.opts.extra['batch_max_size']
        batch_max_bytes = self.opts.extra['batch_max_bytes']

        vbucket_id = 0x0000ffff
        cas, exp, flg = 0, 0, 0

        while (batch.size() < batch_max_size
               and batch.bytes < batch_max_bytes):
            if ratio_sets >= float(self.cur_sets) / float(self.cur_ops or 1):
                self.cur_sets = self.cur_sets + 1
                if xattrs:
                    cmd: int = couchbaseConstants.CMD_SUBDOC_MULTIPATH_MUTATION
                else:
                    cmd = couchbaseConstants.CMD_DCP_MUTATION
                if self.cur_items < max_items:
                    key = str(self.cur_items)
                    self.cur_items = self.cur_items + 1
                else:
                    key = str(self.cur_sets % self.cur_items)
            else:
                self.cur_gets = self.cur_gets + 1
                if xattrs:
                    cmd = couchbaseConstants.CMD_SUBDOC_MULTIPATH_LOOKUP
                else:
                    cmd = couchbaseConstants.CMD_GET
                key = str(self.cur_gets % self.cur_items)
            self.cur_ops = self.cur_ops + 1

            if json_body:
                value = f'{{"name": "{prefix}{key}", "age": {int(key) % 101}, "index": "{key}", "body":"{self.body}"}}'
            else:
                value = self.body

            if xattrs:
                value = json.dumps({"obj": value, "xattr_f": "field1", "xattr_v": "\"value1\""})

            value_bytes: bytes = value.encode()
            # generate a collection key
            if itr:
                try:
                    cid = int(next(itr), 16)
                except StopIteration:
                    itr = iter(collections)
                    cid = int(next(itr), 16)
                except ValueError as e:
                    return f'Invalid collection id, collection id must be a hexadecimal number: {e}', None

                encoded_cid = encode_collection_id(cid)
                # Generate the pack format and pack the key
                doc_key: bytes = struct.pack(
                    ("!" + str(len(encoded_cid)) + "s"
                        + str(len(prefix)) + "s"
                        + str(len(key)) + "s").encode(),
                    encoded_cid,
                    prefix.encode(),
                    key.encode())
            else:
                doc_key = prefix.encode() + key.encode()

            datatype = 0x00
            if json_body and cmd != couchbaseConstants.CMD_GET:
                datatype = 0x01

            msg: couchbaseConstants.BATCH_MSG = (cmd, vbucket_id, doc_key, flg, exp, cas, b'', value_bytes, 0, datatype,
                                                 0, 0)
            batch.append(msg, len(value_bytes))

            if exit_after_creates and self.cur_items >= max_items:
                self.done = True
                return 0, batch

        if batch.size() <= 0:
            return 0, None
        return 0, batch

    @staticmethod
    def total_msgs(opts, source_bucket, source_node, source_map) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[int]]:
        """Returns max-items only if exit-after-creates was specified.
           Else, total msgs is unknown as GenSource does not stop generating."""
        if source_map['cfg']['exit-after-creates'] and source_map['cfg']['ratio-sets'] > 0:
            ratio = source_map['cfg']['ratio-sets']
            ops = 0
            sets = 0
            while sets != source_map['cfg']['max-items']:
                if ratio >= float(sets) / float(ops or 1):
                    sets += 1
                ops += 1
            return 0, ops
        return 0, None class JSONSource(pump.Source):
    """Reads json file or directory or zip file that contains json files."""

    def __init__(self, opts, spec, source_bucket, source_node,
                 source_map, sink_map, ctl, cur):
        super(JSONSource, self).__init__(opts, spec, source_bucket, source_node,
                                         source_map, sink_map, ctl, cur)
        self.done: bool = False
        self.docs = list()
        self.file_iter = None

    @staticmethod
    def can_handle(opts, spec: str) -> bool:
        return spec.startswith(JSON_SCHEME) and \
            (os.path.isfile(spec.replace(JSON_SCHEME, ""))
             or os.path.isdir(spec.replace(JSON_SCHEME, ""))
             or spec.endswith(".zip"))

    @staticmethod
    def check(opts, spec: str) -> Tuple[couchbaseConstants.PUMP_ERROR, Dict[str, Any]]:
        return 0, {'spec': spec,
                   'buckets': [{'name': os.path.normpath(os.path.basename(spec)),
                                'nodes': [{'hostname': 'N/A'}]}]}

    def save_doc(self, batch: pump.Batch, dockey: bytes, docvalue: bytes):
        cmd = couchbaseConstants.CMD_TAP_MUTATION
        vbucket_id = 0x0000ffff
        # common flags:       0x02000000 (JSON)
        # legacy flags:       0x00000006 (JSON)
        cas, exp, flg = 0, 0, 0x02000006
        try:
            doc = json.loads(docvalue)
            if '_id' not in doc:
                msg = (cmd, vbucket_id, dockey, flg, exp, cas, b'', docvalue, 0, 0, 0, 0)
                batch.append(msg, len(docvalue))
        except ValueError as error:
            logging.error(f'Fail to read json file with error: {error!s}')

    @staticmethod
    def gen_dockey(filename: str) -> str:
        return os.path.splitext(os.path.basename(filename))[0]

    @staticmethod
    def enumerate_files(subdir: str, file_candidate: List[str], skip_views: bool, skip_docs: bool):
        for item in os.listdir(subdir):
            path = os.path.join(subdir, item)
            if os.path.isfile(path):
                if (not skip_views and "design_docs" in path.split(os.path.sep)) or \
                   (not skip_docs and "docs" in path.split(os.path.sep)):
                    file_candidate.append(path)
            else:
                if not ((skip_docs and "docs" in path.split(os.path.sep))
                   or (skip_views and "design_docs" in path.split(os.path.sep))):
                    JSONSource.enumerate_files(path, file_candidate, skip_views, skip_docs)

    @staticmethod
    def provide_design(opts, source_spec: str, source_bucket, source_map) -> Tuple[couchbaseConstants.PUMP_ERROR,
                                                                                   List[bytes]]:
        design_files = list()
        f = source_spec.replace(JSON_SCHEME, "")

        if os.path.isfile(f) and f.endswith(".zip"):
            zf = zipfile.ZipFile(f)
            for path in zf.namelist():
                file = os.path.basename(path)
                # Skip the design_docs directory listing
                if file == "design_docs":
                    continue

                dir = os.path.basename(os.path.dirname(path))
                # Skip all files not in the design docs directory
                if dir != "design_docs":
                    continue

                design_files.append(zf.read(path))
            zf.close()
        elif os.path.isdir(f):
            files: List[str] = list()
            JSONSource.enumerate_files(f, files, False, True)
            for path in files:
                if os.path.isfile(path):
                    f = open(path, 'rb')  # type: ignore
                    design_files.append(f.read())  # type: ignore
                    f.close()  # type: ignore

        return 0, design_files

    def provide_batch(self) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[pump.Batch]]:
        if self.done:
            return 0, None

        # During the first iteration load the file names, this is only run once
        if not self.docs:
            self.prepare_docs()

        batch = pump.Batch(self)
        f = self.spec.replace(JSON_SCHEME, "")
        batch_max_size = self.opts.extra['batch_max_size']

        # Each iteration should return a batch or mark the loading a finished
        if os.path.isfile(f) and f.endswith(".zip"):
            zf = zipfile.ZipFile(f)
            while batch.size() < batch_max_size and self.docs:
                path = self.docs.pop()
                key = os.path.basename(path)
                if key.endswith('.json'):
                    key = key[:-5]
                value = zf.read(path)
                self.save_doc(batch, key, value)
            zf.close()
        else:
            while batch.size() < batch_max_size and self.docs:
                path = self.docs.pop()
                key = os.path.basename(path)
                if key.endswith('.json'):
                    key = key[:-5]
                try:
                    fp = open(path, 'rb')
                    value = fp.read()
                    fp.close()
                    self.save_doc(batch, key.encode(), value)
                except IOError as error:
                    logging.error(f'Fail to load json file with error: {error!s}')

        if not self.docs:
            self.done = True

        return 0, batch

    def prepare_docs(self):
        f = self.spec.replace(JSON_SCHEME, "")
        root_name = os.path.basename(f).split('.')[0]
        if os.path.isfile(f) and f.endswith(".zip"):
            zf = zipfile.ZipFile(f)
            for path in zf.namelist():
                file = os.path.basename(path)
                # Skip the docs directory listing
                if file == "docs":
                    continue

                dir = os.path.basename(os.path.dirname(path))

                # This condition is not allowed by the spec, but we allowing it
                # because the training team did properly follow the spec and we
                # don't want to break their training material. Since this tool
                # will be deprecated soon we are making an exception and
                # allowing this.
                if dir == "" or dir == root_name:
                    self.docs.append(path)
                    continue

                # Skip all files not in the docs directory
                # Note that we use a forward slash for all operating systems
                # because this is a zip file and zip paths always use a forward
                # slash
                if "docs" not in path.split("/"):
                    continue

                self.docs.append(path)
            zf.close()
        elif os.path.isdir(f):
            JSONSource.enumerate_files(f, self.docs, True, False)
        else:
            self.docs.append(f) class MCSink(pump.Sink):
    """Dumb client sink using binary memcached protocol.
       Used when moxi or memcached is destination."""

    def __init__(self, opts, spec, source_bucket, source_node,
                 source_map, sink_map, ctl, cur):
        super(MCSink, self).__init__(opts, spec, source_bucket, source_node,
                                     source_map, sink_map, ctl, cur)

        self.op_map = OP_MAP
        if opts.extra.get("try_xwm", 1):
            self.op_map = OP_MAP_WITH_META
        self.conflict_resolve = opts.extra.get("conflict_resolve", 1)
        self.lww_restore = 0
        self.init_worker(MCSink.run)
        self.uncompress = opts.extra.get("uncompress", 0)
        self.txn_warning_issued = False
        if self.get_conflict_resolution_type() == "lww":
            self.lww_restore = 1

    def close(self):
        self.push_next_batch(None, None)

    @staticmethod
    def check_base(opts, spec: str) -> couchbaseConstants.PUMP_ERROR:
        if getattr(opts, "destination_vbucket_state", "active") != "active":
            return f'error: only --destination-vbucket-state=active is supported by this destination: {spec}'

        op = getattr(opts, "destination_operation", None)
        if op not in [None, 'set', 'add', 'get']:
            return f'error: --destination-operation unsupported value: {op}; use set, add, get'
        # Skip immediate superclass Sink.check_base(),
        # since MCSink can handle different destination operations.
        return pump.EndPoint.check_base(opts, spec)

    @staticmethod
    def run(self):
        """Worker thread to asynchronously store batches into sink."""

        mconns: Dict[str, cb_bin_client.MemcachedClient] = {}  # State kept across scatter_gather() calls.
        backoff_cap: int = self.opts.extra.get("backoff_cap", 10)
        while not self.ctl['stop']:
            batch, future = self.pull_next_batch()  # type: Optional[pump.Batch], pump.SinkBatchFuture
            if not batch:
                self.future_done(future, 0)
                self.close_mconns(mconns)
                return

            backoff = 0.1  # Reset backoff after a good batch.

            while batch:  # Loop in case retry is required.
                rv, batch, need_backoff = self.scatter_gather(mconns, batch)
                if rv != 0:
                    self.future_done(future, rv)
                    self.close_mconns(mconns)
                    return

                if batch:
                    self.cur["tot_sink_retry_batch"] = \
                        self.cur.get("tot_sink_retry_batch", 0) + 1

                if need_backoff:
                    backoff = min(backoff * 2.0, backoff_cap)
                    logging.warning(f'backing off, secs: {backoff}')
                    time.sleep(backoff)

            self.future_done(future, 0)

        self.close_mconns(mconns)

    def get_conflict_resolution_type(self) -> str:
        bucket = self.sink_map["buckets"][0]
        conf_res_type = "seqno"
        if "conflictResolutionType" in bucket:
            conf_res_type = bucket["conflictResolutionType"]
        return conf_res_type

    def close_mconns(self, mconns: Dict[str, cb_bin_client.MemcachedClient]):
        for conn in mconns.values():
            self.add_stop_event(conn)
            conn.close()

    def scatter_gather(self, mconns: Dict[str, cb_bin_client.MemcachedClient], batch: pump.Batch) -> \
            Tuple[couchbaseConstants.PUMP_ERROR, Optional[pump.Batch], Optional[bool]]:
        conn: Optional[cb_bin_client.MemcachedClient] = mconns.get("conn")
        if not conn:
            rv, conn = self.connect()
            if rv != 0:
                return rv, None, None
            mconns["conn"] = conn  # type: ignore

        # TODO: (1) MCSink - run() handle --data parameter.

        # Scatter or send phase.
        rv, skipped = self.send_msgs(conn, batch.msgs, self.operation())  # type: ignore
        if rv != 0:
            return rv, None, None

        # Gather or recv phase.
        rv, retry, refresh = self.recv_msgs(conn, batch.msgs, skipped)  # type: ignore
        if refresh:
            self.refresh_sink_map()
        if retry:
            return rv, batch, True

        return rv, None, None

    def send_msgs(self, conn: cb_bin_client.MemcachedClient, msgs: List[couchbaseConstants.BATCH_MSG], operation: str,
                  vbucket_id: Optional[int] = None) -> Tuple[couchbaseConstants.PUMP_ERROR, List[int]]:
        m: List[bytes] = []
        skipped: List[int] = []

        msg_format_length = 0
        for i, msg in enumerate(msgs):
            if not msg_format_length:
                msg_format_length = len(msg)
            cmd, vbucket_id_msg, key, flg, exp, cas, meta, val = msg[:8]
            seqno = dtype = nmeta = conf_res = 0
            if msg_format_length > 8:
                seqno, dtype, nmeta, conf_res = msg[8:12]
            if vbucket_id is not None:
                vbucket_id_msg = vbucket_id

            if self.skip(key, vbucket_id_msg):
                continue

            if cmd == couchbaseConstants.CMD_SUBDOC_MULTIPATH_MUTATION:
                err, req = self.format_multipath_mutation(key, val, vbucket_id_msg, cas, i)
                if err:
                    return err, skipped
                self.append_req(m, req)
                continue
            if cmd == couchbaseConstants.CMD_SUBDOC_MULTIPATH_LOOKUP:
                err, req = self.format_multipath_lookup(key, val, vbucket_id_msg, cas, i)
                if err:
                    return err, skipped
                self.append_req(m, req)
                continue

            rv, translated_cmd = self.translate_cmd(cmd, operation, meta)
            if translated_cmd is None:
                return rv, skipped
            if dtype & couchbaseConstants.DATATYPE_COMPRESSED and self.uncompress and val:
                try:
                    val = snappy.uncompress(val)
                    dtype = dtype ^ couchbaseConstants.DATATYPE_COMPRESSED
                except Exception:
                    pass
            if translated_cmd == couchbaseConstants.CMD_GET:
                val, flg, exp, cas = b'', 0, 0, 0
            if translated_cmd == couchbaseConstants.CMD_NOOP:
                key, val, flg, exp, cas = b'', b'', 0, 0, 0
            if translated_cmd == couchbaseConstants.CMD_DELETE:
                val = b''
            # A tombstone can contain Xattrs
            if translated_cmd == couchbaseConstants.CMD_DELETE_WITH_META and not dtype & couchbaseConstants.DATATYPE_HAS_XATTR:
                val = b''
            # on mutations filter txn related data
            if (translated_cmd in [couchbaseConstants.CMD_SET_WITH_META, couchbaseConstants.CMD_SET]
                    and not getattr(self.opts, 'force_txn', False)):
                skip, val, cas, exp, meta, dtype = self.filter_out_txn(key, val, cas, exp, meta, dtype)
                if skip:
                    skipped.append(i)
                    if not self.txn_warning_issued:
                        self.txn_warning_issued = True
                        print('Mid-transaction data have being rolled backed and restored, but transactional '
                              'atomicity cannot be guaranteed.')
                    continue

            rv, req = self.cmd_request(translated_cmd, vbucket_id_msg, key, val,  # type: ignore
                                       ctypes.c_uint32(flg).value,
                                       exp, cas, meta, i, dtype, nmeta,
                                       conf_res)  # type: ignore
            if rv != 0:
                return rv, skipped

            self.append_req(m, req)

        if m:
            try:
                conn.s.sendall(self.join_str_and_bytes(m))  # type: ignore
            except socket.error as e:
                return f'error: conn.sendall() exception: {e}', skipped

        return 0, skipped

    @staticmethod
    def filter_out_txn(key: bytes, val: bytes, cas: int, exp: int, revid: bytes,
                       data_type: int) -> Tuple[bool, bytes, int, int, bytes, int]:
        """This function return signature is [skip:bool, val:bytes, cas:num, exp:num, revid:num, data_type:num]
           skip is true if the key matches an ATR or a client record, otherwise it will be false.
           If the value has XATTRS they will be checked fot txn object that if it exists will be removed"""

        # If it does not have XATTRS return all values as they where
        if not (data_type & couchbaseConstants.DATATYPE_HAS_XATTR > 0):
            return False, val, cas, exp, revid, data_type

        str_key = key.decode()
        if str_key == '_txn:client-record':
            logging.info('(TXN) Skipped the transfer of the txn-client-record')
            return True, val, cas, exp, revid, data_type

        if ATR_EXP.match(str_key):
            logging.info(f'(TXN) Skipped the transfer of the ATR: {str_key}')
            return True, val, cas, exp, revid, data_type

        # Get the length of the xattrs
        xattr_len_bytes = val[:4]
        xattr_len_int = int.from_bytes(xattr_len_bytes, byteorder='big', signed=False)

        # extended attributes length is invalid return as normal and let the server deal with it
        if xattr_len_int > len(val):
            return False, val, cas, exp, revid, data_type

        # look for a txn on the xattrs
        xattrs = val[4:xattr_len_int + 7]
        txn_marker: Optional[txnMarker] = None
        ix = 0
        while ix < len(xattrs):
            pair_len = int.from_bytes(xattrs[ix: ix + 4], byteorder='big', signed=False)
            ix += 4
            # get the key for the xattr object
            xattr_key = ''
            while ix + len(xattr_key) < len(xattrs) and xattrs[ix + len(xattr_key)] != 0:
                xattr_key += chr(xattrs[ix + len(xattr_key)])

            # check if the key is txn if it matches create a marker and break
            if xattr_key == 'txn':
                txn_marker = txnMarker(ix, pair_len, xattrs
                                       [ix + len(xattr_key) + 1: pair_len + ix - len(xattr_key) + 2])
                break

            # otherwise skip body and move on to next xattr
            ix += pair_len

        if txn_marker is None:
            return False, val, cas, exp, revid, data_type

        # check if txn is valid JSON if not let the server handle it
        try:
            txn_xattr = json.loads(txn_marker.body.decode())
        except Exception:
            logging.debug('(TXN) txn is invalid json')
            return False, val, cas, exp, revid, data_type

        # get op type from the txn
        if 'op' not in txn_xattr or 'type' not in txn_xattr['op']:
            logging.info('(TXN) Transaction extended attributes is missing required field "op"')
            return False, val, cas, exp, revid, data_type

        if txn_xattr['op']['type'] == 'insert':
            logging.info(f'(TXN) Skip document with key {tag_user_data(str_key)} as it is a TXN insert')
            return True, val, cas, exp, revid, data_type

        # get cas, revid and expiry from the txn
        if 'restore' not in txn_xattr:
            logging.info('(TXN) Transaction extended attributes is missing required field "restore"')
            return False, val, cas, exp, revid, data_type

        if 'CAS' in txn_xattr['restore']:
            cas = int(txn_xattr['restore']['CAS'], 16)
        if 'exptime' in txn_xattr['restore']:
            exp = int(txn_xattr['restore']['exptime'])
        if 'revid' in txn_xattr['restore']:
            revid = struct.pack('>Q', int(txn_xattr['restore']['revid']))

        # Remove txn xattrs from value
        logging.info(f'(TXN) Removing transaction extended attributes "{tag_user_data(txn_marker.body)}" '
                     f'from key {tag_user_data(str_key)}')

        new_xattr_len = xattr_len_int - txn_marker.length - 4
        # check if the are any xattrs left
        if new_xattr_len <= 0:
            # copy everything after the xattrs
            val = val[xattr_len_int + 4:]
            data_type = 0x00  # Raw bytes
            try:
                _ = json.loads(val.decode())
                data_type = 0x01  # Raw json
            except Exception:
                data_type = 0x00
            return False, val, cas, exp, revid, data_type

        # create new value package starting with new xattr len
        new_val = new_xattr_len.to_bytes(4, 'big', signed=False)
        if txn_marker.start != 4:
            new_val += val[4: txn_marker.start]

        new_val += val[4 + txn_marker.start + txn_marker.length:]
        return False, new_val, cas, exp, revid, data_type

    @staticmethod
    def format_multipath_mutation(key: bytes, value: bytes, vbucket_id: int, cas: int = 0, opaque: int = 0) -> \
            Tuple[couchbaseConstants.PUMP_ERROR, couchbaseConstants.REQUEST]:
        empty_tuple = (b'', None, None, None, None)
        try:
            json_value: Dict[str, Any] = json.loads(value)
        except Exception:
            return 'value has invalid format for multipath mutation', empty_tuple

        if 'obj' not in json_value:
            return 'value has invalid format for multipath mutation', empty_tuple
        if 'xattr_f' not in json_value:
            return 'value has invalid format for multipath mutation', empty_tuple
        if 'xattr_v' not in json_value:
            return 'value has invalid format for multipath mutation', empty_tuple

        obj: bytes = to_bytes(json_value['obj'])
        xattr_f: bytes = to_bytes(json_value['xattr_f'])
        xattr_v: bytes = to_bytes(json_value['xattr_v'])

        subop_format = ">BBHI"
        sbcmd_len = 8 * 2 + len(obj) + len(xattr_f) + len(xattr_v)
        total_body_len = sbcmd_len + 1 + len(key)
        subcmd_msg0 = struct.pack(subop_format, couchbaseConstants.CMD_SUBDOC_DICT_UPSERT,
                                  couchbaseConstants.SUBDOC_FLAG_XATTR_PATH, len(xattr_f), len(xattr_v))
        subcmd_msg0 += xattr_f + xattr_v
        subcmd_msg1 = struct.pack(subop_format, couchbaseConstants.CMD_SET, 0, 0, len(obj))
        subcmd_msg1 += obj

        msg_head = struct.pack(couchbaseConstants.REQ_PKT_FMT, couchbaseConstants.REQ_MAGIC_BYTE,
                               couchbaseConstants.CMD_SUBDOC_MULTIPATH_MUTATION, len(key),
                               1, 0, vbucket_id, total_body_len, opaque, cas)
        extras = struct.pack(">B", couchbaseConstants.SUBDOC_FLAG_MKDOC)

        return 0, (msg_head + extras + key + subcmd_msg0 + subcmd_msg1, None, None, None, None)

    @staticmethod
    def format_multipath_lookup(key: bytes, value: bytes, vbucket_id: int, cas: int = 0, opaque: int = 0) -> \
            Tuple[couchbaseConstants.PUMP_ERROR, couchbaseConstants.REQUEST]:
        empty_tuple = (b'', None, None, None, None)
        try:
            json_value: Dict[str, Any] = json.loads(value)
        except Exception:
            return 'value has invalid format for multipath mutation', empty_tuple

        if 'xattr_f' not in json_value:
            return 'value has invalid format for multipath lookup', empty_tuple

        field: bytes = to_bytes(json_value['xattr_f'])

        subcmd_fmt = '>BBH'
        subcmd_msg0 = struct.pack(subcmd_fmt, couchbaseConstants.CMD_SUBDOC_GET,
                                  couchbaseConstants.SUBDOC_FLAG_XATTR_PATH, len(field))
        subcmd_msg0 += field
        subcmd_msg1 = struct.pack(subcmd_fmt, couchbaseConstants.CMD_GET, 0, 0)

        total_body_len = len(subcmd_msg0) + len(subcmd_msg1) + len(key)
        msg_head = struct.pack(couchbaseConstants.REQ_PKT_FMT, couchbaseConstants.REQ_MAGIC_BYTE,
                               couchbaseConstants.CMD_SUBDOC_MULTIPATH_LOOKUP, len(key),
                               0, 0, vbucket_id, total_body_len, opaque, cas)

        return 0, (msg_head + key + subcmd_msg0 + subcmd_msg1, None, None, None, None)

    @staticmethod
    def join_str_and_bytes(lst: List[bytes]) -> bytes:
        out = b''
        for x in lst:
            out += x
        return out

    def recv_msgs(self, conn: cb_bin_client.MemcachedClient, msgs: List[couchbaseConstants.BATCH_MSG],
                  skipped: List[int], vbucket_id: Optional[int] = None,
                  verify_opaque: bool = True) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[bool], Optional[bool]]:
        refresh = False
        retry = False
        for i, msg in enumerate(msgs):
            # some messages of the original batch where not sent as they were filter out due to been half-way
            # transaction related documents. The check bellow is to avoid waiting for responses for elements
            # that where not sent
            if len(skipped) > 0 and i == skipped[0]:
                skipped = skipped[1:]
                continue

            cmd, vbucket_id_msg, key, flg, exp, cas, meta, val = msg[:8]
            if vbucket_id is not None:
                vbucket_id_msg = vbucket_id

            if self.skip(key, vbucket_id_msg):
                continue

            try:
                r_cmd, r_status, r_ext, r_key, r_val, r_cas, r_opaque = \
                    self.read_conn(conn)  # type: ignore
                if verify_opaque and i != r_opaque:
                    return f'error: opaque mismatch: {i} {r_opaque}', None, None

                if r_status == couchbaseConstants.ERR_SUCCESS:
                    continue
                elif r_status == couchbaseConstants.ERR_KEY_EEXISTS:
                    continue
                elif r_status == couchbaseConstants.ERR_KEY_ENOENT:
                    if (cmd != couchbaseConstants.CMD_TAP_DELETE and
                            cmd != couchbaseConstants.CMD_GET):
                        logging.warning(f'item not found: {self.spec}, key: {tag_user_data(key)}')
                    continue
                elif (r_status == couchbaseConstants.ERR_ETMPFAIL or
                      r_status == couchbaseConstants.ERR_EBUSY or
                      r_status == couchbaseConstants.ERR_ENOMEM):
                    retry = True  # Retry the whole batch again next time.
                    continue      # But, finish recv'ing current batch.
                elif r_status == couchbaseConstants.ERR_NOT_MY_VBUCKET:
                    str_msg = f'received NOT_MY_VBUCKET; perhaps the cluster is/was rebalancing;' \
                        f' vbucket_id: {vbucket_id_msg}, key: {tag_user_data(key)}, spec: {self.spec},' \
                        f' host:port: {conn.host}:{conn.port}'
                    if self.opts.extra.get("nmv_retry", 1):
                        logging.warning(f'warning: {str_msg}')
                        refresh = True
                        retry = True
                        self.cur["tot_sink_not_my_vbucket"] = \
                            self.cur.get("tot_sink_not_my_vbucket", 0) + 1
                    else:
                        return f'error: {str_msg}', None, None
                elif r_status == couchbaseConstants.ERR_UNKNOWN_COMMAND:
                    if self.op_map == OP_MAP:
                        if not retry:
                            return f'error: unknown command: {r_cmd}', None, None
                    else:
                        if not retry:
                            logging.warning("destination does not take XXX-WITH-META"
                                            " commands; will use META-less commands")
                        self.op_map = OP_MAP
                        retry = True
                elif r_status == couchbaseConstants.ERR_ACCESS:
                    return json.loads(r_val)["error"]["context"], None, None
                elif r_status == couchbaseConstants.ERR_UNKNOWN_COLLECTION:
                    return json.loads(r_val)["error"]["context"], None, None
                else:
                    return "error: MCSink MC error: " + str(r_status), None, None

            except Exception as e:
                logging.error(f'MCSink exception: {e}')
                return f'error: MCSink exception: {e!s}', None, None
        return 0, retry, refresh

    def translate_cmd(self, cmd: int, op: str, meta: bytes) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[int]]:
        if len(meta) == 0:
            # The source gave no meta, so use regular commands.
            self.op_map = OP_MAP

        if cmd in [couchbaseConstants.CMD_TAP_MUTATION, couchbaseConstants.CMD_DCP_MUTATION]:
            m = self.op_map.get(op, None)
            if m:
                return 0, m
            return f'error: MCSink.translate_cmd, unsupported op: {op}', None

        if cmd in [couchbaseConstants.CMD_TAP_DELETE, couchbaseConstants.CMD_DCP_DELETE]:
            if op == 'get':
                return 0, couchbaseConstants.CMD_NOOP
            return 0, self.op_map['delete']

        if cmd == couchbaseConstants.CMD_GET:
            return 0, cmd

        return f'error: MCSink - unknown cmd: {cmd}, op: {op}', None

    def append_req(self, m: List[bytes], req: couchbaseConstants.REQUEST):
        hdr, ext, key, val, extra_meta = req
        m.append(hdr)
        if ext:
            m.append(ext)
        if key:
            m.append(key)
        if val:
            m.append(val)
        if extra_meta:
            m.append(extra_meta)

    @staticmethod
    def can_handle(opts, spec: str) -> bool:
        return (spec.startswith("memcached://") or
                spec.startswith("memcached-binary://"))

    @staticmethod
    def check(opts, spec: str, source_map) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[Dict[str, Any]]]:
        host, port, user, pswd, _ = pump.parse_spec(opts, spec, int(getattr(opts, "port", 11211)))
        if opts.ssl:
            port = couchbaseConstants.SSL_PORT
        rv, conn = MCSink.connect_mc(host, port, user, pswd, None, opts.ssl, not opts.no_ssl_verify, opts.cacert)
        if rv != 0:
            return rv, None
        conn.close()  # type: ignore
        return 0, None

    def refresh_sink_map(self) -> couchbaseConstants.PUMP_ERROR:
        return 0

    @staticmethod
    def consume_design(opts, sink_spec, sink_map,
                       source_bucket, source_map, source_design) -> couchbaseConstants.PUMP_ERROR:
        if source_design:
            logging.warn("warning: cannot restore bucket design"
                         " on a memached destination")
        return 0

    def consume_batch_async(self, batch: Optional[pump.Batch]) -> Tuple[couchbaseConstants.PUMP_ERROR,
                                                                        Optional[pump.SinkBatchFuture]]:
        return self.push_next_batch(batch, pump.SinkBatchFuture(self, batch))

    def connect(self) -> Tuple[couchbaseConstants.PUMP_ERROR, Optional[cb_bin_client.MemcachedClient]]:
        host, port, user, pswd, _ = pump.parse_spec(self.opts, self.spec, int(getattr(self.opts, "port", 11211)))
        if self.opts.ssl:
            port = couchbaseConstants.SSL_PORT
        return MCSink.connect_mc(host, port, user, pswd, self.sink_map["name"],
                                 self.opts.ssl, collections=self.opts.collection)

    @staticmethod
    def connect_mc(host: str, port: int, username: str, password: str, bucket: Optional[str], use_ssl: bool = False,
                   verify: bool = True, ca_cert: str = None, collections: bool = False) -> \
            Tuple[couchbaseConstants.PUMP_ERROR, Optional[cb_bin_client.MemcachedClient]]:
        return pump.get_mcd_conn(host, port, username, password, bucket, use_ssl=use_ssl, verify=verify,
                                 ca_cert=ca_cert, collections=collections)

    def cmd_request(self,
                    cmd: int,
                    vbucket_id: int,
                    key: bytes,
                    val: bytes,
                    flg: int,
                    exp: int,
                    cas: int,
                    meta: Optional[bytes],
                    opaque: int,
                    dtype: int,
                    nmeta: Any,
                    conf_res: Optional[str]) -> Tuple[couchbaseConstants.PUMP_ERROR,
                                                      Tuple[bytes,
                                                            bytes,
                                                            bytes,
                                                            bytes,
                                                            bytes]]:
        ext_meta = b''
        empty_tuple = (b'', b'', b'', b'', b'')
        if (cmd == couchbaseConstants.CMD_SET_WITH_META or
            cmd == couchbaseConstants.CMD_ADD_WITH_META or
                cmd == couchbaseConstants.CMD_DELETE_WITH_META):

            force = 0
            if int(self.conflict_resolve) == 0:
                force |= 1
            if int(self.lww_restore) == 1:
                force |= 2
            if meta is not None and len(meta) != 0:
                try:
                    ext = struct.pack(">IIQQI", flg, exp, int.from_bytes(meta, byteorder='big'), cas, force)
                except ValueError:
                    seq_no = str(meta)
                    if len(seq_no) > 8:
                        seq_no = seq_no[0:8]
                    if len(seq_no) < 8:
                        # The seq_no might be 32-bits from 2.0DP4, so pad with 0x00's.
                        seq_no = ('\x00\x00\x00\x00\x00\x00\x00\x00' + seq_no)[-8:]

                    seq_no_bytes: bytes = seq_no.encode()
                    check_seqno, = struct.unpack(">Q", seq_no_bytes)
                    if check_seqno:
                        ext = (struct.pack(">II", flg, exp) + seq_no_bytes +
                               struct.pack(">QI", cas, force))
                    else:
                        ext = struct.pack(">IIQQI", flg, exp, 1, cas, force)
            else:
                ext = struct.pack(">IIQQI", flg, exp, 1, cas, force)
            if conf_res:
                extra_meta = struct.pack(">BBHH",
                                         couchbaseConstants.DCP_EXTRA_META_VERSION,
                                         couchbaseConstants.DCP_EXTRA_META_CONFLICT_RESOLUTION,
                                         len(conf_res),
                                         conf_res)
                ext += struct.pack(">H", len(extra_meta))
        elif (cmd == couchbaseConstants.CMD_SET or
              cmd == couchbaseConstants.CMD_ADD):
            ext = struct.pack(couchbaseConstants.SET_PKT_FMT, flg, exp)
        elif (cmd == couchbaseConstants.CMD_DELETE or
              cmd == couchbaseConstants.CMD_GET or
              cmd == couchbaseConstants.CMD_NOOP):
            ext = b''
        else:
            return f'error: MCSink - unknown cmd for request: {cmd!s}', empty_tuple

        hdr = self.cmd_header(cmd, vbucket_id, key, val, ext, 0, opaque, dtype)
        return 0, (hdr, ext, key, val, ext_meta)

    def cmd_header(self, cmd: int, vbucket_id: int, key: bytes, val: bytes, ext: bytes, cas: int, opaque: int,
                   dtype: int = 0,
                   fmt: str = couchbaseConstants.REQ_PKT_FMT,
                   magic: int = couchbaseConstants.REQ_MAGIC_BYTE) -> bytes:

        return struct.pack(fmt, magic, cmd,
                           len(key), len(ext), dtype, vbucket_id,
                           len(key) + len(ext) + len(val), opaque, cas)

    def read_conn(self, conn: cb_bin_client.MemcachedClient) -> Tuple[int, int, bytes, bytes, bytes, int, int]:
        ext = b''
        key = b''
        val = b''

        buf, cmd, errcode, extlen, keylen, data, cas, opaque = self.recv_msg(conn.s, getattr(conn, 'buf', b''))
        conn.buf = buf  # type: ignore

        if data:
            ext = data[0:extlen]
            key = data[extlen:extlen + keylen]
            val = data[extlen + keylen:]

        return cmd, errcode, ext, key, val, cas, opaque

    def recv_msg(self, sock: socket.socket, buf: bytes) -> Tuple[bytes, int, int, int, int, bytes, int, int]:
        pkt, buf = self.recv(sock, couchbaseConstants.MIN_RECV_PACKET, buf)
        if len(pkt) == 0:
            raise EOFError()
        magic, cmd, keylen, extlen, dtype, errcode, datalen, opaque, cas = \
            struct.unpack(couchbaseConstants.RES_PKT_FMT, pkt)  # type: int, int, int, int, int, int, int, int, int
        if magic != couchbaseConstants.RES_MAGIC_BYTE:
            raise Exception(f'unexpected recv_msg magic: {magic!s}')
        data, buf = self.recv(sock, datalen, buf)
        return buf, cmd, errcode, extlen, keylen, data, cas, opaque

    def recv(self, skt: socket.socket, nbytes: int, buf: bytes) -> Tuple[bytes, bytes]:
        while len(buf) < nbytes:
            data: bytes = b''
            try:
                data = skt.recv(max(nbytes - len(buf), 4096))
            except socket.timeout:
                logging.error("error: recv socket.timeout")
            except Exception as e:
                logging.error(f'error: recv exception: {e!s}')

            if data == b'':
                return b'', b''
            buf += data

        return buf[:nbytes], buf[nbytes:]
import ast
import csv
import json
import os
import sqlite3
import struct
import tempfile
import time
import unittest
import zipfile
import snappy
import couchbaseConstants as cbcs
from collections import defaultdict
from mock_binary_server import MockMemcachedServer
from mock_server import MockRESTServer
from cb_bin_client import MemcachedClient
from pump import Batch, filter_bucket_nodes
from pump_bfd import BFD, CBB_VERSION, DDOC_FILE_NAME, FTS_FILE_NAME, INDEX_FILE_NAME, BFDSource
from pump_bfd2 import BFDSinkEx
from pump_csv import CSVSink, CSVSource
from pump_dcp import DCPStreamSource
from pump_gen import GenSource
from pump_json import JSONSource
from pump_mc import MCSink
            f1 = open(os.path.join(tmpdirname, 'design_docs'), 'w')
            f1.write('design doc')
            f1.close()
            rv, des = self.source.provide_design(None, 'json://' + tmpdirname, None, None)
            self.assertEqual(rv, 0)
            self.assertListEqual(des, [b'design doc'])
            os.remove(os.path.join(tmpdirname, 'design_docs'))

        # Test scenario where no design doc in the dir
        with tempfile.TemporaryDirectory() as tmpdirname:
            rv, des = self.source.provide_design(None, 'json://' + tmpdirname, None, None)
            self.assertEqual(rv, 0)
            self.assertListEqual(des, [])

    def test_provide_design_zip(self):
        # Test scenario where design is in a zip structure should be :
        with tempfile.TemporaryDirectory() as tmpdirname:
            f1 = open(os.path.join(tmpdirname, 'design_docs'), 'w')
            f1.write('design doc')
            fname = f1.name
            f1.close()

            zipF = zipfile.ZipFile(os.path.join(tmpdirname, 'test_zip.zip'), 'w')
            zipF.write(fname, 'design_docs' + os.path.sep + 'design_doc', zipfile.ZIP_DEFLATED)
            zipF.close()

            rv, des = self.source.provide_design(None, 'json://' + os.path.join(tmpdirname, 'test_zip.zip'), None, None)
            self.assertEqual(rv, 0)
            self.assertListEqual(des, [b'design doc'])
            os.remove(fname)
            os.remove(os.path.join(tmpdirname, 'test_zip.zip'))

    def test_provide_batch_dir(self):
        with tempfile.TemporaryDirectory() as tmpdirname:
            # create directory
            os.mkdir(os.path.join(tmpdirname, 'docs'))
            # create json data
            f1 = open(os.path.join(tmpdirname, 'docs', 'doc0.json'), 'w')
            json.dump({'string': 'value', 'number': 1}, f1)
            f1.close()

            f2 = open(os.path.join(tmpdirname, 'docs', 'doc1'), 'w')
            json.dump({'bool': True, 'list': [0, 1, 2]}, f2)
            f2.close()

            self.source = JSONSource(self.opts, 'json://' + tmpdirname, None, None, None, None, None, None)
            rv, batch = self.source.provide_batch()
            self.assertEqual(rv, 0)
            self.assertEqual(batch.size(), 2)
            expected_msgs = [
                (cbcs.CMD_TAP_MUTATION, 0x0000ffff, 'doc0', 0x02000006, 0, 0, '',
                 json.dumps({'string': 'value', 'number': 1}), 0, 0, 0, 0),
                (cbcs.CMD_TAP_MUTATION, 0x0000ffff, 'doc1', 0x02000006, 0, 0, '',
                 json.dumps({'bool': True, 'list': [0, 1, 2]}), 0, 0, 0, 0)
            ]
            self.assertEqual(batch.msgs.sort(key=lambda tup: tup[2]), expected_msgs.sort(key=lambda tup: tup[2]))

    def test_provide_batch_zip(self):
        with tempfile.TemporaryDirectory() as tmpdirname:
            # create json data
            f1 = open(os.path.join(tmpdirname, 'doc0.json'), 'w')
            json.dump({'string': 'value', 'number': 1}, f1)
            f1.close()

            f2 = open(os.path.join(tmpdirname, 'doc1'), 'w')
            json.dump({'bool': True, 'list': [0, 1, 2]}, f2)
            f2.close()

            # create zip
            zipF = zipfile.ZipFile(os.path.join(tmpdirname, 'test_zip.zip'), 'w')
            zipF.write(os.path.join(tmpdirname, 'doc0.json'), 'docs' + os.path.sep + 'doc0.json', zipfile.ZIP_DEFLATED)
            zipF.write(os.path.join(tmpdirname, 'doc1'), 'docs' + os.path.sep + 'doc1', zipfile.ZIP_DEFLATED)
            zipF.close()

            self.source = JSONSource(self.opts, 'json://' + os.path.join(tmpdirname, 'test_zip.zip'), None, None, None,
                                     None, None, None)
            rv, batch = self.source.provide_batch()
            self.assertEqual(rv, 0)
            self.assertEqual(batch.size(), 2)
            expected_msgs = [
                (cbcs.CMD_TAP_MUTATION, 0x0000ffff, 'doc0', 0x02000006, 0, 0, '',
                 json.dumps({'string': 'value', 'number': 1}), 0, 0, 0, 0),
                (cbcs.CMD_TAP_MUTATION, 0x0000ffff, 'doc1', 0x02000006, 0, 0, '',
                 json.dumps({'bool': True, 'list': [0, 1, 2]}), 0, 0, 0, 0)
            ]
            self.assertEqual(batch.msgs.sort(key=lambda tup: tup[2]), expected_msgs.sort(key=lambda tup: tup[2]))


class TestGenSource(unittest.TestCase):
    def setUp(self):
        self.source_map = {'cfg':
                           {'cur-ops': 0,
                            'cur-gets': 0,
                            'cur-sets': 0,
                            'cur-items': 0,
                            'exit-after-creates': True,
                            'max-items': 5,
                            'min-value-size': 10,
                            'prefix': "",
                            'ratio-sets': 1,
                            'json': True,
                            'low-compression': False,
                            'xattr': False}
                           }
        self.opts = Ditto({'collection': False, 'extra': {'batch_max_size': 100, 'batch_max_bytes': 40000}})
        self.source = GenSource(self.opts, None, None, None, self.source_map, None, None, None)

    def test_can_handle(self):
        self.assertFalse(self.source.can_handle(None, 'http://localhost:9000'))
        self.assertTrue(self.source.can_handle(None, 'gen:'))

    def test_provide_batch_json_all_sets_stop_after_creation(self):
        rv, batch = self.source.provide_batch()
        self.assertEqual(rv, 0)
        self.assertEqual(batch.size(), 5)
        for v in batch.msgs:
            self.assertEqual(v[0], cbcs.CMD_DCP_MUTATION)

    def test_provide_batch_json_all_sets_stop_after_creation_xattr(self):
        self.source_map['cfg']['xattr'] = True

        self.source = GenSource(self.opts, None, None, None, self.source_map, None, None, None)