How to interpret an XCOM within a task_group to keep a history of dynamically generated tasks within it? - airflow-2.x

Question

I implemented a dag with a task_group that loops on the contents of a file.. This isn't ideal, I wish I could loop through the contents of an XCom, unless I'm mistaken, that isn't possible.
import json
import logging
from airflow.decorators import task_group
from airflow.operators.dummy import DummyOperator
from dtm.migration.helpers.tasks import (
check_compliance_of_objects_sizes,
)
from airflow.providers.google.cloud.transfers.gcs_to_gcs import GCSToGCSOperator
from airflow.providers.google.cloud.sensors.gcs import (
GCSObjectExistenceSensor,
)
#task_group(group_id="copy_task_group")
def copy_and_verify(
inp_parameters_path: str,
impersonated_service_account: str,
):
try:
with open(inp_parameters_path, "r") as f:
inp_parameters = json.load(f)
for entry in inp_parameters:
sensor_source_file = GCSObjectExistenceSensor(
task_id=f"sensor_source_file_{entry}",
bucket=inp_parameters[entry]["source"]["bucket"],
object=(
f"{inp_parameters[entry]['source']['prefix']}"
f"/{inp_parameters[entry]['source']['object']}"
),
impersonation_chain=impersonated_service_account,
)
copy_file = GCSToGCSOperator(
task_id=f"copy_file_{entry}",
source_bucket=inp_parameters[entry]["source"]["bucket"],
source_object=(
f"{inp_parameters[entry]['source']['prefix']}"
f"/{inp_parameters[entry]['source']['object']}"
),
destination_bucket=inp_parameters[entry]["destination"][
"bucket"
],
destination_object=(
f"{inp_parameters[entry]['destination']['prefix']}"
f"/{inp_parameters[entry]['destination']['object']}"
),
impersonation_chain=impersonated_service_account,
)
check_size = check_compliance_of_objects_sizes(
date_of_execution="{{ds_nodash}}",
data_to_check=inp_parameters[entry],
impersonation_chain=impersonated_service_account,
)
sensor_destination_file = GCSObjectExistenceSensor(
task_id=f"sensor_destination_file_{entry}",
bucket=inp_parameters[entry]["destination"]["bucket"],
object=(
f"{inp_parameters[entry]['destination']['prefix']}"
f"/{inp_parameters[entry]['destination']['object']}"
),
impersonation_chain=impersonated_service_account,
)
end_op = DummyOperator(task_id=f"end_{entry}")
(
sensor_source_file
>> copy_file
>> check_size
>> sensor_destination_file
>> end_op
)
except FileNotFoundError:
logging.info(
f"File {inp_parameters_path} is generated in a prior task."
)
As you can see, I used try...except FileNotFoundError: because the error (file not found) is raised on the dag creation since the inp_parameters_path hadn't been created yet by the upstream tasks.
The inp_parameters_path is a file that is created upstream, but unfortunately, it seems that in order to display the logged steps of the task_group post execution I need to keep this file.. Which I don't because it will change from one day to another.
Example of removal of task_group contents after execution:
The contents were well present before the remove_mapping_file has run.
How to pass an XCom to a task_group the same way I can do it with a task? If that's not possible, how do I archive those created inp_parameters_path files per execution so that I can come back and browse the execution of the dag?
If it can help, here's how my dag operates :
#dag(
catchup=False,
schedule_interval="#daily",
max_active_runs=1,
dag_id="migrate_dtm_gcs_data_history",
start_date=datetime(2022, 12, 1),
dagrun_timeout=timedelta(minutes=20),
tags=[
"migration",
"dtm",
],
default_args=default_args,
)
def template_dag():
map = convert_csv_migration_to_map(
source_csv=f"{COMPOSER_GCS_LOC_PATH}/dags/dtm/migration/helpers/gcs_migration_mapping_archetype.csv",
delimiter=",",
)
updated_map_1 = update_mapping_with_gcp_project_and_buckets_ids(
src_mig_map=map,
env="dev",
impersonation_chain=IMPERSONATED_SERVICE_ACCOUNT,
)
list_files_and_prefixes = GCSListObjectsOperator(
task_id="list_files_and_prefixes",
bucket=LEGACY_DTM_HISTORY_BUCKET,
prefix="raw_data/datamart/{{ds_nodash}}",
impersonation_chain=IMPERSONATED_SERVICE_ACCOUNT,
)
updated_map_2 = match_data_with_migration_map(
src_mig_map=updated_map_1,
files_and_prefixes=list_files_and_prefixes.output,
)
list_of_maps = flatten_to_input_files_granularity(
src_mig_map=updated_map_2,
)
gcs_migration_maps_path = (
f"{COMPOSER_GCS_LOC_PATH}/dags/dtm/migration"
f"/helpers/gcs_migration_list_of_maps_{{ds_nodash}}.json"
)
write_list_of_maps_into_dict = gen_dict_file_from_list_of_map(
file_name=gcs_migration_maps_path,
list_of_maps=list_of_maps,
)
copy_task_group = copy_and_verify(
inp_parameters_path=gcs_migration_maps_path,
impersonated_service_account=IMPERSONATED_SERVICE_ACCOUNT,
)
remove_mapping_file = BashOperator(
task_id="remove_mapping_file",
bash_command=f"rm -vf {gcs_migration_maps_path}",
)
(
map
>> updated_map_1
>> list_files_and_prefixes
>> updated_map_2
>> list_of_maps
>> write_list_of_maps_into_dict
>> copy_task_group
>> remove_mapping_file
)
dag = template_dag()