Make airflow dags test be able to execute Mapped Tasks

[ashb]

In order to do this there were a few required:

• The BackfillJob needs to know about mapped tasks, both to expand them, and in order to update it's TI tracking
• The DebugExecutor, task_command.py, LocalTaskJob needed to "unmap" the mapped task to get the real operator back at the right time.

I was testing this with the following dag:

fromairflowimportDAGfromairflow.decoratorsimporttaskfromairflow.operators.pythonimportPythonOperatorimportpendulum@taskdefmake_list():
returnlist(map(lambdaa: f'echo "{a!r}"', [1, 2, {'a': 'b'}]))
defconsumer(*args):
print(repr(args))
withDAG(dag_id='maptest', start_date=pendulum.DateTime(2022, 1, 18)) asdag:
PythonOperator(task_id='consumer', python_callable=consumer).map(op_args=make_list())

It can't "unmap" decorated operators successfully yet, so we're using old-school PythonOperator

We also just pass the whole value to the operator, not just the current mapping value(s)

---

^ 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]

No tests yet -- that's what I'm working on now.

[uranusjr]

I wonder if it’d be nicer to also implement unmap() on BaseOperator and have it simply return self.

[ashb]

I made it raise an exception (as it should never be called?) but exist for typing.

I could have it return self and then remove the if task.is_mapped -- WDYT?

[uranusjr]

My OOP training tells me polymorphism is better than an if clause, so my vote goes to return self. But ultimately I guess either makes sense (depending on how your mental model is of the logic).

[uranusjr]

Is this guaranteed (and why?) or just a temporary measure to get pass Mypy?

[ashb]

Temporary for now until we support mapping task groups (the DAG api supports it, but the runtime doesn't yet)

[ashb]

The why: Because expand_mapped_task doesn't exist on BaseOperator, only MappedOperator

[uranusjr]

Makes sense

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

[ashb]

Not to self: if you want CI to run a test dag file, it helps if you commit it too.

[ashb]

A lot of SSH timeouts, and the MSSQL failures.

[uranusjr]

Yeah those SSH things are popping up a lot recently, not sure what had changed 😞

[potiuk]

Yeah those SSH things are popping up a lot recently, not sure what had changed disappointed

Yeah. I tried to address it via #21262 and it seemed to help a bit but I believe this is a side effect of some unrelated test, which we should track down and fix.

[ashb]

Maybe have had some un-expected luck via #21292 (that, or what ever environmental issue was causing it has vanished)

[ashb]

Right, only SSH failures (which I have a PR open for) so I'm going to merge this now.

[ashb]

In order to do this there were a few required:

• The BackfillJob needs to know about mapped tasks, both to expand them, and in order to update it's TI tracking
• The DebugExecutor, task_command.py, LocalTaskJob needed to "unmap" the mapped task to get the real operator back at the right time.

I was testing this with the following dag:

fromairflowimportDAGfromairflow.decoratorsimporttaskfromairflow.operators.pythonimportPythonOperatorimportpendulum@taskdefmake_list():
returnlist(map(lambdaa: f'echo "{a!r}"', [1, 2, {'a': 'b'}]))
defconsumer(*args):
print(repr(args))
withDAG(dag_id='maptest', start_date=pendulum.DateTime(2022, 1, 18)) asdag:
PythonOperator(task_id='consumer', python_callable=consumer).map(op_args=make_list())

It can't "unmap" decorated operators successfully yet, so we're using old-school PythonOperator

We also just pass the whole value to the operator, not just the current mapping value(s)

---

^ 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]

No tests yet -- that's what I'm working on now.

[uranusjr]

I wonder if it’d be nicer to also implement unmap() on BaseOperator and have it simply return self.

[ashb]

I made it raise an exception (as it should never be called?) but exist for typing.

I could have it return self and then remove the if task.is_mapped -- WDYT?

[uranusjr]

My OOP training tells me polymorphism is better than an if clause, so my vote goes to return self. But ultimately I guess either makes sense (depending on how your mental model is of the logic).

