class Context(object): def __init__(self): # Flags self.verbose = False self.debug = False self.insecure_ssl = False self.json = False self.flags_set = False # Paths self.home = os.getcwd() self.config_file = None # JSON Config (optional for dcos) self.config = None # Conditional options self.framework = 'riak' self.cluster = 'default' self.node = 'riak-default-1' self.timeout = 60 # RiakMesosClient self.client = None def cli_error(self, message): raise CliError(message) def _init_flags(self, verbose, debug, info, version, config_schema, json, insecure_ssl, **kwargs): # Exit immediately if any of these are found if self.flags_set: return args = sys.argv[1:] if info or '--info' in args: click.echo('Start and manage Riak nodes in Mesos.') exit(0) if version or '--version' in args: click.echo('Riak Mesos Framework Version ' + constants.version) exit(0) if config_schema or '--config-schema' in args: click.echo('{}') exit(0) # Process remaining flags for all future command invocations if self.verbose or '--verbose' in args or '-v' in args: self.verbose = True else: self.verbose = verbose if self.insecure_ssl or '--insecure-ssl' in args: self.insecure_ssl = True else: self.insecure_ssl = insecure_ssl self.json = True if self.json or '--json' in args else json self.debug = True if self.debug or '--debug' in args else debug # Configure logging for 3rd party libs if self.debug: logging.basicConfig(level=0) self.verbose = True elif self.verbose: logging.basicConfig(level=20) else: logging.basicConfig(level=50) self.flags_set = True self.vlog("Insecure SSL Mode: " + str(self.insecure_ssl)) self.vlog("Verbose Mode: " + str(self.verbose)) self.vlog("Debug Mode: " + str(self.debug)) self.vlog("JSON Mode: " + str(self.json)) def init_args(self, home, config, framework, cluster, node, **kwargs): self._init_flags(**kwargs) if home is not None: self.home = home if self.config is None or config is not None: if config is not None: self.config_file = config else: usr_conf_file = self.home + '/.config/riak-mesos/config.json' sys_conf_file = '/etc/riak-mesos/config.json' usr_home = expanduser("~") usr_home_conf_file = \ usr_home + '/.config/riak-mesos/config.json' if os.path.isfile(usr_conf_file): self.config_file = usr_conf_file elif os.path.isfile(usr_home_conf_file): self.config_file = usr_home_conf_file elif os.path.isfile(sys_conf_file): self.config_file = sys_conf_file else: self.config_file = None if self.config_file is not None: self.vlog('Using config file: ' + self.config_file) else: self.vlog('Couldn\'t find config file') self.config = RiakMesosConfig(self.config_file) if cluster is not None: self.cluster = cluster if node is not None: self.node = node if framework is not None: self.framework = framework _framework = self.config.get('framework-name') if framework is None and _framework != '': self.framework = _framework if 'timeout' in kwargs and kwargs['timeout'] is not None: self.timeout = kwargs['timeout'] def log(self, msg, *args): """Logs a message to stderr.""" if args: msg %= args click.echo(msg, file=sys.stderr) def vlog(self, msg, *args): """Logs a message to stderr only if verbose is enabled.""" if self.verbose: self.log(msg, *args) def vlog_request(self, r): """Logs request info to stderr only if verbose is enabled.""" if self.debug: self.vlog('HTTP URL: ' + r.url) self.vlog('HTTP Method: ' + r.request.method) self.vlog('HTTP Body: ' + str(r.request.body)) self.vlog('HTTP Status: ' + str(r.status_code)) self.vlog('HTTP Response Text: ' + r.text) def vtraceback(self): if self.verbose: traceback.print_exc() def _init_client(self): ctx = self if self.config_file is None: try: _client = RiakMesosClient(ctx, RiakMesosDCOSStrategy) self.client = _client return except Exception as e: self.vlog(e.message) _client = RiakMesosClient(ctx) self.client = _client return def get_framework_url(self): if self.client is None: self._init_client() return self.client.framework_url() def api_request(self, method, path, exit_on_failure=True, **kwargs): return self.framework_request(method, 'api/v1/' + path, exit_on_failure, **kwargs) def framework_request(self, method, path, exit_on_failure=True, **kwargs): if self.client is None: self._init_client() try: framework_url = self.client.framework_url() return self.http_request(method, framework_url + path, exit_on_failure, **kwargs) except Exception as e: if exit_on_failure: raise e else: self.vlog(e) return FailedRequest(0, method, 'framework_url_not_available/' + path) def master_request(self, method, path, exit_on_failure=True, **kwargs): if self.client is None: self._init_client() try: master_url = self.client.master_url() return self.http_request(method, master_url + path, exit_on_failure, **kwargs) except Exception as e: if exit_on_failure: raise e else: self.vlog(e) return FailedRequest(0, method, 'master_url_not_available/' + path) def node_request(self, method, node, path, exit_on_failure=True, **kwargs): return self.framework_request(method, 'riak/nodes/' + node + '/' + path, exit_on_failure, **kwargs) def marathon_client(self): if self.client is None: self._init_client() marathon_url = self.client.marathon_url() return marathon.Client(marathon_url) def zk_command(self, command, path): if self.client is None: self._init_client() zk_url = self.client.zk_url() try: zk = KazooClient(hosts=zk_url) zk.start() res = False if command == 'get': data, stat = zk.get(path) res = data.decode("utf-8") elif command == 'delete': zk.delete(path, recursive=True) res = 'Successfully deleted ' + path zk.stop() return res except Exception as e: self.vlog(e) return False def http_request(self, method, url, exit_on_failure=True, **kwargs): try: verify = True if self.insecure_ssl: verify = False r = http.request(method, url, verify=verify, is_success=_default_is_success, **kwargs) self.vlog_request(r) if r.status_code == 404: return FailedRequest( 404, method, url, 'Resource at ' + url + ' was not found (Status Code: 404)') return r except Exception as e: if exit_on_failure: raise e else: self.vlog(e) return FailedRequest(0, method, url)
class Context(object): def __init__(self): # Flags self.verbose = False self.debug = False self.insecure_ssl = False self.json = False self.flags_set = False # Paths self.home = os.getcwd() self.config_file = None # JSON Config (optional for dcos) self.config = None # Conditional options self.framework = "riak" self.cluster = "default" self.node = "riak-default-1" self.timeout = 60 # RiakMesosClient self.client = None def cli_error(self, message): raise CliError(message) def _init_flags(self, verbose, debug, info, version, config_schema, json, insecure_ssl, **kwargs): # Exit immediately if any of these are found if self.flags_set: return args = sys.argv[1:] if info or "--info" in args: click.echo("Start and manage Riak nodes in Mesos.") exit(0) if version or "--version" in args: click.echo("Riak Mesos Framework Version " + constants.version) exit(0) if config_schema or "--config-schema" in args: click.echo("{}") exit(0) # Process remaining flags for all future command invocations if self.verbose or "--verbose" in args or "-v" in args: self.verbose = True else: self.verbose = verbose if self.insecure_ssl or "--insecure-ssl" in args: self.insecure_ssl = True else: self.insecure_ssl = insecure_ssl self.json = True if self.json or "--json" in args else json self.debug = True if self.debug or "--debug" in args else debug # Configure logging for 3rd party libs if self.debug: logging.basicConfig(level=0) self.verbose = True elif self.verbose: logging.basicConfig(level=20) else: logging.basicConfig(level=50) self.flags_set = True self.vlog("Insecure SSL Mode: " + str(self.insecure_ssl)) self.vlog("Verbose Mode: " + str(self.verbose)) self.vlog("Debug Mode: " + str(self.debug)) self.vlog("JSON Mode: " + str(self.json)) def init_args(self, home, config, framework, cluster, node, **kwargs): self._init_flags(**kwargs) if home is not None: self.home = home if self.config is None or config is not None: if config is not None: self.config_file = config else: usr_conf_file = self.home + "/.config/riak-mesos/config.json" sys_conf_file = "/etc/riak-mesos/config.json" usr_home = expanduser("~") usr_home_conf_file = usr_home + "/.config/riak-mesos/config.json" if os.path.isfile(usr_conf_file): self.config_file = usr_conf_file elif os.path.isfile(usr_home_conf_file): self.config_file = usr_home_conf_file elif os.path.isfile(sys_conf_file): self.config_file = sys_conf_file else: self.config_file = None if self.config_file is not None: self.vlog("Using config file: " + self.config_file) else: self.vlog("Couldn't find config file") self.config = RiakMesosConfig(self.config_file) if cluster is not None: self.cluster = cluster if node is not None: self.node = node if framework is not None: self.framework = framework _framework = self.config.get("framework-name") if self.framework is None and _framework != "": self.framework = _framework if "timeout" in kwargs and kwargs["timeout"] is not None: self.timeout = kwargs["timeout"] def log(self, msg, *args): """Logs a message to stderr.""" if args: msg %= args click.echo(msg, file=sys.stderr) def vlog(self, msg, *args): """Logs a message to stderr only if verbose is enabled.""" if self.verbose: self.log(msg, *args) def vlog_request(self, r): """Logs request info to stderr only if verbose is enabled.""" if self.debug: self.vlog("HTTP URL: " + r.url) self.vlog("HTTP Method: " + r.request.method) self.vlog("HTTP Body: " + str(r.request.body)) self.vlog("HTTP Status: " + str(r.status_code)) self.vlog("HTTP Response Text: " + r.text) def vtraceback(self): if self.verbose: traceback.print_exc() def _init_client(self): ctx = self if self.config_file is None: try: _client = RiakMesosClient(ctx, RiakMesosDCOSStrategy) self.client = _client return except Exception as e: self.vlog(e.message) _client = RiakMesosClient(ctx) self.client = _client return def get_framework_url(self): if self.client is None: self._init_client() return self.client.framework_url() def api_request(self, method, path, exit_on_failure=True, **kwargs): return self.framework_request(method, "api/v1/" + path, exit_on_failure, **kwargs) def framework_request(self, method, path, exit_on_failure=True, **kwargs): if self.client is None: self._init_client() try: framework_url = self.client.framework_url() return self.http_request(method, framework_url + path, exit_on_failure, **kwargs) except Exception as e: if exit_on_failure: raise e else: self.vlog(e) return FailedRequest(0, method, "framework_url_not_available/" + path) def master_request(self, method, path, exit_on_failure=True, **kwargs): if self.client is None: self._init_client() try: master_url = self.client.master_url() return self.http_request(method, master_url + path, exit_on_failure, **kwargs) except Exception as e: if exit_on_failure: raise e else: self.vlog(e) return FailedRequest(0, method, "master_url_not_available/" + path) def node_request(self, method, node, path, exit_on_failure=True, **kwargs): return self.framework_request(method, "riak/nodes/" + node + "/" + path, exit_on_failure, **kwargs) def marathon_client(self): if self.client is None: self._init_client() marathon_url = self.client.marathon_url() return marathon.Client(marathon_url) def zk_command(self, command, path): if self.client is None: self._init_client() zk_url = self.client.zk_url() try: zk = KazooClient(hosts=zk_url) zk.start() res = False if command == "get": data, stat = zk.get(path) res = data.decode("utf-8") elif command == "delete": zk.delete(path, recursive=True) res = "Successfully deleted " + path zk.stop() return res except Exception as e: self.vlog(e) return False def http_request(self, method, url, exit_on_failure=True, **kwargs): try: verify = True if self.insecure_ssl: verify = False r = http.request(method, url, verify=verify, is_success=_default_is_success, **kwargs) self.vlog_request(r) if r.status_code == 404: return FailedRequest(404, method, url, "Resource at " + url + " was not found (Status Code: 404)") return r except Exception as e: if exit_on_failure: raise e else: self.vlog(e) return FailedRequest(0, method, url)