rethinkdb-python/rethinkdb/utils_common.py at master · rethinkdb/rethinkdb-python

476 lines (386 loc) · 15.5 KB
# Copyright 2018 RethinkDB
# Licensed under the Apache License, Version 2.0 (the 'License');
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# This file incorporates work covered by the following copyright:
# Copyright 2010-2016 RethinkDB, all rights reserved.
from __future__ import print_function
import collections
import copy
from looseversion import LooseVersion
import getpass
import inspect
import optparse
import threading
from rethinkdb import ast, errors, net, query, version
default_batch_size = 200
class RetryQuery(object):
    __connectOptions = None
    __local = None
    def __init__(self, connect_options):
        if "host" not in connect_options:
            raise AssertionError("Hostname is a required connection parameter")
        if "port" not in connect_options:
            raise AssertionError("Port number is a required connection parameter")
        connect_options["port"] = int(connect_options["port"])
        if connect_options["port"] <= 0:
            raise AssertionError("Port number can not be less than one")
        self.__connectOptions = copy.deepcopy(connect_options)
        self.__local = threading.local()
    def conn(self, test_connection=True):
        if not hasattr(self.__local, "connCache"):
            self.__local.connCache = {}
        # check if existing connection is still good
        if os.getpid() in self.__local.connCache and test_connection:
            try:
                ast.expr(0).run(self.__local.connCache[os.getpid()])
            except errors.ReqlError:
                del self.__local.connCache[os.getpid()]
        # cache a new connection
        if not os.getpid() in self.__local.connCache:
            self.__local.connCache[os.getpid()] = net.make_connection(
                net.DefaultConnection, **self.__connectOptions
        # return the connection
        return self.__local.connCache[os.getpid()]
    def __call__(
        self, name, query_str, times=5, run_options=None, test_connection=True
        # Try a query multiple times to guard against bad connections
        if name is None:
            raise AssertionError("Name can not be none")
        name = str(name)
        if not isinstance(query_str, ast.RqlQuery):
            raise AssertionError(
                "Query must be a ReQL query instead of {value}".format(value=query_str)
        if not isinstance(times, int) or times < 1:
            raise ValueError(
                "Times must be a positive integer instead of {value}".format(
                    value=times
        if run_options is None:
            run_options = {}
        if not isinstance(run_options, dict):
            raise ValueError(
                "Run option must be a dict instead of {value}".format(value=run_options)
        last_error = None
        test_connection = False
        for _ in range(times):
            try:
                conn = self.conn(
                    test_connection=test_connection
                )  # we are already guarding for this
            except errors.ReqlError as e:
                last_error = RuntimeError(
                    "Error connecting for during '%s': %s" % (name, str(e))
                test_connection = True
            try:
                return query_str.run(conn, **run_options)
            except (errors.ReqlTimeoutError, errors.ReqlDriverError) as e:
                last_error = RuntimeError(
                    "Connection error during '%s': %s" % (name, str(e))
            # other errors immediately bubble up
        if last_error is not None:
            raise last_error
def print_progress(ratio, indent=0, read=None, write=None):
    total_width = 40
    done_width = min(int(ratio * total_width), total_width)
    sys.stdout.write(
        "\r%(indent)s[%(done)s%(undone)s] %(percent)3d%%%(readRate)s%(writeRate)s\x1b[K"
            "indent": " " * indent,
            "done": "=" * done_width,
            "undone": " " * (total_width - done_width),
            "percent": int(100 * ratio),
            "readRate": (" r: %d" % read) if read is not None else "",
            "writeRate": (" w: %d" % write) if write is not None else "",
    sys.stdout.flush()
def check_minimum_version(options, minimum_version="1.6", raise_exception=True):
    minimum_version = LooseVersion(minimum_version)
    version_string = options.retryQuery(
        "get server version",
        query.db("rethinkdb").table("server_status")[0]["process"]["version"],
    matches = re.match(
        r"(rethinkdb|rebirthdb) (?P<version>(\d+)\.(\d+)\.(\d+)).*", version_string
    if not matches:
        raise RuntimeError("invalid version string format: %s" % version_string)
    if LooseVersion(matches.group("version")) < minimum_version:
        if raise_exception:
            raise RuntimeError(
                "Incompatible version, expected >= %s got: %s"
                % (minimum_version, version_string)
        return False
    return True
DbTable = collections.namedtuple("DbTable", ["db", "table"])
_tableNameRegex = re.compile(r"^(?P<db>[\w-]+)(\.(?P<table>[\w-]+))?$")
class CommonOptionsParser(optparse.OptionParser, object):
    __retryQuery = None
    __connectRegex = re.compile(r"^\s*(?P<hostname>[\w.-]+)(:(?P<port>\d+))?\s*$")
    def format_epilog(self, formatter):
        return self.epilog or ""
    def __init__(self, *args, **kwargs):
        # -- Type Checkers
        def check_tls_option(_, opt_str, value):
            value = str(value)
            if os.path.isfile(value):
                return {"ca_certs": os.path.realpath(value)}
            else:
                raise optparse.OptionValueError(
                    "Option %s value is not a file: %r" % (opt_str, value)
        def check_db_table_option(_, _opt_str, value):
            res = _tableNameRegex.match(value)
            if not res:
                raise optparse.OptionValueError(
                    "Invalid db or db.table name: %s" % value
            if res.group("db") == "rethinkdb":
                raise optparse.OptionValueError(
                    "The `rethinkdb` database is special and cannot be used here"
            return DbTable(res.group("db"), res.group("table"))
        def check_positive_int(_, opt_str, value):
            try:
                value = int(value)
                if value < 1:
                    raise ValueError
            except ValueError:
                raise optparse.OptionValueError(
                    "%s value must be an integer greater than 1: %s" % (opt_str, value)
            return value
        def check_existing_file(_, opt_str, value):
            if not os.path.isfile(value):
                raise optparse.OptionValueError(
                    "%s value was not an existing file: %s" % (opt_str, value)
            return os.path.realpath(value)
        def check_new_file_location(_, opt_str, value):
            try:
                real_value = os.path.realpath(value)
            except Exception:
                raise optparse.OptionValueError(
                    "Incorrect value for %s: %s" % (opt_str, value)
            if os.path.exists(real_value):
                raise optparse.OptionValueError(
                    "%s value already exists: %s" % (opt_str, value)
            return real_value
        def file_contents(_, opt_str, value):
            if not os.path.isfile(value):
                raise optparse.OptionValueError(
                    "%s value is not an existing file: %r" % (opt_str, value)
            try:
                with open(value, "r") as passwordFile:
                    return passwordFile.read().rstrip("\n")
            except IOError:
                raise optparse.OptionValueError(
                    "bad value for %s: %s" % (opt_str, value)
        # -- Callbacks
        def combined_connect_action(obj, opt, value, parser, *args, **kwargs):
            """optparse.takeaction() calls the callback (which this is set as)
               with the following args: self, opt, value, parser *args, **kwargs
            res = self.__connectRegex.match(value)
            if not res:
                raise optparse.OptionValueError(
                    "Invalid 'host:port' format: %s" % value
            if res.group("hostname"):
                parser.values.hostname = res.group("hostname")
            if res.group("port"):
                parser.values.driver_port = int(res.group("port"))
        # -- setup custom Options object
        class CommonOptionChecker(optparse.Option, object):
            TYPES = optparse.Option.TYPES + (
                "tls_cert",
                "db_table",
                "pos_int",
                "file",
                "new_file",
                "file_contents",
            TYPE_CHECKER = copy.copy(optparse.Option.TYPE_CHECKER)
            TYPE_CHECKER["tls_cert"] = check_tls_option
            TYPE_CHECKER["db_table"] = check_db_table_option
            TYPE_CHECKER["pos_int"] = check_positive_int
            TYPE_CHECKER["file"] = check_existing_file
            TYPE_CHECKER["new_file"] = check_new_file_location
            TYPE_CHECKER["file_contents"] = file_contents
            ACTIONS = optparse.Option.ACTIONS + ("add_key", "get_password")
            STORE_ACTIONS = optparse.Option.STORE_ACTIONS + ("add_key", "get_password")
            TYPED_ACTIONS = optparse.Option.TYPED_ACTIONS + ("add_key",)
            ALWAYS_TYPED_ACTIONS = optparse.Option.ALWAYS_TYPED_ACTIONS + ("add_key",)
            def take_action(self, action, dest, opt, value, values, parser):
                if dest is None:
                    raise AssertionError("Destination can not be none")
                if action == "add_key":
                    if self.metavar is None:
                        raise AssertionError("Metavar can not be none")
                    values.ensure_value(dest, {})[self.metavar.lower()] = value
                elif action == "get_password":
                    values.ensure_value('password', getpass.getpass("Password for `admin`: "))
                else:
                    super(CommonOptionChecker, self).take_action(
                        action, dest, opt, value, values, parser
        kwargs["option_class"] = CommonOptionChecker
        # - default description to the module's __doc__
        if "description" not in kwargs:
            # get calling module
            caller = inspect.getmodule(inspect.stack()[1][0])
            if caller.__doc__:
                kwargs["description"] = caller.__doc__
        # -- add version
        if "version" not in kwargs:
            kwargs["version"] = "%%prog %s" % version.VERSION
        # -- call super
        super(CommonOptionsParser, self).__init__(*args, **kwargs)
        # -- add common options
        self.add_option(
            "-q",
            "--quiet",
            dest="quiet",
            default=False,
            action="store_true",
            help="suppress non-error messages",
        self.add_option(
            "--debug",
            dest="debug",
            default=False,
            action="store_true",
            help=optparse.SUPPRESS_HELP,
        connection_group = optparse.OptionGroup(self, "Connection options")
        connection_group.add_option(
            "-c",
            "--connect",
            dest="driver_port",
            metavar="HOST:PORT",
            help="host and client port of a rethinkdb node to connect (default: localhost:%d)"
            % net.DEFAULT_PORT,
            action="callback",
            callback=combined_connect_action,
            type="str",
        connection_group.add_option(
            "--driver-port",
            dest="driver_port",
            metavar="PORT",
            help="driver port of a rethinkdb server",
            type="int",
            default=os.environ.get("RETHINKDB_DRIVER_PORT", net.DEFAULT_PORT),
        connection_group.add_option(
            "--host-name",
            dest="hostname",
            metavar="HOST",
            help="host and driver port of a rethinkdb server",
            default=os.environ.get("RETHINKDB_HOSTNAME", "localhost"),
        connection_group.add_option(
            "-u",
            "--user",
            dest="user",
            metavar="USERNAME",
            help="user name to connect as",
            default=os.environ.get("RETHINKDB_USER", "admin"),
        connection_group.add_option(
            "-p",
            "--password",
            dest="password",
            help="interactively prompt for a password for the connection",
            action="get_password",
        connection_group.add_option(
            "--password-file",
            dest="password",
            metavar="PSWD_FILE",
            help="read the connection password from a file",
            type="file_contents",
        connection_group.add_option(
            "--tls-cert",
            dest="ssl",
            metavar="TLS_CERT",
            help="certificate file to use for TLS encryption.",
            type="tls_cert",
        self.add_option_group(connection_group)
    def parse_args(self, *args, **kwargs):
        # - validate options
        connect = True
        if "connect" in kwargs:
            connect = kwargs["connect"]
            del kwargs["connect"]
        # - validate ENV variables
        if "RETHINKDB_DRIVER_PORT" in os.environ:
            driver_port = os.environ["RETHINKDB_DRIVER_PORT"]
            if not isinstance(driver_port, int) or driver_port < 1:
                self.error(
                    "ENV variable RETHINKDB_DRIVER_PORT is not a useable "
                    "integer: %s" % os.environ["RETHINKDB_DRIVER_PORT"]
        # - parse options
        options, args = super(CommonOptionsParser, self).parse_args(*args, **kwargs)
        # - setup a RetryQuery instance
        self.__retryQuery = RetryQuery(
            connect_options={
                "host": options.hostname,
                "port": options.driver_port,
                "user": options.user,
                "password": options.password,
                "ssl": options.ssl,
        options.retryQuery = self.__retryQuery
        # - test connection
        if connect:
            try:
                options.retryQuery.conn()
            except errors.ReqlError as e:
                self.error("Unable to connect to server: %s" % str(e))
        return options, args
_interrupt_seen = False
def abort(pools, exit_event):
    global _interrupt_seen
    if _interrupt_seen:
        # second time
        print("\nSecond terminate signal seen, aborting ungracefully")
        for pool in pools:
            for worker in pool:
                worker.terminate()
                worker.join(0.1)
        print("\nTerminate signal seen, aborting")
        _interrupt_seen = True
        exit_event.set()
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

FilesExpand file tree

utils_common.py

Latest commit

History

utils_common.py

File metadata and controls