[uranusjr]

Is this guaranteed (and why?) or just a temporary measure to get pass Mypy?

[ashb]

Temporary for now until we support mapping task groups (the DAG api supports it, but the runtime doesn't yet)

[ashb]

The why: Because expand_mapped_task doesn't exist on BaseOperator, only MappedOperator

[uranusjr]

Makes sense

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

[ashb]

Not to self: if you want CI to run a test dag file, it helps if you commit it too.

[ashb]

A lot of SSH timeouts, and the MSSQL failures.

[uranusjr]

Yeah those SSH things are popping up a lot recently, not sure what had changed 😞

[potiuk]

Yeah those SSH things are popping up a lot recently, not sure what had changed disappointed

Yeah. I tried to address it via #21262 and it seemed to help a bit but I believe this is a side effect of some unrelated test, which we should track down and fix.

[ashb]

Maybe have had some un-expected luck via #21292 (that, or what ever environmental issue was causing it has vanished)

[ashb]

Right, only SSH failures (which I have a PR open for) so I'm going to merge this now.

[ashb]

In order to do this there were a few required:

• The BackfillJob needs to know about mapped tasks, both to expand them, and in order to update it's TI tracking
• The DebugExecutor, task_command.py, LocalTaskJob needed to "unmap" the mapped task to get the real operator back at the right time.

I was testing this with the following dag:

fromairflowimportDAGfromairflow.decoratorsimporttaskfromairflow.operators.pythonimportPythonOperatorimportpendulum@taskdefmake_list():
returnlist(map(lambdaa: f'echo "{a!r}"', [1, 2, {'a': 'b'}]))
defconsumer(*args):
print(repr(args))
withDAG(dag_id='maptest', start_date=pendulum.DateTime(2022, 1, 18)) asdag:
PythonOperator(task_id='consumer', python_callable=consumer).map(op_args=make_list())

It can't "unmap" decorated operators successfully yet, so we're using old-school PythonOperator

We also just pass the whole value to the operator, not just the current mapping value(s)

---

^ 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]

No tests yet -- that's what I'm working on now.

[uranusjr]

I wonder if it’d be nicer to also implement unmap() on BaseOperator and have it simply return self.

[ashb]

I made it raise an exception (as it should never be called?) but exist for typing.

I could have it return self and then remove the if task.is_mapped -- WDYT?

[uranusjr]

My OOP training tells me polymorphism is better than an if clause, so my vote goes to return self. But ultimately I guess either makes sense (depending on how your mental model is of the logic).

[uranusjr]

Is this guaranteed (and why?) or just a temporary measure to get pass Mypy?

[ashb]

Temporary for now until we support mapping task groups (the DAG api supports it, but the runtime doesn't yet)

[ashb]

The why: Because expand_mapped_task doesn't exist on BaseOperator, only MappedOperator

[uranusjr]

Makes sense

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

[ashb]

Not to self: if you want CI to run a test dag file, it helps if you commit it too.

[ashb]

A lot of SSH timeouts, and the MSSQL failures.

[uranusjr]

Yeah those SSH things are popping up a lot recently, not sure what had changed 😞

[potiuk]

Yeah those SSH things are popping up a lot recently, not sure what had changed disappointed

Yeah. I tried to address it via #21262 and it seemed to help a bit but I believe this is a side effect of some unrelated test, which we should track down and fix.

[ashb]

Maybe have had some un-expected luck via #21292 (that, or what ever environmental issue was causing it has vanished)

[ashb]

Right, only SSH failures (which I have a PR open for) so I'm going to merge this now.

[ashb]

In order to do this there were a few required:

• The BackfillJob needs to know about mapped tasks, both to expand them, and in order to update it's TI tracking
• The DebugExecutor, task_command.py, LocalTaskJob needed to "unmap" the mapped task to get the real operator back at the right time.

I was testing this with the following dag:

fromairflowimportDAGfromairflow.decoratorsimporttaskfromairflow.operators.pythonimportPythonOperatorimportpendulum@taskdefmake_list():
returnlist(map(lambdaa: f'echo "{a!r}"', [1, 2, {'a': 'b'}]))
defconsumer(*args):
print(repr(args))
withDAG(dag_id='maptest', start_date=pendulum.DateTime(2022, 1, 18)) asdag:
PythonOperator(task_id='consumer', python_callable=consumer).map(op_args=make_list())

It can't "unmap" decorated operators successfully yet, so we're using old-school PythonOperator

We also just pass the whole value to the operator, not just the current mapping value(s)

---

^ 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]

