diff --git a/update_sqlite.py b/update_sqlite.py new file mode 100755 index 0000000..37019c3 --- /dev/null +++ b/update_sqlite.py @@ -0,0 +1,98 @@ +#! /usr/bin/env python3 +import sys +import os +import logging +import sqlite3 +from logging.handlers import SysLogHandler +from elasticsearch import Elasticsearch + +log_level = logging.INFO + +host = "https://scitastaloa.epfl.ch:9200" +api_key = ["XXXXXXXXXXXXXXXXXXXX", "XXXXXXXXXXXXXXXXXXXXXX"] +indexes = { + "slurm_pricing_cost": ".slurm_pricing_cost", + "slurm_account_pricing": ".slurm_account_pricing" +} + +pivot_idx = 'pricing' + +sqlite_db = '.slurm.db' +sql_table = { + 'name': 'data', + 'fields': { + 'account': 'text', + 'cluster': 'text', + 'value': 'real' + } +} + + +def create_data(): + data = [] + es = Elasticsearch(hosts=host, timeout=10, api_key=api_key) + temp_a = {} + index = indexes["slurm_pricing_cost"] + resp = es.search(index=index, + query={"match_all": {}}, size=1000) + for hit in resp['hits']['hits']: + source = hit['_source'] + pivot = source[pivot_idx] + temp_b = [] + for cluster, value in source['clusters'].items(): + temp_b.append([cluster, value]) + temp_a[pivot] = temp_b + + index = indexes["slurm_account_pricing"] + resp = es.search(index=index, + query={"match_all": {}}, size=1000) + for hit in resp['hits']['hits']: + source = hit['_source'] + pivot = source[pivot_idx] + for cluster, value in temp_a[pivot]: + data.append((source['account'], cluster, value)) + + return data + + +def sqlite_init(cur): + sql = "DROP TABLE %s" % sql_table['name'] + try: + cur.execute(sql) + except sqlite3.OperationalError: + pass + col = ["%s %s" % (k, v) for k, v in sql_table['fields'].items()] + sql = "CREATE TABLE %s (%s)" % (sql_table['name'], ",".join(col)) + cur.execute(sql) + + +def insert_data(): + data = create_data() + con = sqlite3.connect(sqlite_db) + cur = con.cursor() + sqlite_init(cur) + cur.executemany(f"INSERT INTO {sql_table['name']} VALUES(?, ?, ?)", data) + con.commit() + con.close() + + +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) + + insert_data() + + logging.info("Done")