Cascade update of taskinstance to TaskMap table (#31445)

(cherry picked from commit f6bb474)

Author: ephraimbuddy

airflow/migrations/versions/0125_2_6_2_add_onupdate_cascade_to_taskmap.py

@@ -0,0 +1,62 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Add ``onupdate`` cascade to ``task_map`` table
+
+Revision ID: c804e5c76e3e
+Revises: 98ae134e6fff
+Create Date: 2023-05-19 23:30:57.368617
+
+"""
+from __future__ import annotations
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "c804e5c76e3e"
+down_revision = "98ae134e6fff"
+branch_labels = None
+depends_on = None
+airflow_version = "2.6.2"
+
+
+def upgrade():
+    """Apply Add onupdate cascade to taskmap"""
+    with op.batch_alter_table("task_map") as batch_op:
+        batch_op.drop_constraint("task_map_task_instance_fkey", type_="foreignkey")
+        batch_op.create_foreign_key(
+            "task_map_task_instance_fkey",
+            "task_instance",
+            ["dag_id", "task_id", "run_id", "map_index"],
+            ["dag_id", "task_id", "run_id", "map_index"],
+            ondelete="CASCADE",
+            onupdate="CASCADE",
+        )
+
+
+def downgrade():
+    """Unapply Add onupdate cascade to taskmap"""
+    with op.batch_alter_table("task_map") as batch_op:
+        batch_op.drop_constraint("task_map_task_instance_fkey", type_="foreignkey")
+        batch_op.create_foreign_key(
+            "task_map_task_instance_fkey",
+            "task_instance",
+            ["dag_id", "task_id", "run_id", "map_index"],
+            ["dag_id", "task_id", "run_id", "map_index"],
+            ondelete="CASCADE",
+        )

airflow/models/taskmap.py

@@ -72,6 +72,7 @@ class TaskMap(Base):
             ],
             name="task_map_task_instance_fkey",
             ondelete="CASCADE",
+            onupdate="CASCADE",
         ),
     )
 

docs/apache-airflow/img/airflow_erd.sha256

@@ -1 +1 @@
-c9c60f41925503ca504eca8b6688ca7ec2d8d7cbbd711285ae5c711cc2599f0e
+b8e715d65aba5965f3ae0134ee7c38dcff539ebf27442db9771145ed7fb18186

docs/apache-airflow/img/airflow_erd.svg

@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
 | Revision ID                     | Revises ID        | Airflow Version   | Description                                                  |
 +=================================+===================+===================+==============================================================+
-| ``98ae134e6fff`` (head)         | ``6abdffdd4815``  | ``2.6.0``         | Increase length of user identifier columns in ``ab_user``    |
+| ``c804e5c76e3e`` (head)         | ``98ae134e6fff``  | ``2.6.2``         | Add ``onupdate`` cascade to ``task_map`` table               |
++---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
+| ``98ae134e6fff``                | ``6abdffdd4815``  | ``2.6.0``         | Increase length of user identifier columns in ``ab_user``    |
 |                                 |                   |                   | and ``ab_register_user`` tables                              |
 +---------------------------------+-------------------+-------------------+--------------------------------------------------------------+
 | ``6abdffdd4815``                | ``290244fb8b83``  | ``2.6.0``         | add dttm index on log table                                  |

docs/apache-airflow/migrations-ref.rst

@@ -69,7 +69,7 @@
 from airflow.operators.python import PythonOperator
 from airflow.sensors.base import BaseSensorOperator
 from airflow.sensors.python import PythonSensor
-from airflow.serialization.serialized_objects import SerializedBaseOperator
+from airflow.serialization.serialized_objects import SerializedBaseOperator, SerializedDAG
 from airflow.settings import TIMEZONE
 from airflow.stats import Stats
 from airflow.ti_deps.dep_context import DepContext
@@ -3484,6 +3484,39 @@ def pull_something(value):
         assert task_map.length == expected_length
         assert task_map.keys == expected_keys
 
+    def test_no_error_on_changing_from_non_mapped_to_mapped(self, dag_maker, session):
+        """If a task changes from non-mapped to mapped, don't fail on integrity error."""
+        with dag_maker(dag_id="test_no_error_on_changing_from_non_mapped_to_mapped") as dag:
+
+            @dag.task()
+            def add_one(x):
+                return [x + 1]
+
+            @dag.task()
+            def add_two(x):
+                return x + 2
+
+            task1 = add_one(2)
+            add_two.expand(x=task1)
+
+        dr = dag_maker.create_dagrun()
+        ti = dr.get_task_instance(task_id="add_one")
+        ti.run()
+        assert ti.state == TaskInstanceState.SUCCESS
+        dag._remove_task("add_one")
+        with dag:
+            task1 = add_one.expand(x=[1, 2, 3]).operator
+        serialized_dag = SerializedDAG.from_dict(SerializedDAG.to_dict(dag))
+
+        dr.dag = serialized_dag
+        dr.verify_integrity(session=session)
+        ti = dr.get_task_instance(task_id="add_one")
+        assert ti.state == TaskInstanceState.REMOVED
+        dag.clear()
+        ti.refresh_from_task(task1)
+        # This should not raise an integrity error
+        dr.task_instance_scheduling_decisions()
+
 
 class TestMappedTaskInstanceReceiveValue:
     @pytest.mark.parametrize(