diff --git a/site-modules/profile/files/prometheus/update-prometheus-config b/site-modules/profile/files/prometheus/update-prometheus-config index a0ec8e13..c111608e 100755 --- a/site-modules/profile/files/prometheus/update-prometheus-config +++ b/site-modules/profile/files/prometheus/update-prometheus-config @@ -1,115 +1,116 @@ #!/usr/bin/env python3 # # This generates a static configuration for Prometheus # # Copyright © 2020 The Software Heritage Developers. # This file is released under the Apache-2.0 License. # from collections import defaultdict import copy from dataclasses import asdict, dataclass, fields import datetime import os import stat import sys from typing import Any, Dict, Iterable, List, Optional import yaml @dataclass(frozen=True) class JobGroup: """Job parameters from which to group prometheus jobs""" job_name: str scrape_interval: Optional[int] scrape_timeout: Optional[int] metrics_path: Optional[str] scheme: Optional[str] @classmethod def from_dict(cls, dict): init_vars = {field.name: dict.get(field.name) for field in fields(cls)} + if init_vars.get('metrics_path') == '/metrics': - del init_vars['metrics_path'] + init_vars['metrics_path'] = None if init_vars.get('scheme') == 'http': - del init_vars['scheme'] + init_vars['scheme'] = None return cls(**init_vars) def load_yaml_from_dir(dirname: str) -> Iterable[Dict[str, Any]]: """Load all yaml files from a given directory""" for filename in os.listdir(dirname): if not filename.endswith((".yml", ".yaml")): continue path = os.path.join(dirname, filename) with open(path, "r") as f: yield from yaml.safe_load(f) def generate_scrape_configs(configs: Dict[JobGroup, List[Dict[str, Any]]]): """Generate a scrape_configs entry from a dict""" for params, targets in configs.items(): yield { **{ param: value for param, value in asdict(params).items() if value is not None }, "static_configs": targets, } def merge_prometheus_config( base_config: Dict[str, Any], scrape_configs: Iterable[Dict[str, Any]] ) -> Dict[str, Any]: """Merge the main prometheus config with scrape configs""" config = copy.deepcopy(base_config) config.setdefault("scrape_configs", []).extend(scrape_configs) return config def replace_file(old_file, new_file): """Replace old_file with new_file, ensuring permissions are the same""" try: info = os.stat(old_file) os.chown(new_file, info.st_uid, info.st_gid) os.chmod(new_file, stat.S_IMODE(info.st_mode)) except FileNotFoundError: pass os.rename(new_file, old_file) if __name__ == "__main__": base_conffile = sys.argv[1] exported_dir = sys.argv[2] output = sys.argv[3] config_groups: Dict[JobGroup, List[Dict[str, Any]]] = defaultdict(list) for conf in load_yaml_from_dir(exported_dir): group = JobGroup.from_dict(conf) for key in asdict(group): conf.pop(key, None) config_groups[group].append(conf) with open(base_conffile, "r") as f: base_config = yaml.safe_load(f) full_config = merge_prometheus_config( base_config, generate_scrape_configs(config_groups), ) now = datetime.datetime.now(tz=datetime.timezone.utc).isoformat() with open(output + ".tmp", "w") as f: print(f"# This file was generated by {sys.argv[0]} on {now}.", file=f) print(f"# Changes will be lost", file=f) print(f"", file=f) yaml.dump(full_config, f, default_flow_style=False) replace_file(output, output + ".tmp")