Fix issue with cross-join in DB in function.

Fix issue with xcom value is not returned as an object
robinhood · Jun 29, 2021 · 08c650b · 08c650b
1 parent 078a151
commit 08c650b
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/airflow_prometheus_exporter/prometheus_exporter.py b/airflow_prometheus_exporter/prometheus_exporter.py
@@ -121,8 +121,8 @@ def get_dag_duration_info():
                 ),
             )
             .filter(
-                TaskInstance.start_date.isnot(None),
-                TaskInstance.end_date.isnot(None),
+                dag_start_dt_query.c.start_date.isnot(None),
+                DagRun.end_date.isnot(None),
             )
             .all()
         )
@@ -446,7 +446,8 @@ def collect(self):
         for tasks in xcom_config.get("xcom_params", []):
             for param in get_xcom_params(tasks["task_id"]):
                 xcom_value = extract_xcom_parameter(param.value)
-
+                if not isinstance(xcom_value, dict):
+                    xcom_value = {xcom_value: xcom_value}
                 if tasks["key"] in xcom_value:
                     xcom_params.add_metric(
                         [param.dag_id, param.task_id], xcom_value[tasks["key"]]