[uranusjr]

This is another PR that is likely much too big. Many (sort of interrelated) things done in one.

Rewrite @task mapped operators expansion so it behave correctly

Previously, a mapped @task was not expanded correctly because its mapped argument would look something like this:

mapped_arguments={
    "op_kwargs": {
        "arg": [1, 2, 3],
    },
}

this cannot be expanded using normal logic, because PythonOperator natively expects something like this instead:

mapped_arguments={
    "op_kwargs": [
        {"arg": 1},
        {"arg": 2},
        {"arg": 3},
    ],
}

so additional logic is implemented to modify expand_mapped_task to correctly expand.

Re-implement expand_mapped_task logic to be able to expand literals

Previously, expand_mapped_task only looks in TaskMap and cannot expand e.g. .map(arg=[1, 2, 3]). This implements a more sophisticated logic to collect information from mapped arguments and expand literal and XComArg inputs properly. Some refactoring was also done so the same logic can be reused for task-runtime value unpacking.

Additional checks added to ensure correct literal types are passed to .map()

Basically making sure it only receives lists and dicts (or XComArg, which is already checked separately when the upstream pushes to XCom).

Extend TaskInstance.render_templates to “unpack” mapped values for task execution

This is the main thing. After all template fields are rendered, additional process is done to “unpack” mapped values into individual ones based on map_index, which further mutates the operator object held by a TaskInstance. This reuses similar logic from expand_mapped_task to calculate the total length of map, so it can locate its map_index inside the series.

Previously, render_template_fields was only implemented on BaseOperator, and a MappedOperator is first unmapped before templates are rendered. This approach was unforuantely wrong, since value unpacking (which needs to happen after template rendering resolves XComArg) needs to be aware of the original MappedOperator (mainly to access the user-supplied mapped kwargs). So the new logic delays unmapping until during render_templates. If a TaskInstance’s task is a MappedOperator, its render_templates would call the MappedOperator’s render_template_fields, which unmaps itself, calls render_template_fields on the unmapped operator, unpacks values for the unmapped operator, and then returns the unmapped operator. The TaskInstance would reassign its task to the unmapped, rendered, unpacked operator. This means that TaskInstance.render_templates now may have a side effect of mutating its own task attribute. I don’t particularly like this, but couldn’t find a cleaner approach without breaking much of the existing interface (complicated by the fact that both BaseOperator.render_template_fields and BaseOperator.render_template are public API). I think this is close to the best possible interface considering existing constraints (and I try to document this as clearly as possible for future maintainability).

A few tests were added to check this is working (see TestMappedTaskInstanceReceiveValue).

[uranusjr]

This is actually working

[ashb]

LGTM. couple of possible changes, but those can be done later (or not at all)

[ashb]

I do wonder if we'll need an escape hatch for this.

I guess we can add that when someone finds a specific case.

[uranusjr]

Escape hatch for if an operator cannot correctly run __init__ with mocks?

I think eventually we may want to create a separate API for validation without actually creating the operator, perhaps a classmethod called validate that can be called here.

[ashb]

Yeah, I was thinking an escape hatch for when an Operator ctor does something that causes it to fail when given a Mock.

A proper class-level validate(**kwargs) sounds like a better idea.

[ashb]

Future possible enhancement (absolutely not in this pr):

Set spec for the mock based on the type hint of the param we are are creating the mock for.

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.