Support Task Flow and enhance dynamic task mapping

[tatiana]

Implement support for Airflow TaskFlow, available since 2.0.

How to test

The following example defines a task that generates a list of numbers and another that consumes this list and creates dynamically (using Airflow dynamic task mapping) an independent task that doubles each individual number.

example_taskflow:
  default_args:
    owner: "custom_owner"
    start_date: 2 days
  description: "Example of TaskFlow powered DAG that includes dynamic task mapping."
  schedule_interval: "0 3 * * *"
  default_view: "graph"
  tasks:

    numbers_list:
      decorator: airflow.decorators.task
      python_callable: sample.build_numbers_list

    double_number_with_dynamic_task_mapping_taskflow:
      decorator: airflow.decorators.task
      python_callable: sample.double
      expand:
          number: +numbers_list  # the prefix + tells DagFactory to resolve this value as the task `numbers_list`, previously defined

For the sample.py file below:

def build_numbers_list():
    return [2, 4, 6]


def double(number: int):
    result = 2 * number
    print(result)
    return result

In the UI, it is shown as:

And:

Scope

This PR includes several use cases of dynamic task mapping:

1. Simple mapping
2. Task-generated mapping
3. Repeated mapping
4. Adding parameters that do not expand (partial)
5. Mapping over multiple parameters
6. Named mapping (map_index_template)

The following dynamic task mapping cases were not tested but are expected to work:

• Mapping with non-TaskFlow operators
• Mapping over the result of classic operators
• Filtering items from a mapped task

The following dynamic task mapping cases were not tested and should not work (they were considered outside of the scope of the current ticket):

• Assigning multiple parameters to a non-TaskFlow operator
• Mapping over a task group
• Transforming expanding data
• Combining upstream data (aka “zipping”)

Tests

The feature is being tested by running the example DAGs introduced in this PR, which validate various scenarios of task flow and dynamic task mapping and serve as documentation.

As with other parts of DAG Factory, we can and should improve the overall unit test coverage.

Two example DAG files were added, containing multiple examples of TaskFlow and Dynamic Task mapping. This is how they are displayed in the AIrflow UI:

Docs

This PR does not contain user-facing docs other than the README. However, we'll address this as part of #278.

Related issues

This PR closes two open tickets:

Closes: #302 (support named mapping, via the map_index_template argument)

Example of usage of map_index_template:

    dynamic_task_with_named_mapping:
      decorator: airflow.decorators.task
      python_callable: sample.extract_last_name
      map_index_template: "{{ custom_mapping_key }}"
      expand:
        full_name:
          - Lucy Black
          - Vera Santos
          - Marks Spencer

Closes: #301 (Mapping over multiple parameters)

Example of multiple parameters:

    multiply_with_multiple_parameters:
      decorator: airflow.decorators.task
      python_callable: sample.multiply
      expand:
          a: +numbers_list  # the prefix + tells DagFactory to resolve this value as the task `numbers_list`, previously defined
          b: +another_numbers_list # the prefix + tells DagFactory to resolve this value as the task `another_numbers_list`, previously defined

[pankajastro]

LGTM, I'm really excited for for this feature

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 89.47368% with 10 lines in your changes missing coverage. Please review.

Project coverage is 93.29%. Comparing base (80b885e) to head (aa0ef3e).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
dagfactory/dagbuilder.py	89.47%	10 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #314      +/-   ##
==========================================
+ Coverage   92.69%   93.29%   +0.59%     
==========================================
  Files          10       10              
  Lines         726      776      +50     
==========================================
+ Hits          673      724      +51     
+ Misses         53       52       -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.