Flush the session between writing and deletion of RTIF (#42928) (#43012)

* FLush the session before deleting the RTIF data

Previously, this was how it was done, but now,
a session was used for both the writing and deletion of RTIF,
which we suspect caused StaleDataError. The related PR: #38565

This PR brings back the old behaviour of using different sessions for writing/deleting RTIFs

* fixup! Use different sessions in writing and deletion of RTIF

* add test and use flush

(cherry picked from commit ced319f)

airflow/models/taskinstance.py

@@ -1634,11 +1634,12 @@ def _get_previous_ti(
 
 @internal_api_call
 @provide_session
-def _update_rtif(ti, rendered_fields, session: Session | None = None):
+def _update_rtif(ti, rendered_fields, session: Session = NEW_SESSION):
     from airflow.models.renderedtifields import RenderedTaskInstanceFields
 
     rtif = RenderedTaskInstanceFields(ti=ti, render_templates=False, rendered_fields=rendered_fields)
     RenderedTaskInstanceFields.write(rtif, session=session)
+    session.flush()
     RenderedTaskInstanceFields.delete_old_records(ti.task_id, ti.dag_id, session=session)
 
 

tests/models/test_renderedtifields.py

@@ -24,14 +24,17 @@
 from datetime import date, timedelta
 from unittest import mock
 
+import pendulum
 import pytest
+from sqlalchemy import select
 
 from airflow import settings
 from airflow.configuration import conf
 from airflow.decorators import task as task_decorator
-from airflow.models import Variable
+from airflow.models import DagRun, Variable
 from airflow.models.renderedtifields import RenderedTaskInstanceFields as RTIF
 from airflow.operators.bash import BashOperator
+from airflow.operators.python import PythonOperator
 from airflow.utils.task_instance_session import set_current_task_instance_session
 from airflow.utils.timezone import datetime
 from tests.test_utils.asserts import assert_queries_count
@@ -385,3 +388,48 @@ def test_redact(self, redact, dag_maker):
             "env": "val 2",
             "cwd": "val 3",
         }
+
+    @pytest.mark.skip_if_database_isolation_mode
+    def test_rtif_deletion_stale_data_error(self, dag_maker, session):
+        """
+        Here we verify bad behavior.  When we rerun a task whose RTIF
+        will get removed, we get a stale data error.
+        """
+        with dag_maker(dag_id="test_retry_handling"):
+            task = PythonOperator(
+                task_id="test_retry_handling_op",
+                python_callable=lambda a, b: print(f"{a}\n{b}\n"),
+                op_args=[
+                    "dag {{dag.dag_id}};",
+                    "try_number {{ti.try_number}};yo",
+                ],
+            )
+
+        def run_task(date):
+            run_id = f"abc_{date.to_date_string()}"
+            dr = session.scalar(select(DagRun).where(DagRun.execution_date == date, DagRun.run_id == run_id))
+            if not dr:
+                dr = dag_maker.create_dagrun(execution_date=date, run_id=run_id)
+            ti = dr.task_instances[0]
+            ti.state = None
+            ti.try_number += 1
+            session.commit()
+            ti.task = task
+            ti.run()
+            return dr
+
+        base_date = pendulum.datetime(2021, 1, 1)
+        exec_dates = [base_date.add(days=x) for x in range(40)]
+        for date_ in exec_dates:
+            run_task(date=date_)
+
+        session.commit()
+        session.expunge_all()
+
+        # find oldest date
+        date = session.scalar(
+            select(DagRun.execution_date).join(RTIF.dag_run).order_by(DagRun.execution_date).limit(1)
+        )
+        date = pendulum.instance(date)
+        # rerun the old date. this will fail
+        run_task(date=date)