cbworkloadgen

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

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

import pump_transfer
import pump
import pump_gen
import sys
import optparse
import platform
import os

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\n" \
            "  %prog -n 10.3.121.192:8091 -r .9 -i 100000 \\\n" \
            "         -s 100 -b my-other-bucket --threads=10\n\n"

    def main(self, argv):

        src, sink, common_opts, count_opts, bool_opts = self.opt_construct(argv)

        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"
                   }

        sink_opts = {"node" : "http://"}

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

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

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

        bool_opts = {"ssl": ["-s", None],
                     "force_txn": ["--force-txn", 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]
        sink_str = ""
        for key in sink_opts.keys():
            val = getattr(opts, key, None)
            if val:
                sink_str += sink_opts[key] + val

        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

        return gen_str, sink_str, common_opts, count_opts, bool_opts

    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("-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",
                     help=optparse.SUPPRESS_HELP)
        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))