Refactore report.py. Simpler progress reports posting

Author: michael-kotliar

cwl_airflow/utilities/report.py

@@ -2,7 +2,6 @@
 import json
 import jwt
 import logging
-import requests
 
 from airflow.models import Variable
 from airflow.utils.state import State
@@ -11,31 +10,22 @@
 
 CONN_ID = "process_report"
 ROUTES = {
-    "progress": "progress",
-    "results":  "results",
-    "status":   "status"
+    "progress": "airflow/progress",
+    "results":  "airflow/results",
+    "status":   "airflow/status"
 }
 PRIVATE_KEY = "process_report_private_key"
 ALGORITHM = "process_report_algorithm"
 
+http_hook = HttpHook(method="POST", http_conn_id=CONN_ID)  # won't fail even if CONN_ID doesn't exist
 
-def prepare_connection(conn_id, route):
-    http_hook = HttpHook(http_conn_id=conn_id)
-    session = http_hook.get_conn()
-    url = "/".join(
-        [
-            u.strip("/") for u in [http_hook.base_url, session.headers["endpoint"], route]
-        ]
-    )
-    return http_hook, session, url
 
-
-def sign_with_jwt(data, private_key=None, algorithm=None):
+def sign_with_jwt(data):
     try:
         data = jwt.encode(
             payload=data,
-            key=private_key or Variable.get(PRIVATE_KEY),
-            algorithm=algorithm or Variable.get(ALGORITHM)
+            key=Variable.get(PRIVATE_KEY),
+            algorithm=Variable.get(ALGORITHM)
         ).decode("utf-8")
     except Exception as err:
         logging.debug(f"Failed to sign data with JWT key. \n {err}")
@@ -45,27 +35,19 @@ def sign_with_jwt(data, private_key=None, algorithm=None):
 def post_progress(context, from_task=None):
     from_task = False if from_task is None else from_task
     try:
-        http_hook, session, url = prepare_connection(CONN_ID, ROUTES["progress"])
         dag_run = context["dag_run"]
         len_tis = len(dag_run.get_task_instances())
         len_tis_success = len(dag_run.get_task_instances(state=State.SUCCESS)) + int(from_task)
         data = sign_with_jwt(
-            data={
-                "state":    dag_run.state,
-                "dag_id":   dag_run.dag_id,
-                "run_id":   dag_run.run_id,
+            {
+                "state": dag_run.state,
+                "dag_id": dag_run.dag_id,
+                "run_id": dag_run.run_id,
                 "progress": int(len_tis_success / len_tis * 100),
-                "error":    context["reason"] if dag_run.state == State.FAILED else ""
+                "error": context["reason"] if dag_run.state == State.FAILED else ""
             }
         )
-        prepped_request = session.prepare_request(
-            requests.Request(
-                "POST",
-                url,
-                json={"payload": data}
-            )
-        )
-        http_hook.run_and_check(session, prepped_request, {})
+        http_hook.run(endpoint=ROUTES["progress"], json={"payload": data})
     except Exception as err:
         logging.debug(f"Failed to POST progress updates. \n {err}")
 
@@ -76,11 +58,10 @@ def post_results(context):
     isinstance(task, CWLJobGatherer) to find the proper task because of the
     endless import loop (file where we define CWLJobGatherer class import this
     file). If CWLDAG is contsructed with custom gatherer node, posting results
-    might not work.
+    might not work. We need to except missing results file as the same callback
+    is used for clean_dag_run DAG
     """
-    
     try:
-        http_hook, session, url = prepare_connection(CONN_ID, ROUTES["results"])
         dag_run = context["dag_run"]
         results = {}
         try:
@@ -89,47 +70,31 @@ def post_results(context):
                 results = json.load(input_stream)
         except Exception as err:
             logging.debug(f"Failed to read results. \n {err}")
-
         data = sign_with_jwt(
-            data={
-                "dag_id":  dag_run.dag_id,
-                "run_id":  dag_run.run_id,
+            {
+                "dag_id": dag_run.dag_id,
+                "run_id": dag_run.run_id,
                 "results": results
             }
         )
-        prepped_request = session.prepare_request(
-            requests.Request(
-                "POST",
-                url,
-                json={"payload": data}
-            )
-        )
-        http_hook.run_and_check(session, prepped_request, {})
+        http_hook.run(endpoint=ROUTES["results"], json={"payload": data})
     except Exception as err:
         logging.debug(f"Failed to POST results. \n {err}")
 
 
 def post_status(context):
     try:
-        http_hook, session, url = prepare_connection(CONN_ID, ROUTES["status"])
         dag_run = context["dag_run"]
         ti = context["ti"]
         data = sign_with_jwt(
-            data={
-                "state":    ti.state,
-                "dag_id":   dag_run.dag_id,
-                "run_id":   dag_run.run_id,
-                "task_id":  ti.task_id
+            {
+                "state": ti.state,
+                "dag_id": dag_run.dag_id,
+                "run_id": dag_run.run_id,
+                "task_id": ti.task_id
             }
         )
-        prepped_request = session.prepare_request(
-            requests.Request(
-                "POST",
-                url,
-                json={"payload": data}
-            )
-        )
-        http_hook.run_and_check(session, prepped_request, {})
+        http_hook.run(endpoint=ROUTES["status"], json={"payload": data})
     except Exception as err:
         logging.debug(f"Failed to POST status updates. \n {err}")