Skip to content

Commit

Permalink
Refactore report.py. Simpler progress reports posting
Browse files Browse the repository at this point in the history
  • Loading branch information
michael-kotliar committed Nov 6, 2020
1 parent 9dd3125 commit b62c9c7
Showing 1 changed file with 25 additions and 60 deletions.
85 changes: 25 additions & 60 deletions cwl_airflow/utilities/report.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
import json
import jwt
import logging
import requests

from airflow.models import Variable
from airflow.utils.state import State
Expand All @@ -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}")
Expand All @@ -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}")

Expand All @@ -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:
Expand All @@ -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}")

Expand Down

0 comments on commit b62c9c7

Please sign in to comment.