Add more metadata

Author: ividito

dags/veda_data_pipeline/groups/processing_tasks.py

@@ -4,6 +4,7 @@
 from copy import deepcopy
 import smart_open
 from airflow.models.variable import Variable
+from airflow.models.xcom import LazyXComSelectSequence
 from airflow.decorators import task
 from airflow.datasets import Dataset, DatasetAlias
 from airflow.datasets.metadata import Metadata
@@ -88,7 +89,7 @@ def build_stac_task(payload, ti=None):
             DatasetAlias("VEDA-Datasets")
         ],
 )
-def post_ingest_report(ti, logical_date):  # params are Airflow kwargs - use this task without input
+def post_ingest_dataset_event(ti, logical_date, built_items = {}):  # params are Airflow kwargs - use this task without input
     """
     Logs a Dataset event, saving the config used as a versioned object in s3, and creating a Metadata object visible in Airflow.
     
@@ -113,8 +114,22 @@ def post_ingest_report(ti, logical_date):  # params are Airflow kwargs - use thi
         json.dump(payload, f, indent=2)
     log_task(f"Payload written to {key}")
 
+    # built items can be either a dict or a list of dicts
+    if isinstance(built_items, LazyXComSelectSequence):
+        built_items = list(built_items)
+    elif not isinstance(built_items, list):
+        built_items = [built_items]
+    print(f"Built items: {built_items}")
+    success_count = sum(item.get("payload", {}).get("status", {}).get("successes", 0) for item in built_items)
+    failure_count = sum(item.get("payload", {}).get("status", {}).get("failures", 0) for item in built_items)
+
     yield Metadata(
         Dataset(f"{collection}"),
-        extra={"ingest_datetime": str(logical_date) },  # extra has to be provided, can be {}
+        extra={
+            "ingest_datetime": str(logical_date),
+            "ingest_configuration": key,
+            "successful_items": success_count,
+            "failed_items": failure_count,
+        },  # extra has to be provided, can be {}
         alias="VEDA-Datasets",
     )

dags/veda_data_pipeline/veda_dataset_pipeline.py

@@ -4,7 +4,7 @@
 from veda_data_pipeline.groups.discover_group import discover_from_s3_task, get_files_task
 from airflow.operators.empty import EmptyOperator
 from veda_data_pipeline.groups.collection_group import collection_task_group
-from veda_data_pipeline.groups.processing_tasks import submit_to_stac_ingestor_task, build_stac_task, extract_discovery_items_from_payload, remove_thumbnail_asset, post_ingest_report
+from veda_data_pipeline.groups.processing_tasks import submit_to_stac_ingestor_task, build_stac_task, extract_discovery_items_from_payload, remove_thumbnail_asset, post_ingest_dataset_event
 
 template_dag_run_conf = {
     "collection": "<collection-id>",
@@ -53,4 +53,4 @@
     discover = discover_from_s3_task.partial(payload=mutated_payloads).expand(event=discovery_items)
     get_files = get_files_task(payload=discover)
     build_stac = build_stac_task.expand(payload=get_files)
-    submit_stac = submit_to_stac_ingestor_task.expand(built_stac=build_stac) >> post_ingest_report() >> end
+    submit_stac = submit_to_stac_ingestor_task.expand(built_stac=build_stac) >> post_ingest_dataset_event(built_items=build_stac) >> end