Implement AiiDA provenance tracking for Airflow via XCom backend and listeners

Add comprehensive AiiDA provenance graph creation for Airflow DAG executions:

XCom Backend (src/airflow_provider_aiida/xcom/backend.py):
- Custom XCom backend that creates AiiDA nodes during serialization/deserialization
- WorkChainNode represents entire DAG run (dag_id + run_id)
- CalcJobNode represents individual tasks (task_id + run_id + map_index)
- Data nodes store XCom values as AiiDA-typed data
- Establishes provenance links:
  * CALL_CALC: WorkChain → CalcJob (workflow calls calculation)
  * CREATE: CalcJob → Data (task produces output)
  * INPUT_CALC: Data → CalcJob (data flows to consuming task)
- Smart link labeling:
  * For PythonOperator: extracts parameter names via inspect.signature()
  * For other operators: deterministic hash based on producer task info
- Handles duplicate link prevention and stored node constraints
- Monkey-patches link validation for stored nodes when necessary

Provenance Listener (src/airflow_provider_aiida/plugins/provenance_listener.py):
- Airflow listener plugin that updates AiiDA node states in real-time
- Hooks into DAG run lifecycle (success, failure, running)
- Hooks into task instance lifecycle (success, failure, running)
- Maps Airflow states to AiiDA ProcessStates:
  * QUEUED/SCHEDULED → CREATED
  * RUNNING → RUNNING
  * SUCCESS → FINISHED
  * FAILED → EXCEPTED
  * SKIPPED → KILLED
- Creates nodes proactively if they don't exist (handles tasks starting before XCom)
- Registered as ProvenanceListenerPlugin for automatic discovery

Common utilities to handle aiida nodes(src/airflow_provider_aiida/common/utils.py):
- _get_or_create_workchain_node: Query by unique_id or create new WorkChainNode
- _get_or_create_calcjob_node: Query by unique_id or create new CalcJobNode
- _sanitize_link_label: Ensure AiiDA-compatible link labels (alphanumeric + underscore)
- All new nodes initialized with ProcessState.CREATED

Caveats:
- When deserializing no information about the input key is given, so an
  educated guess has to be made which for the moment fails when maps are used
- on_dag_run_running is not called in test run environment, therefore the
  workchain node is created in on_task_run_running function
- Because we have no guarantee from airflow for the order of callbacks (executed by the task instance) and xcom backend (executed by the scheduler) we have to make logic redundant in the xcom backend and listeners

Result: Complete AiiDA provenance graph mirroring Airflow DAG structure with
real-time state synchronization and proper data lineage tracking.

Author: agoscinski

pyproject.toml

@@ -43,6 +43,9 @@ provider_info = "airflow_provider_aiida.__init__:get_provider_info"
 [project.entry-points."aiida.dags"]
 aiida-standard = "airflow_provider_aiida.example_dags"
 
+[project.entry-points."airflow.plugins"]
+aiida_dag_run_listener = "airflow_provider_aiida.plugins.provenance_listener:ProvenanceListenerPlugin"
+
 [tool.hatch.version]
 path = "src/airflow_provider_aiida/__init__.py"
 

src/airflow_provider_aiida/common/__init__.py

@@ -0,0 +1 @@
+"""Common utilities for the AiiDA Airflow provider."""

src/airflow_provider_aiida/common/utils.py

