cbrecovery

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

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

import copy
import logging
import os
import platform
import re
import string
import sys
import time
import urllib.request, urllib.parse, urllib.error

import json

from collections import defaultdict
import pump_transfer
import pump
import pump_dcp
import pump_cb

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": "No recovery needed",
                "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"
            }
        }

    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 = self.pre_transfer(opts, source, sink)
        if err:
            return self.explain_error(err, cmd)
        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 i, 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, user, pwd, path = \
            pump.parse_spec(opts, sink, 8091)

        #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:
            return err, cmd
        response = json.loads(response)
        recoveryMap = response.get('recoveryMap', None)
        if not recoveryMap:
            return "Missing recovery map from response", None
        else:
            self.recovery_map = recoveryMap
        self.recovery_uuid = response.get('uuid', None)

        return None, None

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

        host, port, user, pwd, path = \
            pump.parse_spec(opts, sink, 8091)
        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
        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, conn, response = \
            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
        elif 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):
        if cmd not in self.err_map:
            logging.error(err)
            return
        pattern = ';\sresponse:\s(.+);\sreason'
        m = re.search(pattern, err)
        if m:
            errcode = m.groups()[0]
            if errcode in self.err_map[cmd]:
                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))