Map and Partial DAG authoring interface for Dynamic Task Mapping

[ashb]

Closes #19889

---

^ Add meaningful description above

Read the Pull Request Guidelines for more information.
In case of fundamental code change, Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in UPDATING.md.

[ashb]

This should now have all of the DAG authoring interface changes for AIP-42, including (some, if not extensive) tests.

Still to do:

• Fix mypy issues in new code (there will be lots, I was running with SKIP=mypy
• [ ] Add docs. That could be a follow up PR but it feels nicer to me to add at least basic docs here as part of this PR.

Docs will come later -- the interface might change yet still as we work out more details.

[ashb]

I've got another small change coming to better validate mapped/partial argument names (I was missing a few cases)

[uranusjr]

Not finished, but submitting this batch because my comments already went stale once before I finished.

[uranusjr]

This makes me wonder, what should MyOperator.map(x=something).map(x=another) do? If I understand this correctly, this would currently discard something and just map to another. We should likely add something in to prevent this from happening, perhaps in _validate_arg_names?

[ashb]

Yeah, that's an unspecified point of the API. I think I'm leaning towards .map().map() being an error on general grounds.

But yes, in my head updating the params was my intent.

[potiuk]

I also think map().map() should be an error. We already agreed on suporting map(arg1, arg2) and:

There should be one-- and preferably only one --obvious way to do it.

[ashb]

I'm unliekly to be able to make this change before stopping for Christmas, so either someone else can make this, or we can merge it with this and fix it later.

(I think that having a .map() function that returns an error would be clearer than having no map method, similar to how I have .partial() on a Task object throw an error still.)

[uranusjr]

What’s the purpose of this?

[ashb]

From the tests, to catch this:

    @task_group_decorator
    def tg():
        ...

    with pytest.warns(UserWarning, match='was never mapped'):
        with DAG("test-dag", start_date=DEFAULT_DATE):
            tg.partial()

I.e. tg = tg.partial().map() was needed. You doesn't make sense to have a partial TG without also mapping it once.

[uranusjr]

What would happen in this case, nothing? (i.e. this is likely a user mistake)

[ashb]

Correct -- nothing, and likely a user mistake as the tasks in the task group wouldn't appear in the DAG

[uranusjr]

I wonder if _validate_arg_names and this have similar logic to be refactored out.

[ashb]

Similar, but the way this gets the possible names (most of the function) is quite differnt, so I couldn't work out a way of really sharing it

[uranusjr]

Do we have a test to ensure get_serialized_fields() is up-to-date? Otherwise we may accidentally add a field on BaseOperator but forget to serialize it, and this test would fail to reflect that.

[ashb]

We don't, no, this is it. This will pick up anything added via self.x in the constructor, but won't pick up things like x: int = 0 at the class level.

Think it's worth it?

[uranusjr]

Didn’t read the tests.

[uranusjr]

Is Mypy OK with this? I’d make _set_relatives() only accept the sequence case, and do the coercing outside in the public functions.

[ashb]

It is, yes.

[uranusjr]

Can we not change the parent class definition to be plain class-level variable declarations?

[ashb]

No, because on BaseOperator these are defined as properties.

[ashb]

Hmmm, this one doesn't behave correctly right now, as:

@taks(task_id='x')
def y():
    ...

This currently stores task_id in the self.kwargs

[ashb]

Removed this check only in deb54df

[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.