def __init__(self, cluster: str, config_dir: Path, *, dry_run: bool = True) -> None: """Initialize the instance. Arguments: cluster (str): the name of the cluster to connect to. config_dir (str): path to the directory containing the configuration files for the Redis clusters. dry_run (bool, optional): whether this is a DRY-RUN. """ self._dry_run = dry_run self._shards: DefaultDict[str, Dict] = defaultdict(dict) config = load_yaml_config(config_dir / f"{cluster}.yaml") for datacenter, shards in config["shards"].items(): for shard, data in shards.items(): self._shards[datacenter][shard] = RedisInstance( host=data["host"], port=data["port"], password=config["password"], decode_responses=True, )
def reposync(self, name: str) -> RepoSync: """Get a Reposync instance. Arguments: name (str): the name of the repo to sync. Returns: spicerack.reposync.RepoSync: the reposync instance. """ config = load_yaml_config(self.config_dir / "reposync" / "config.yaml") if name not in config["repos"]: raise SpicerackError(f"Unknown repo {name}") repo_dir = Path(config["base_dir"], name) query = ",".join(config["remotes"]) remote_hosts = self.remote().query(query) if not repo_dir.is_dir(): raise SpicerackError( f"The repo directory ({repo_dir}) does not exist") repo = Repo(repo_dir) if not repo.bare: raise SpicerackError( f"The repo directory ({repo_dir}) is not a bare git repository" ) return RepoSync(repo, self.username, remote_hosts, dry_run=self._dry_run)
def test_load_yaml_config_no_raise(caplog, name): """Loading an invalid config with raises=False should return an empty dictionary.""" with caplog.at_level(DEBUG): config_dict = load_yaml_config(get_fixture_path('config', name), raises=False) assert {} == config_dict check_logs(caplog, 'Could not load config file', DEBUG)
def test_offset_transfer_dry_run(self, consumer_patch, _, func_arguments): """It should read but not transfer offsets between consumer groups.""" kafka = Kafka(kafka_config=load_yaml_config( get_fixture_path("kafka", "config.yaml")), dry_run=True) to_consumer_mock = self._setup_consumer_mocks( consumer_patch, _answer_offsets_for_times, _answer_partitions_for_topic) kafka.transfer_consumer_position(*func_arguments[0]) assert not to_consumer_mock.commit.called
def load_services(): """Load the dc-local hostnames for all active-active services.""" # TODO: find a way to use spicerack.config_dir here # It's not easy as we need this when parsing CLI arguments. config_full_path = '/etc/spicerack/cookbooks/sre.switchdc.services.yaml' every_service = load_yaml_config(config_full_path) # Only select active-active services services = {} for srv, data in every_service.items(): if data['active_active']: services[srv] = data['rec'] return services
def kafka(self) -> Kafka: """Get an instance to interact with Kafka. Returns: spicerack.kafka.Kafka: the instance Raises: KeyError: If the configuration file does not contain the correct keys. """ configuration = load_yaml_config(self._spicerack_config_dir / "kafka" / "config.yaml") return Kafka(kafka_config=configuration, dry_run=self._dry_run)
def test_no_topic_partitions(self, consumer_patch, _): """It should read but not transfer offsets between consumer groups.""" kafka = Kafka(kafka_config=load_yaml_config( get_fixture_path("kafka", "config.yaml")), dry_run=False) self._setup_consumer_mocks(consumer_patch, _answer_offsets_for_times, lambda _: None) with pytest.raises( expected_exception=KafkaError, match="Partitions not found for topic eqiad.wikidata."): kafka.transfer_consumer_position( ["wikidata"], ConsumerDefinition("eqiad", "jumbo", "consumer_jumbo_1"), ConsumerDefinition("eqiad", "jumbo", "consumer_jumbo_2"), )
def netbox(self, *, read_write: bool = False) -> Netbox: """Get a Netbox instance to interact with Netbox's API. Arguments: read_write (bool, optional): whether to use a read-write token. Returns: spicerack.netbox.Netbox: the instance """ config = load_yaml_config(self._spicerack_config_dir / "netbox" / "config.yaml") if read_write and not self._dry_run: token = config["api_token_rw"] else: token = config["api_token_ro"] return Netbox(config["api_url"], token, dry_run=self._dry_run)
def ganeti(self) -> Ganeti: """Get an instance to interact with Ganeti. Returns: spicerack.ganeti.Ganeti: the instance Raises: KeyError: If the configuration file does not contain the correct keys. """ configuration = load_yaml_config(self._spicerack_config_dir / "ganeti" / "config.yaml") return Ganeti( configuration["username"], configuration["password"], configuration["timeout"], self.remote(), )
def main(argv: Optional[Sequence[str]] = None) -> Optional[int]: """Entry point, run the tool. Arguments: argv (list, optional): the list of command line arguments to parse. Returns: int: the return code, zero on success, non-zero on failure. """ args = argument_parser().parse_args(argv) if not args.cookbook: args.cookbook = "" config = load_yaml_config(args.config_file) cookbooks_base_dir = Path(config["cookbooks_base_dir"]).expanduser() sys.path.append(str(cookbooks_base_dir)) def get_cookbook( spicerack: Spicerack, cookbook_path: str, cookbook_args: Sequence[str] = () ) -> Optional[BaseItem]: """Run a single cookbook. Arguments: argv (sequence, optional): a sequence of strings of command line arguments to parse. Returns: None: on success. int: the return code, zero on success, non-zero on failure. """ cookbooks = CookbookCollection(cookbooks_base_dir, cookbook_args, spicerack, path_filter=cookbook_path) return cookbooks.get_item(cookbook_path) params = config.get("instance_params", {}) params.update({ "verbose": args.verbose, "dry_run": args.dry_run, "get_cookbook_callback": get_cookbook }) try: spicerack = Spicerack(**params) except TypeError as e: print( "Unable to instantiate Spicerack, check your configuration:", e, file=sys.stderr, ) return 1 cookbooks = CookbookCollection(cookbooks_base_dir, args.cookbook_args, spicerack, path_filter=args.cookbook) if args.list: print(cookbooks.menu.get_tree(), end="") return 0 cookbook_item = cookbooks.get_item(args.cookbook) if cookbook_item is None: logger.error("Unable to find cookbook %s", args.cookbook) return cookbook.NOT_FOUND_RETCODE base_path = Path(config["logs_base_dir"]) / cookbook_item.path.replace( ".", os.sep) _log.setup_logging( base_path, cookbook_item.name, spicerack.username, dry_run=args.dry_run, host=config.get("tcpircbot_host", None), port=int(config.get("tcpircbot_port", 0)), ) logger.debug("Executing cookbook %s with args: %s", args.cookbook, args.cookbook_args) return cookbook_item.run()
def run(args, spicerack): """Required by Spicerack API.""" config_full_path = os.path.join(spicerack.config_dir, CONFIG_PATH) logger.info('Attempting to read secrets from %s', config_full_path) config = load_yaml_config(config_full_path, raises=False) account_id = get_secret('account_id', config) api_token = get_secret('api_token', config) # Configure the session object, reused across requests. session = requests.Session() # Construct the User-Agent per WMF internal policy (a good generic choice). Include the python-requests version. user_agent = ( 'WMF spicerack/sre.network.cf (https://wikitech.wikimedia.org/wiki/Spicerack; [email protected]) ' '{}').format(session.headers['User-Agent']) headers = { 'Content-Type': 'application/json', 'Authorization': 'Bearer {}'.format(api_token), 'User-Agent': user_agent } session.headers.update(headers) session.proxies = spicerack.requests_proxies base_url = CF_BASE_URL.format(account_id) # Get all the configured prefixes response = session.get( '{base_url}/addressing/prefixes'.format(base_url=base_url)) list_prefixes = parse_cf_response('list all prefixes', response) # Store the update action result, so we don't interrupt the run if there is an error # But still return a cookbook failure if at least one update fails. return_code = 0 # Iterate over all prefixes for prefix in list_prefixes['result']: # Only care about the ones we select (or all) if prefix['cidr'] != args.query and not prefix[ 'description'].startswith(args.query) and args.query != 'all': logger.debug( 'Skipping prefix %(cidr)s ("%(description)s"), not matching query', prefix) continue log_prefix_status(prefix) if args.action == 'status': continue advertise = args.action == 'start' if prefix['advertised'] == advertise: logger.info( 'Prefix %(cidr)s ("%(description)s") already in desired state', prefix) continue if spicerack.dry_run: logger.debug( 'Skipping update of prefix %(cidr)s ("%(description)s") in DRY-RUN mode', prefix) continue try: update_prefix_status(session, base_url, prefix, advertise) except Exception as e: # pylint: disable=broad-except # Don't interrupt the run if there is an error logger.error('тЪая╕П Failed to update prefix %s: %s', prefix['cidr'], e) return_code = 1 return return_code
def test_load_yaml_config_valid(): """Loading a valid config should return its content.""" config_dir = get_fixture_path('config', 'valid.yaml') config_dict = load_yaml_config(config_dir) assert 'key' in config_dict
def test_load_yaml_config_raise(name, message): """Loading an invalid config should raise Exception.""" with pytest.raises(WmflibError, match=message): load_yaml_config(get_fixture_path('config', name))
def test_load_yaml_config_empty(): """Loading an empty config should return an empty dictionary.""" config_dict = load_yaml_config(get_fixture_path('config', 'empty.yaml')) assert {} == config_dict
def setup_method(self): """Set up the module under test.""" # pylint: disable=attribute-defined-outside-init self.kafka = Kafka(kafka_config=load_yaml_config( get_fixture_path("kafka", "config.yaml")), dry_run=False)