feat: Add a decorator function to enable streaming functinality.

2 files changed, 25 insertions, 16 deletions

diff --git a/gn3/api/rqtl.py b/gn3/api/rqtl.py
index 9fd7017..7b29cef 100644
--- a/gn3/api/rqtl.py
+++ b/gn3/api/rqtl.py
@@ -15,14 +15,15 @@ from gn3.computations.rqtl import (
     process_rqtl_pairscan,
     process_perm_output,
 )
-from gn3.computations.streaming import run_process
+from gn3.computations.streaming import run_process, enable_streaming
 from gn3.fs_helpers import assert_path_exists, get_tmpdir
 
 rqtl = Blueprint("rqtl", __name__)
 
 
 @rqtl.route("/compute", methods=["POST"])
-def compute():
+@enable_streaming
+def compute(stream_ouput_file):
     """Given at least a geno_file and pheno_file, generate and
     run the rqtl_wrapper script and return the results as JSON
 
@@ -31,17 +32,6 @@ def compute():
     phenofile = request.form["pheno_file"]
     assert_path_exists(genofile)
     assert_path_exists(phenofile)
-
-    run_id = request.args.get("id")
-    with open(
-        os.path.join(current_app.config.get("TMPDIR"), f"{run_id}.txt"),
-        "w+",
-        encoding="utf-8",
-    ):
-        # TODO thos should  be refactored
-        pass
-    # Split kwargs by those with values and boolean ones
-    # that just convert to True/False
     kwargs = ["covarstruct", "model", "method", "nperm", "scale", "control"]
     boolean_kwargs = ["addcovar", "interval", "pstrata", "pairscan"]
     all_kwargs = kwargs + boolean_kwargs
@@ -75,7 +65,6 @@ def compute():
     )
 
     rqtl_output = {}
-    #  get the stdout file
     run_id = request.args.get("id", str(uuid.uuid4()))
     if not os.path.isfile(
         os.path.join(
@@ -83,8 +72,6 @@ def compute():
         )
     ):
         pass
-    stream_ouput_file = os.path.join(current_app.config.get("TMPDIR"), f"{run_id}.txt")
-
     run_process(rqtl_cmd.get("rqtl_cmd").split(), stream_ouput_file, run_id)
 
     if "pairscan" in rqtl_bool_kwargs:
diff --git a/gn3/computations/streaming.py b/gn3/computations/streaming.py
index d39aa7f..771076c 100644
--- a/gn3/computations/streaming.py
+++ b/gn3/computations/streaming.py
@@ -1,5 +1,8 @@
 """Module contains streaming procedures  for genenetwork. """
+import os
 import subprocess
+from functools import wraps
+from flask import current_app, has_app_context, request
 
 
 def run_process(cmd, output_file, run_id):
@@ -32,3 +35,22 @@ def run_process(cmd, output_file, run_id):
     except subprocess.CalledProcessError as error:
         return {"msg": "error occurred",
                 "error": str(error), "run_id": run_id}
+
+
+def enable_streaming(func):
+    """Decorator function to enable streaming for an endpoint
+    Note: should be used only in an app context
+    """
+    @wraps(func)
+    def decorated_function(*args, **kwargs):
+        if not has_app_context:
+            raise RuntimeError("This decorator must be used within an app context.")
+        run_id = request.args.get("id")
+        stream_ouput_file = os.path.join(current_app.config.get("TMPDIR"),
+                                         f"{run_id}.txt")
+        with open(stream_ouput_file, "w+", encoding="utf-8",
+                  ) as file_handler:
+            file_handler.write("File created for streaming\n"
+                               )
+        return func(stream_ouput_file, *args, **kwargs)
+    return decorated_function