Enable XCom pickling in Airflow (#421)

* Enable xcom pickling * Update duration reporting
WordPress · Mar 21, 2022 · 0e3675c · 0e3675c
1 parent 22a8965
commit 0e3675c
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 5 deletions.
diff --git a/docker/airflow/Dockerfile b/docker/airflow/Dockerfile
@@ -18,6 +18,8 @@ ENV OUTPUT_DIR=/var/workflow_output/
 ENV AIRFLOW__CORE__DAGS_FOLDER=${AIRFLOW_HOME}/openverse_catalog/dags
 ENV AIRFLOW__CORE__LOAD_EXAMPLES=False
 ENV AIRFLOW__CORE__LOAD_DEFAULT_CONNECTIONS=False
+ENV AIRFLOW__CORE__ENABLE_XCOM_PICKLING=True
+
 # TODO: Test if moving this to .env changes anything!
 ENV AIRFLOW__LOGGING__REMOTE_LOGGING=True
 ENV AIRFLOW__LOGGING__REMOTE_LOG_CONN_ID=aws_default

diff --git a/openverse_catalog/dags/common/dag_factory.py b/openverse_catalog/dags/common/dag_factory.py
@@ -208,6 +208,7 @@ def create_provider_api_workflow(
         catchup=False,
         doc_md=doc_md,
         tags=["provider"] + [f"provider: {media_type}" for media_type in media_types],
+        render_template_as_native_obj=True,
     )
 
     with dag:

diff --git a/openverse_catalog/dags/common/loader/reporting.py b/openverse_catalog/dags/common/loader/reporting.py
@@ -13,15 +13,11 @@ def report_completion(provider_name, media_type, duration, record_count):
     In all cases the data is logged.
     """
 
-    # This happens when the task is manually set to `success` in Airflow before
-    # completing.
-    duration = "_No data_" if duration == "None" else duration
-
     message = f"""
 *Provider*: `{provider_name}`
 *Media Type*: `{media_type}`
 *Number of Records Upserted*: {record_count}
-*Duration of data pull task*: {duration}
+*Duration of data pull task*: {duration or "_No data_"}
 
 * _Duration includes time taken to pull data of all media types._
 """