Homec4science
Diffusion sausage-binary (master) Browse

sausage-binary/update_es.py
586dce4c4a6emaster

Last Change
Raw File

update_es.py
View Options

#! /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
log_level = logging.INFO
filepath = "/etc/accounts.yaml"
to_remove = ["clusters", "free", "premium", "root", "courses"]
free = ["free", "courses", "root"]
host = "https://scitastaloa.epfl.ch:9200"
# API key generated with this ACL :
# {
# "viewer": {
# "cluster": [
# "monitor"
# ],
# "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"
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
a_type = v["parent"] if "parent" in v else "commitment"
if a_type == "root":
a_type = k
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")
accounts = get_accounts()
set_es_index(accounts)
logging.info(f"Done")
c4science · Help