added caching

conf/config.py.sample

@@ -34,6 +34,8 @@ job_uid_comm = {
 
 # DB
 #############################
+CACHEJSON = False # cache aggregated results as json file
+CACHEDIR = "/path/to/caching/directory"
 
 job_info = {
     "fetch_job_info_lsf": "bjobs -o \"{:s}\" {:s}",

data.py

@@ -7,20 +7,59 @@ import conf.config as conf
 import subprocess
 import re
 import os
+import gzip
+import io
 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)
+def cache_file(jobid, filename, content):
+    if conf.CACHEJSON is False or conf.CACHEDIR == "":
+        return False
 
-    formatted["recommendations"] = rcm
+    if not os.path.exists(conf.CACHEDIR):
+        os.makedirs(conf.CACHEDIR, mode=0o770)
+
+    cachefile = conf.CACHEDIR + "/" + filename
+
+    fd = os.open(cachefile, os.O_CREAT | os.O_WRONLY, 0o660)
+    f = gzip.GzipFile(fileobj=io.FileIO(fd, mode='wb'))
+    f.write(content.encode())
+    f.close()
+
+    return True
+
+def cache_data(jobid, type, content):
+    if conf.CACHEJSON is False or conf.CACHEDIR == "":
+        return False
+
+    filename = ""
+    if type == "text":
+        filename = "{:s}.txt.json.gz".format(jobid)
+    elif type == "pdf":
+        filename = "{:s}.pdf.json.gz".format(jobid)
+    else:
+        raise RuntimeError("invalid type of report: {}".format(type))
+
+    cache_file(jobid, filename, content)
+
+
+def merge_and_out(job_id, aggr, rcm, type, cache, out_dir=None):
+    if cache[type] is False:
+        formatted = format.format(aggr, type)
+        formatted["recommendations"] = rcm
+        jsonstr = json.dumps(formatted)
+
+        if conf.CACHEJSON is True:
+            cache_data(job_id, type, jsonstr)
+    else:
+        jsonstr = cache[type]
 
     if out_dir is None:
-        print(json.dumps(formatted))
+        print(jsonstr)
     else:
         filename = "{:s}/{:s}.{:s}.json".format(out_dir, job_id, type)
         with open(filename, 'w') as outfile:
-            json.dump(formatted, outfile)
+            outfile.write(jsonstr)
 
         print("{:s} data was exported in {:s}".format(type, filename))
 
@@ -56,6 +95,37 @@ def check_user(jobid):
 
     return
 
+def check_cache_file(jobid, filename):
+    if conf.CACHEJSON is False or conf.CACHEDIR == "":
+        return False
+
+    if not os.path.exists(conf.CACHEDIR):
+        os.makedirs(conf.CACHEDIR, mode=0o770)
+
+    cachefile = conf.CACHEDIR + "/" + filename
+
+    res = False
+    if os.path.exists(cachefile):
+        f = gzip.open(cachefile, "r")
+        content = f.read()
+        if content != "init" and content != "":
+            res = content
+    else:
+        fd = os.open(cachefile, os.O_CREAT | os.O_WRONLY, 0o660)
+        f = gzip.GzipFile(fileobj=io.FileIO(fd, mode='wb'))
+        f.write("init".encode())
+        f.close()
+
+    return res
+
+def check_cache_txt(jobid):
+    txtfilename = "{:s}.txt.json.gz".format(jobid)
+    return check_cache_file(jobid, txtfilename)
+
+def check_cache_pdf(jobid):
+    txtfilename = "{:s}.pdf.json.gz".format(jobid)
+    return check_cache_file(jobid, txtfilename)
+
 def main():
     parser = argparse.ArgumentParser(description="""
         Gets the job information required for generating text or PDF reports
@@ -85,18 +155,37 @@ def main():
             exit()
     # End of errors in arguments
 
-    aggr = aggregator.get_aggregator(job_id, "pdf")
+    # initial values
+    need_to_aggregate = True
+    cache = {"text": False, "pdf": False}
+    aggr = None
+    recommendations = None
 
-    recommendations = rcm.get_recommendations(aggr)
+    # Check if we can use cached data
+    if conf.CACHEJSON is True:
+        if args.type == "all":
+            cache["text"] = check_cache_txt(job_id)
+            cache["pdf"] = check_cache_pdf(job_id)
+            need_to_aggregate = (cache["text"] is False) or (cache["pdf"] is False)
+        elif args.type == "text":
+            cache["text"] = check_cache_txt(job_id)
+            need_to_aggregate = cache["text"] is False
+        elif args.type == "pdf":
+            cache["pdf"] = check_cache_pdf(job_id)
+            need_to_aggregate = cache["pdf"] is False
+
+    if need_to_aggregate is True:
+        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)
+        merge_and_out(job_id, aggr, recommendations, args.type, cache)
     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)
+            merge_and_out(job_id, aggr, recommendations, "text", cache, args.output_dir)
+            merge_and_out(job_id, aggr, recommendations, "pdf", cache, args.output_dir)
         else:
-            merge_and_out(job_id, aggr, recommendations, args.type, args.output_dir)
+            merge_and_out(job_id, aggr, recommendations, args.type, cache, args.output_dir)
 
     return 0