diff --git a/update_es.py b/update_es.py index 1f15dde..fb14e08 100644 --- a/update_es.py +++ b/update_es.py @@ -1,110 +1,114 @@ #! /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", "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 generated with this ACL : +{ + "viewer": { + "cluster": [ + "monitor", + "manage_enrich" + ], + "indices": [ + { + "names": [ + ".slurm_account_pricing" + ], + "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" +index = ".slurm_account_pricing" +policy = "slurm_account_pricing_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 + # determine account pricing 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 + if "pricing" in v: + doc["pricing"] = v["pricing"] + elif "parent" in v: + doc["pricing"] = accounts[v["parent"]]['pricing'] 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: + resp["_source"] = None + pass + try: + if resp["_source"] != account: 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})") - + 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")