No tests yet -- that's what I'm working on now.

[uranusjr]

I wonder if it’d be nicer to also implement unmap() on BaseOperator and have it simply return self.

[ashb]

I made it raise an exception (as it should never be called?) but exist for typing.

I could have it return self and then remove the if task.is_mapped -- WDYT?

[uranusjr]

My OOP training tells me polymorphism is better than an if clause, so my vote goes to return self. But ultimately I guess either makes sense (depending on how your mental model is of the logic).

[uranusjr]

Is this guaranteed (and why?) or just a temporary measure to get pass Mypy?

[ashb]

Temporary for now until we support mapping task groups (the DAG api supports it, but the runtime doesn't yet)

[ashb]

The why: Because expand_mapped_task doesn't exist on BaseOperator, only MappedOperator

[uranusjr]

Makes sense

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

[ashb]

Not to self: if you want CI to run a test dag file, it helps if you commit it too.

[ashb]

A lot of SSH timeouts, and the MSSQL failures.

[uranusjr]

Yeah those SSH things are popping up a lot recently, not sure what had changed 😞

[potiuk]

Yeah those SSH things are popping up a lot recently, not sure what had changed disappointed

Yeah. I tried to address it via #21262 and it seemed to help a bit but I believe this is a side effect of some unrelated test, which we should track down and fix.

[ashb]

Maybe have had some un-expected luck via #21292 (that, or what ever environmental issue was causing it has vanished)

[ashb]

Right, only SSH failures (which I have a PR open for) so I'm going to merge this now.

[ashb]

In order to do this there were a few required:

• The BackfillJob needs to know about mapped tasks, both to expand them, and in order to update it's TI tracking
• The DebugExecutor, task_command.py, LocalTaskJob needed to "unmap" the mapped task to get the real operator back at the right time.

I was testing this with the following dag:

fromairflowimportDAGfromairflow.decoratorsimporttaskfromairflow.operators.pythonimportPythonOperatorimportpendulum@taskdefmake_list():
returnlist(map(lambdaa: f'echo "{a!r}"', [1, 2, {'a': 'b'}]))
defconsumer(*args):
print(repr(args))
withDAG(dag_id='maptest', start_date=pendulum.DateTime(2022, 1, 18)) asdag:
PythonOperator(task_id='consumer', python_callable=consumer).map(op_args=make_list())

It can't "unmap" decorated operators successfully yet, so we're using old-school PythonOperator

We also just pass the whole value to the operator, not just the current mapping value(s)

---

^ 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]

No tests yet -- that's what I'm working on now.

[uranusjr]

I wonder if it’d be nicer to also implement unmap() on BaseOperator and have it simply return self.

[ashb]

I made it raise an exception (as it should never be called?) but exist for typing.

I could have it return self and then remove the if task.is_mapped -- WDYT?

[uranusjr]

My OOP training tells me polymorphism is better than an if clause, so my vote goes to return self. But ultimately I guess either makes sense (depending on how your mental model is of the logic).

[uranusjr]

Is this guaranteed (and why?) or just a temporary measure to get pass Mypy?

[ashb]

