data.py

#! /usr/bin/env python3

import json
import argparse
import db.aggregator as aggregator
import conf.config as conf
import subprocess
import re
import os
from format import format
from rcm import rcm

def merge_and_out(job_id, aggr, rcm, type, out_dir=None):
    formatted = format.format(aggr, type)

    formatted["recommendations"] = rcm

    if out_dir is None:
        print(json.dumps(formatted))
    else:
        filename = "{:s}/{:s}.{:s}.json".format(out_dir, job_id, type)
        with open(filename, 'w') as outfile:
            json.dump(formatted, outfile)

        print("{:s} data was exported in {:s}".format(type, filename))

def check_user(jobid):
    euid = os.geteuid()
    uid = os.getuid()

    if euid == 0 or uid == 0:
        return

    if conf.BATCH_SYSTEM == "SLURM":
        job_uid_comm = conf.job_uid_comm["slurm"].format(jobid = jobid)
    elif conf.BATCH_SYSTEM == "LSF":
        job_uid_comm = conf.job_uid_comm["lsf"].format(jobid = jobid)
    else:
        print("Cannot check the user ID, no batch system specified")
        exit(1)

    result = subprocess.run(job_uid_comm, stdout=subprocess.PIPE, shell=True,
                            executable='/bin/bash')
    out = result.stdout.decode("utf-8")
    m = re.search("[0-9]+", out)

    if m is None or m.group(0) is None:
        print("Cannot parse UID. {:s}".format(out))
        exit(1)

    job_uid = m.group(0)

    if int(job_uid) != int(uid):
        print("Access denied. UIDs of the user and the job do not match.")
        exit(1)

    return

def main():
    parser = argparse.ArgumentParser(description="""
        Gets the job information required for generating text or PDF reports
        and outputs it in JSON format.
    """, formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument("-t", "--type", help="type of the output",
                        choices=['text', 'pdf', 'all'], default="text")

    parser.add_argument("-o", "--output-dir", help="output directory",
                        const="./out", required=False, nargs='?')

    parser.add_argument("JOBID",
                        help="job ID used in the batch system")

    args = parser.parse_args()

    job_id = args.JOBID

    if conf.SECUSER:
        check_user(job_id)

    # Errors in arguments
    if args.output_dir is None:
        if args.type == "all":
            print("Cannot print both data in STDOUT")
            exit()
    # End of errors in arguments

    aggr = aggregator.get_aggregator(job_id, "pdf")

    recommendations = rcm.get_recommendations(aggr)

    if args.output_dir is None:
        merge_and_out(job_id, aggr, recommendations, args.type)
    else:
        if args.type == "all":
            merge_and_out(job_id, aggr, recommendations, "text", args.output_dir)
            merge_and_out(job_id, aggr, recommendations, "pdf", args.output_dir)
        else:
            merge_and_out(job_id, aggr, recommendations, args.type, args.output_dir)

    return 0

if __name__ == "__main__":
    main()