cbworkloadgen

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

"""
Workload generator tool for Couchbase 2.0.
"""

import optparse
import os
import platform
import sys

import pump
import pump_gen
import pump_transfer
import couchbaseConstants

from urllib.parse import urlparse


class OptConstructError(Exception):
    """Class of validation exceptions that can be thrown by opt_construct function"""
    pass

class WorkloadGen(pump_transfer.Transfer):
    """Entry point for 2.0 cbworkloadgen."""

    def __init__(self):
        self.name = "cbworkloadgen"
        self.source_alias = "generator"
        self.sink_alias = "destination"

        self.usage = \
            "%prog [options]\n\n" \
            "Generate workload to destination.\n\n" \
            "Examples:\n" \
            "  %prog -n localhost:8091 -u Administrator -p password\n" \
            "  %prog -n 10.3.121.192:8091 -r .9 -i 100000 -u Administrator -p password \\\n" \
            "         -s 100 -b my-other-bucket --threads=10\n\n"

    def main(self, argv):
        try:
            src, sink, common_opts, count_opts, bool_opts = self.opt_construct(argv)
        except OptConstructError as e:
            sys.exit(f"{self.name}: error: {e}")

        local_args = [argv[0]]
        local_args.append(src)
        local_args.append(sink)
        for v in common_opts.values():
            local_args.append(v[0])
            local_args.append(v[1])

        for v in count_opts.values():
            if v[1] is not None:
                for i in range(v[1]):
                    local_args.append(v[0])

        for v in bool_opts.values():
            if v[1]:
                local_args.append(v[0])
        return pump_transfer.Transfer.main(self, local_args)

    def opt_construct(self, argv):
        gen_opts = {"ratio_sets": "ratio-sets",
                    "loop" : "exit-after-creates",
                    "max_items" : "max-items",
                    "size" : "min-value-size",
                    "prefix" : "prefix",
                    "json" : "json",
                    "low_compression" : "low-compression",
                    "xattr":"xattr"
                   }

        common_opts = {"bucket" : ["-B", None],
                       "threads" : ["-t", None],
                       "username" : ["-u", None],
                       "password" : ["-p", None],
                       "extra" : ["-x", "backoff_cap=0.1"],
                       "collection" : ["-c", None],
                       "cacert" : ["--cacert", None]
                      }

        list_opts =   {"collection" : ["-c", None]
                      }

        count_opts = {"verbose" : ["-v", None]}

        bool_opts = {"ssl": ["-s", None],
                     "force_txn": ["--force-txn", None],
                     "no_ssl_verify": ["--no-ssl-verify", None]
                    }

        p = optparse.OptionParser(usage=self.usage)
        self.add_parser_options(p)
        opts, rest = p.parse_args(argv[1:])

        # optimisation to avoid having to check the xattrs for txn when loading the data using cbworkloadgen
        setattr(opts, 'force_txn', True)

        gen_str = "gen:"
        for key in gen_opts.keys():
            val = getattr(opts, key, None)
            if val is not None:
                if key in ("loop", "json"):
                    val = int(val)
                gen_str += f'{gen_opts[key]}={val!s},'
        if gen_str[-1] == ",":
            gen_str = gen_str[:-1]

        for key in common_opts.keys():
            val = getattr(opts, key, None)
            if val:
                common_opts[key][1] = str(val)

        for key in list_opts.keys():
            val = getattr(opts, key, None)
            if val:
                common_opts[key][1] =  list(val)

        for key in count_opts.keys():
            val = getattr(opts, key, None)
            if val:
                count_opts[key][1] = int(val)

        for key in bool_opts.keys():
            val = getattr(opts, key, None)
            if val:
                bool_opts[key][1] = True

        sink_str = ""
        val = getattr(opts, 'node', None)
        if val:
            sink_str = self.construct_sink(val, bool_opts['ssl'][1])

        return gen_str, sink_str, common_opts, count_opts, bool_opts

    def construct_sink(self, sink_str, ssl):
        if sink_str == "":
            # This error case should not be reachable since the flag parser should not allow empty values for the
            # --node/-n flag, added for consistency and completeness
            raise OptConstructError("the ip:port of a node seems to not be supplied, please check your CLI command")

        sink_str_split = sink_str.split('://')

        if len(sink_str_split) > 2:
            raise OptConstructError("the supplied ip:port of a node has incorrect format, check if a scheme is " \
                "supplied more than once")

        if len(sink_str_split) == 1:
            sink_str = f"http{'s' if ssl else ''}://" + sink_str
            self.validate_urlparse(urlparse(sink_str), ssl)
        else:
            sink_str = sink_str.replace("couchbase://", "http://")
            sink_str = sink_str.replace("couchbases://", "https://")

            sink_str_parsed = urlparse(sink_str)
            scheme = sink_str_parsed.scheme

            if scheme == "https":
                ssl = True
            elif scheme == "http":
                if ssl:
                    raise OptConstructError("the ip:port of a node is supplied with 'http' scheme, which is not " \
                        "supported with the --ssl flag, please use 'https' or remove the scheme entirely")
            else:
                raise OptConstructError(f"the ip:port of a node is supplied with an unsupported scheme '{scheme}' " \
                    "(only 'http', 'https', 'couchbase' and 'couchbases' are supported)")

            self.validate_urlparse(sink_str_parsed, ssl)

        return sink_str

    def validate_urlparse(self, urlparse_result, ssl):
        try:
            port = urlparse_result.port
        except Exception as e:
            raise OptConstructError(f"failed to get the port from the supplied ip:port of a node: {e}")

        self.validate_port(port, ssl)

    def validate_port(self, port, ssl):
        if port is None:
            raise OptConstructError("the ip:port of a node is supplied without a port")
        elif (ssl and port != couchbaseConstants.SSL_REST_PORT):
                raise OptConstructError(f"the ip:port of a node is supplied with a bad port '{port}', you " \
                    f"must use '{couchbaseConstants.SSL_REST_PORT}' if you are using SSL encryption")

    def add_parser_options(self, p):
        p.add_option("-r", "--ratio-sets",
                     action="store", type="float", default=.95,
                     metavar=".95",
                     help="""set/get operation ratio""")
        p.add_option("-n", "--node", type="string", default="127.0.0.1:8091",
                     metavar="127.0.0.1:8091",
                     help="node's ns_server ip:port")
        p.add_option("-b", "--bucket",
                     action="store", type="string", default="default",
                     metavar="default",
                     help="""insert data to a different bucket other than default """)
        p.add_option("--ssl",
                     action="store_true", default=False,
                     help="Transfer data with SSL enabled")
        p.add_option("--no-ssl-verify", default=False, action="store_true",
                     help="Skips SSL verification of certificates against the CA")
        p.add_option("--cacert", default=None, action="store",
                     help="Verifies the cluster identity with this certificate")
        p.add_option("-i", "--max-items",
                     action="store", type="int", default=10000,
                     metavar="10000",
                     help="""number of items to be inserted""")
        p.add_option("-s", "--size",
                     action="store", type="int", default=10,
                     metavar="10",
                     help="""minimum value size""")
        p.add_option("--prefix",
                     action="store", type="string", default="pymc",
                     metavar="pymc",
                     help="""prefix to use for memcached keys or json ids""")
        p.add_option("-j", "--json",
                     action="store_true", default=False,
                     help="""insert json data""")
        p.add_option("-l", "--loop",
                     action="store_false", default=True,
                     help="""loop forever until interrupted by users""")
        p.add_option("-u", "--username",
                     action="store", type="string", default=None,
                     help="REST username for cluster or server node")
        p.add_option("-p", "--password",
                     action="store", type="string", default=None,
                     help="REST password for cluster or server node")
        p.add_option("-t", "--threads",
                     action="store", type="int", default=1,
                     metavar="1",
                     help="""number of concurrent workers""")
        p.add_option("-v", "--verbose",
                     action="count", default=0,
                     help="verbose logging; more -v's provide more verbosity")
        p.add_option("--low-compression",
                     action="store_true", default=False,
                     help="generate document data that is difficult to compress")
        p.add_option("-c", "--collection", type="string", action="append", metavar="id",
                     help='The collection ID, in hex format, to do the operations on')
        p.add_option('--xattr', action="store_true", default=False,
                     help='generate extended attributes for inserted documents')
        return p

    def find_handlers(self, opts, source, sink):
        return pump_gen.GenSource, pump.PumpingStation.find_handler(opts, sink, pump_transfer.SINKS)


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(WorkloadGen().main(sys.argv))