Temporary for now until we support mapping task groups (the DAG api supports it, but the runtime doesn't yet)

[ashb]

The why: Because expand_mapped_task doesn't exist on BaseOperator, only MappedOperator

[uranusjr]

Makes sense

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

[ashb]

Not to self: if you want CI to run a test dag file, it helps if you commit it too.

[ashb]

A lot of SSH timeouts, and the MSSQL failures.

[uranusjr]

Yeah those SSH things are popping up a lot recently, not sure what had changed 😞

[potiuk]

Yeah those SSH things are popping up a lot recently, not sure what had changed disappointed

Yeah. I tried to address it via #21262 and it seemed to help a bit but I believe this is a side effect of some unrelated test, which we should track down and fix.

[ashb]

Maybe have had some un-expected luck via #21292 (that, or what ever environmental issue was causing it has vanished)

[ashb]

Right, only SSH failures (which I have a PR open for) so I'm going to merge this now.

[ashb]

In order to do this there were a few required:

• The BackfillJob needs to know about mapped tasks, both to expand them, and in order to update it's TI tracking
• The DebugExecutor, task_command.py, LocalTaskJob needed to "unmap" the mapped task to get the real operator back at the right time.

I was testing this with the following dag:

fromairflowimportDAGfromairflow.decoratorsimporttaskfromairflow.operators.pythonimportPythonOperatorimportpendulum@taskdefmake_list():
returnlist(map(lambdaa: f'echo "{a!r}"', [1, 2, {'a': 'b'}]))
defconsumer(*args):
print(repr(args))
withDAG(dag_id='maptest', start_date=pendulum.DateTime(2022, 1, 18)) asdag:
PythonOperator(task_id='consumer', python_callable=consumer).map(op_args=make_list())

It can't "unmap" decorated operators successfully yet, so we're using old-school PythonOperator

We also just pass the whole value to the operator, not just the current mapping value(s)

---

^ 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]

No tests yet -- that's what I'm working on now.

[uranusjr]

I wonder if it’d be nicer to also implement unmap() on BaseOperator and have it simply return self.

[ashb]

I made it raise an exception (as it should never be called?) but exist for typing.

I could have it return self and then remove the if task.is_mapped -- WDYT?

[uranusjr]

My OOP training tells me polymorphism is better than an if clause, so my vote goes to return self. But ultimately I guess either makes sense (depending on how your mental model is of the logic).

[uranusjr]

Is this guaranteed (and why?) or just a temporary measure to get pass Mypy?

[ashb]

Temporary for now until we support mapping task groups (the DAG api supports it, but the runtime doesn't yet)

[ashb]

The why: Because expand_mapped_task doesn't exist on BaseOperator, only MappedOperator

[uranusjr]

Makes sense

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

[ashb]

Not to self: if you want CI to run a test dag file, it helps if you commit it too.

[ashb]

A lot of SSH timeouts, and the MSSQL failures.

[uranusjr]

Yeah those SSH things are popping up a lot recently, not sure what had changed 😞

[potiuk]

Yeah those SSH things are popping up a lot recently, not sure what had changed disappointed

Yeah. I tried to address it via #21262 and it seemed to help a bit but I believe this is a side effect of some unrelated test, which we should track down and fix.

[ashb]

Maybe have had some un-expected luck via #21292 (that, or what ever environmental issue was causing it has vanished)

[ashb]

Right, only SSH failures (which I have a PR open for) so I'm going to merge this now.

[ashb]

In order to do this there were a few required:

• The BackfillJob needs to know about mapped tasks, both to expand them, and in order to update it's TI tracking
• The DebugExecutor, task_command.py, LocalTaskJob needed to "unmap" the mapped task to get the real operator back at the right time.

I was testing this with the following dag:

fromairflowimportDAGfromairflow.decoratorsimporttaskfromairflow.operators.pythonimportPythonOperatorimportpendulum@taskdefmake_list():
returnlist(map(lambdaa: f'echo "{a!r}"', [1, 2, {'a': 'b'}]))
defconsumer(*args):
print(repr(args))
withDAG(dag_id='maptest', start_date=pendulum.DateTime(2022, 1, 18)) asdag:
PythonOperator(task_id='consumer', python_callable=consumer).map(op_args=make_list())

It can't "unmap" decorated operators successfully yet, so we're using old-school PythonOperator

We also just pass the whole value to the operator, not just the current mapping value(s)

---

^ 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]

