diff --git a/update_es.py b/update_es.py index 057655c..1f15dde 100644 --- a/update_es.py +++ b/update_es.py @@ -1,110 +1,110 @@ #! /usr/bin/env python3 import sys import os import logging import yaml from logging.handlers import SysLogHandler from elasticsearch import Elasticsearch, exceptions as es_e, client as es_client log_level = logging.INFO filepath = "/etc/accounts.yaml" -to_remove = ["clusters", "free", "premium", "root", "courses"] +to_remove = ["clusters", "root"] free = ["free", "courses", "root"] host = "https://scitastaloa.epfl.ch:9200" # API key generated with this ACL : # { # "viewer": { # "cluster": [ # "monitor", # "manage_enrich" # ], # "indices": [ # { # "names": [ # ".slurm_account_type" # ], # "privileges": [ # "read", # "write", # "view_index_metadata" # ], # "allow_restricted_indices": false # } # ], # "run_as": [ # "*" # ] # } # } api_key = ["XXXXXXXXXXXXXXXXXXXX", "XXXXXXXXXXXXXXXXXXXXXX"] index = ".slurm_account_type" - +policy = "slurm_account_type_enrich" def get_accounts(): docs = [] with open(filepath, encoding="utf-8") as f: accounts = yaml.safe_load(f) # remove unecessaries items for tr in to_remove: accounts.pop(tr) # determine account type based on parent for k, v in accounts.items(): doc = {} if "groups" not in v: next doc["account"] = k.lower() a_type = v["parent"] if "parent" in v else "commitment" if a_type == "root": a_type = k.lower() doc["type"] = a_type docs.append(doc) return docs def set_es_index(accounts): es = Elasticsearch(hosts=host, timeout=10, api_key=api_key) for account in accounts: name = account["account"] try: resp = es.get(index=index, id=name) except es_e.NotFoundError: logging.info(f"[{name}] missing") # print(resp["_source"]) if resp["_source"] != account: try: es.index(index=index, id=name, document=account) logging.info(f"[{name}] created or updated") except es_e.AuthorizationException as e: logging.warning(f"[{name}] update or creation impossible ({e})") if __name__ == "__main__": assert sys.version_info >= (3, 6) pid = os.getpid() level = log_level fmt = f"[%(levelname)s] ES {pid} %(name)s - %(message)s" # create a syslog handler handler = SysLogHandler(address="/dev/log") handler.setLevel(level) handler.setFormatter(logging.Formatter(fmt)) # add syslog handler to root logging.basicConfig(level=level, format=fmt, handlers=[handler]) logging.getLogger("elasticsearch").setLevel(logging.WARNING) logging.getLogger("elastic_transport").setLevel(logging.WARNING) logging.info(f"Begin updating '{index}' index based on accounts.yaml file") # create list account based on accounts.yaml file accounts = get_accounts() # update ES index es = Elasticsearch(hosts=host, timeout=10, api_key=api_key) set_es_index(es, accounts) # execute associated policy es_client.EnrichClient(es).execute_policy(name=policy) logging.info(f"Done")