@@ -0,0 +1,162 @@
+"""Common ORM utilities for AiiDA nodes in the Airflow provider."""
+
+import re
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+# Add file handler for detailed ORM logging
+log_file = os.path.expanduser('~/airflow_provider_aiida_orm.log')
+file_handler = logging.FileHandler(log_file, mode='a')
+file_handler.setLevel(logging.DEBUG)
+formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+file_handler.setFormatter(formatter)
+logger.addHandler(file_handler)
+# TODO for prototyping phase we put it on DEBUG
+logger.setLevel(logging.DEBUG)
+
+logger.info(f"ORM logging initialized, writing to: {log_file}")
+
+
+def _sanitize_link_label(label: str) -> str:
+    """Sanitize a string to be a valid AiiDA link label.
+
+    AiiDA link labels must contain only alphanumeric characters and underscores.
+
+    :param label: The label to sanitize
+    :return: Sanitized label with only valid characters
+    """
+    # Replace any non-alphanumeric, non-underscore characters with underscores
+    sanitized = re.sub(r'[^a-zA-Z0-9_]', '_', label)
+    # Ensure it doesn't start with a number (prepend 'x' if it does)
+    if sanitized and sanitized[0].isdigit():
+        sanitized = 'x' + sanitized
+    return sanitized or 'result'
+
+
+def _get_workchain_node(dag_id: str, run_id: str):
+    """Get an existing WorkChainNode for an Airflow DAG run.
+
+    Returns None if no node exists.
+
+    :param dag_id: Airflow DAG ID
+    :param run_id: Airflow run ID
+    :return: WorkChainNode instance or None
+    """
+    from aiida.orm import WorkChainNode, QueryBuilder
+
+    # Create a unique identifier for this DAG run
+    unique_id = f"{dag_id}__{run_id}"
+
+    # Check if node already exists using the unique_id stored in extras
+    qb = QueryBuilder()
+    qb.append(WorkChainNode, filters={'extras.airflow_unique_id': unique_id})
+    results = qb.all()
+
+    if results:
+        existing_node = results[0][0]
+        logger.debug(f"Found existing WorkChainNode: dag_id={dag_id}, run_id={run_id}, pk={existing_node.pk}")
+        return existing_node
+
+    logger.debug(f"No WorkChainNode found: dag_id={dag_id}, run_id={run_id}")
+    return None
+
+def _get_or_create_workchain_node(dag_id: str, run_id: str):
+    """Get or create a WorkChainNode representing an Airflow DAG run.
+
+    Sets initial state to CREATED when creating a new node.
+
+    :param dag_id: Airflow DAG ID
+    :param run_id: Airflow run ID
+    :return: Tuple of (WorkChainNode instance, created: bool)
+             created is True if a new node was created, False if existing node was found
+    """
+    from aiida import load_profile
+    from aiida.orm import WorkChainNode, QueryBuilder
+    from plumpy.process_states import ProcessState
+    load_profile()
+
+    # Create a unique identifier for this DAG run
+    unique_id = f"{dag_id}__{run_id}"
+
+    # Check if node already exists using the unique_id stored in extras
+    qb = QueryBuilder()
+    qb.append(WorkChainNode, filters={'extras.airflow_unique_id': unique_id})
+    results = qb.all()
+
+    if results:
+        existing_node = results[0][0]
+        logger.debug(f"Found existing WorkChainNode: dag_id={dag_id}, run_id={run_id}, pk={existing_node.pk}")
+        return existing_node, False
+
+    # Create new WorkChainNode for the DAG run
+    logger.debug(f"Creating NEW WorkChainNode: dag_id={dag_id}, run_id={run_id}, unique_id={unique_id}")
+    workchain_node = WorkChainNode()
+    workchain_node.set_process_label(f"{dag_id}[{run_id}]")
+    workchain_node.label = dag_id
+    workchain_node.description = f"Airflow DAG run: {dag_id}, run_id: {run_id}"
+    workchain_node.base.extras.set('airflow_unique_id', unique_id)
+    workchain_node.base.extras.set('airflow_dag_id', dag_id)
+    workchain_node.base.extras.set('airflow_run_id', run_id)
+    workchain_node.base.extras.set('airflow_xcom_backend', True)
+
+    # Set initial process state to CREATED
+    workchain_node.set_process_state(ProcessState.CREATED)
+
+    logger.info(f"Created NEW WorkChainNode: dag_id={dag_id}, run_id={run_id}, node={workchain_node}")
+
+    return workchain_node, True
+
+
+def _get_or_create_calcjob_node(task_id: str, dag_id: str, run_id: str, map_index: int = -1):
+    """Get or create a CalcJobNode representing an Airflow task.
+
+    Sets initial state to CREATED when creating a new node.
+
+    :param task_id: Airflow task ID
+    :param dag_id: Airflow DAG ID
+    :param run_id: Airflow run ID
+    :param map_index: Airflow map index for mapped tasks
+    :return: CalcJobNode instance
+    """
+    from aiida import load_profile
+    from aiida.orm import CalcJobNode, QueryBuilder
+    from plumpy.process_states import ProcessState
+    load_profile()
+
+    # Create a unique identifier for this task execution (used in extras for lookup)
+    unique_id = f"{dag_id}__{task_id}__{run_id}"
+    if map_index >= 0:
+        unique_id += f"__{map_index}"
+
+    # Check if node already exists using the unique_id stored in extras
+    qb = QueryBuilder()
+    qb.append(CalcJobNode, filters={'extras.airflow_unique_id': unique_id})
+    results = qb.all()
+
+    if results:
+        existing_node = results[0][0]
+        logger.debug(f"Found existing CalcJobNode: dag_id={dag_id}, task_id={task_id}, run_id={run_id}, map_index={map_index}, pk={existing_node.pk}")
+        return existing_node
+
+    # Create new CalcJobNode with task_id as label (more readable)
+    logger.warning(f"Creating NEW CalcJobNode: dag_id={dag_id}, task_id={task_id}, run_id={run_id}, map_index={map_index}, unique_id={unique_id}")
+    calc_node = CalcJobNode()
+    # Use task_id as label for readability, with map_index if applicable
+    calc_node.set_process_label(f"{task_id}[{map_index}]" if map_index >= 0 else task_id)
+    calc_node.label = task_id
+    calc_node.description = f"Airflow task: {task_id} from DAG: {dag_id}, run: {run_id}"
+    calc_node.base.extras.set('airflow_unique_id', unique_id)
+    calc_node.base.extras.set('airflow_dag_id', dag_id)
+    calc_node.base.extras.set('airflow_task_id', task_id)
+    calc_node.base.extras.set('airflow_run_id', run_id)
+    calc_node.base.extras.set('airflow_map_index', map_index)
+    calc_node.base.extras.set('airflow_xcom_backend', True)
+
+    # Set initial process state to CREATED
+    calc_node.set_process_state(ProcessState.CREATED)
+
+    logger.info(f"Created NEW CalcJobNode: dag_id={dag_id}, task_id={task_id}, run_id={run_id}, map_index={map_index}, node={calc_node}")
+
+    return calc_node