No tests yet -- that's what I'm working on now.

[uranusjr]

I wonder if it’d be nicer to also implement unmap() on BaseOperator and have it simply return self.

[ashb]

I made it raise an exception (as it should never be called?) but exist for typing.

I could have it return self and then remove the if task.is_mapped -- WDYT?

[uranusjr]

My OOP training tells me polymorphism is better than an if clause, so my vote goes to return self. But ultimately I guess either makes sense (depending on how your mental model is of the logic).

[uranusjr]

Is this guaranteed (and why?) or just a temporary measure to get pass Mypy?

[ashb]

Temporary for now until we support mapping task groups (the DAG api supports it, but the runtime doesn't yet)

[ashb]

The why: Because expand_mapped_task doesn't exist on BaseOperator, only MappedOperator

[uranusjr]

Makes sense

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

[ashb]

Not to self: if you want CI to run a test dag file, it helps if you commit it too.

[ashb]

A lot of SSH timeouts, and the MSSQL failures.

[uranusjr]

Yeah those SSH things are popping up a lot recently, not sure what had changed 😞

[potiuk]

Yeah those SSH things are popping up a lot recently, not sure what had changed disappointed

Yeah. I tried to address it via #21262 and it seemed to help a bit but I believe this is a side effect of some unrelated test, which we should track down and fix.

[ashb]

Maybe have had some un-expected luck via #21292 (that, or what ever environmental issue was causing it has vanished)

[ashb]

Right, only SSH failures (which I have a PR open for) so I'm going to merge this now.

[ashb]

In order to do this there were a few required:

• The BackfillJob needs to know about mapped tasks, both to expand them, and in order to update it's TI tracking
• The DebugExecutor, task_command.py, LocalTaskJob needed to "unmap" the mapped task to get the real operator back at the right time.

I was testing this with the following dag:

fromairflowimportDAGfromairflow.decoratorsimporttaskfromairflow.operators.pythonimportPythonOperatorimportpendulum@taskdefmake_list():
returnlist(map(lambdaa: f'echo "{a!r}"', [1, 2, {'a': 'b'}]))
defconsumer(*args):
print(repr(args))
withDAG(dag_id='maptest', start_date=pendulum.DateTime(2022, 1, 18)) asdag:
PythonOperator(task_id='consumer', python_callable=consumer).map(op_args=make_list())

It can't "unmap" decorated operators successfully yet, so we're using old-school PythonOperator

We also just pass the whole value to the operator, not just the current mapping value(s)

---

^ 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]

No tests yet -- that's what I'm working on now.

[uranusjr]

I wonder if it’d be nicer to also implement unmap() on BaseOperator and have it simply return self.

[ashb]

I made it raise an exception (as it should never be called?) but exist for typing.

I could have it return self and then remove the if task.is_mapped -- WDYT?

[uranusjr]

My OOP training tells me polymorphism is better than an if clause, so my vote goes to return self. But ultimately I guess either makes sense (depending on how your mental model is of the logic).

[uranusjr]

Is this guaranteed (and why?) or just a temporary measure to get pass Mypy?

[ashb]

Temporary for now until we support mapping task groups (the DAG api supports it, but the runtime doesn't yet)

[ashb]

The why: Because expand_mapped_task doesn't exist on BaseOperator, only MappedOperator

[uranusjr]

Makes sense

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

[ashb]

Not to self: if you want CI to run a test dag file, it helps if you commit it too.

[ashb]

A lot of SSH timeouts, and the MSSQL failures.

[uranusjr]

Yeah those SSH things are popping up a lot recently, not sure what had changed 😞

[potiuk]

Yeah those SSH things are popping up a lot recently, not sure what had changed disappointed

Yeah. I tried to address it via #21262 and it seemed to help a bit but I believe this is a side effect of some unrelated test, which we should track down and fix.

[ashb]

Maybe have had some un-expected luck via #21292 (that, or what ever environmental issue was causing it has vanished)

[ashb]

Right, only SSH failures (which I have a PR open for) so I'm going to merge this now.