cbrecovery

#!/usr/bin/env python3
# -*-python-*-

"""
Recover missing vbuckets from remote cluster due to failed rebalancing operation
"""

import copy
import json
import logging
import os
import platform
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from collections import defaultdict

import pump
import pump_cb
import pump_dcp
import pump_transfer


class Recovery(pump_transfer.Transfer):
    """Entry point for 2.0 cbrecovery."""

    def __init__(self):
        self.name = "cbrecovery"
        self.source_alias = "cluster to recover from"
        self.sink_alias = "cluster to recover to"
        self.recovery_map = None
        self.recovery_uuid = None
        self.sink_bucket = None
        self.source_bucket = None
        self.stats = defaultdict(int)
        self.start_time = time.time()

        self.usage = \
            "%prog [options] source destination\n\n" \
            "Recover missing vbuckets from remote cluster .\n\n" \
            "Examples:\n" \
            "  %prog [options] http://SOURCE:8091 http://DEST:8091"

        self.err_map = {
            "startRecovery": {
                "400": "Could not start recovery",
                "404": "Specified bucket not found",
                "500": "Could not start recovery because some nodes failed",
                "503": "Could not start recovery because rebalance is running"
            },
            "commitVBucket": {
                "400": "recovery_uuid query parameter is missing",
                "402": "Vbucket is either unspecified or could not be converted",
                "404": "Specified vbucket is not part of recovery map",
                "406": "Either no recovery in progress or wrong recovery_uuid provided",
                "500": "Could not commit vbucket basuse some nodes failed"
            }
        }

        self.bad_request_code_err_map = {
            "not_needed": "No recovery needed",
            "recovery_running": "Recovery is already running against one of the buckets on the destination cluster," \
                "only one can run at a time",
        }

    def main(self, argv):

        err, opts, source, sink = self.opt_parse(argv)
        if err:
            return err
        self.sink_bucket = getattr(opts, "bucket_destination", "default")
        self.source_bucket = getattr(opts, "bucket_source", "default")

        # communicate with destination before data transferring
        err, cmd, bad_request_code = self.pre_transfer(opts, source, sink)
        if err:
            return self.explain_error(err, cmd, bad_request_code)
        if self.recovery_map:
            sys.stderr.write(f'Missing vbuckets to be recovered:{json.dumps(self.recovery_map)}\n')
        else:
            sys.stderr.write("No missing vbuckets to be recovered.\n")
            return 0

        if opts.dry_run:
            err, cmd = self.post_transfer(opts, source, sink, None)
            return self.explain_error(err, cmd)

        # only care about data recovery
        argv.append("--extra")
        argv.append("data_only=1,conflict_resolve=0,allow_recovery_vb_remap=1")

        for _, node in enumerate(self.recovery_map):
            for vbucket in node["vbuckets"]:
                if opts.id is None or int(opts.id) == int(vbucket):
                    node_map = {node["node"]: [vbucket]}
                    temp_argv = copy.deepcopy(argv)
                    temp_argv.append("--vbucket-list")
                    temp_argv.append(json.dumps(node_map))
                    err = pump_transfer.Transfer.main(self, temp_argv)
                    if err:
                        return err
                    self.post_transfer(opts, source, sink, vbucket)

        self.report_stats()
        return 0

    def aggregate_stats(self, cur):
        for k, val in cur.items():
            self.stats[k] += val
        return 0

    def report_stats(self):
        delta = time.time() - self.start_time
        x = sorted([k for k in self.stats.keys() if "_sink_" in k])

        width_k = max([5] + [len(k.replace("tot_sink_", "")) for k in x] + [len("Recovery")])
        width_v = max([20] + [len(str(self.stats[k])) for k in x])
        width_s = max([10] + [len("%0.1f" % (self.stats[k] / delta)) for k in x])
        sys.stderr.write(f'\n {"Recovery".ljust(width_k)} : {"Total".rjust(width_v)} | {"Per sec".rjust(width_s)}\n')
        for k in x:
            per_sec = f'{self.stats[k]/delta:0.1f}'
            sys.stderr.write(f' {k.replace("tot_sink_", ""):<{width_k}} : {self.stats[k]!s:>{width_v}} |'
                             f' {per_sec:>{width_s}}\n')

        total_vbucket = 0
        if self.recovery_map:
            for k, node in enumerate(self.recovery_map):
                total_vbucket += len(node["vbuckets"])
        sys.stderr.write(f'{total_vbucket} vbuckets recovered with elapsed time {delta:0.2f} seconds\n')

    def pre_transfer(self, opts, source, sink):
        host, port, _, _, path = pump.parse_spec(opts, sink, 8091)

        # NOTE: Request is going to the destination cluster e.g. the one being recovered. Use the correct credentials.
        user, pwd = opts.username_dest, opts.password_dest

        # retrieve a list of missing vbucket
        cmd = "startRecovery"
        url = f'/pools/default/buckets/{self.sink_bucket}/controller/{cmd}'
        err, conn, response = pump.rest_request(host, int(port), user, pwd, opts.ssl, url, method='POST',
                                                reason='start_recovery')

        if err:
            if response != b'':
                response = json.loads(response)
                bad_request_code = response.get('code', None)
                if bad_request_code:
                    return err, cmd, bad_request_code

            return err, cmd, None

        response = json.loads(response)
        recovery_map = response.get('recoveryMap', None)
        if not recovery_map:
            return "Missing recovery map from response", None, None

        self.recovery_map = recovery_map
        self.recovery_uuid = response.get('uuid', None)

        return None, None, None

    def post_transfer(self, opts, source, sink, vbucket):
        if not self.sink_bucket:
            return "Should specify destination bucket for restore", None

        if opts.dry_run:
            cmd = "stopRecovery"
            reason_msg = "Stop recovery"
        else:
            cmd = "commitVBucket"
            reason_msg = f'Activate vbucket:{vbucket}'
        if vbucket is not None:
            params = urllib.parse.urlencode({"vbucket": vbucket})
        else:
            params = None

        host, port, _, _, path = pump.parse_spec(opts, sink, 8091)

        # NOTE: Request is going to the destination cluster e.g. the one being recovered. Use the correct credentials.
        user, pwd = opts.username_dest, opts.password_dest

        url = f'/pools/default/buckets/{self.sink_bucket}/controller/{cmd}?recovery_uuid={self.recovery_uuid}'
        post_headers = {"Content-type": "application/x-www-form-urlencoded"}
        err, _, _ = pump.rest_request(host, int(port), user, pwd, False, url, method='POST', body=params,
                                                reason=reason_msg,
                                                headers=post_headers)
        if err:
            return err, cmd
        if vbucket:
            logging.debug(f'Recover and activate vbucket: {vbucket}')
        return None, None

    def opt_parser_options(self, p):

        p.add_option("-b", "--bucket-source",
                     action="store", type="string", default="default",
                     metavar="default",
                     help="""source bucket to recover from """)
        p.add_option("-B", "--bucket-destination",
                     action="store", type="string", default="default",
                     metavar="default",
                     help="""destination bucket to recover to """)
        self.opt_parser_options_common(p)

        p.remove_option("-k")

    def find_handlers(self, opts, source, sink):
        return pump_dcp.DCPStreamSource, pump_cb.CBSink

    def explain_error(self, err, cmd, bad_request_code=None):
        if cmd not in self.err_map:
            logging.error(err)
            return
        pattern = r';\sresponse:\s(.+);\sreason'
        m = re.search(pattern, err)
        if m:
            errcode = m.groups()[0]
            if errcode in self.err_map[cmd]:
                if errcode == "400":
                    if bad_request_code in self.bad_request_code_err_map:
                        logging.error(f'Error: {self.bad_request_code_err_map[bad_request_code]}')
                    else:
                        logging.error(f'Error: could not start recovery for unknown reason: {bad_request_code}')
                else:
                    logging.error(f'Error: {self.err_map[cmd][errcode]}')
            else:
                logging.error(f'Error: unknown error code:{errcode}')
        else:
            logging.error(err)


if __name__ == '__main__':
    if platform.system() == "Windows":
        python_lib = os.path.join(os.path.dirname(sys.argv[0]), '..')
        sys.path.append(python_lib)

    pump_transfer.exit_handler(Recovery().main(sys.argv))