Read fresh Dag bundle state during refresh

airflow-core/src/airflow/dag_processing/manager.py

@@ -749,11 +749,11 @@ def get_bundle_state(self, bundle_name: str, *, session: Session = NEW_SESSION)
 
         Returns ``None`` if the bundle has no database record.
         """
-        row = session.scalar(
-            select(DagBundleModel)
-            .where(DagBundleModel.name == bundle_name)
-            .options(load_only(DagBundleModel.last_refreshed, DagBundleModel.version))
-        )
+        row = session.execute(
+            select(DagBundleModel.last_refreshed, DagBundleModel.version).where(
+                DagBundleModel.name == bundle_name
+            )
+        ).one_or_none()
         if row is None:
             return None
         return BundleState(last_refreshed=row.last_refreshed, version=row.version)
@@ -815,7 +815,6 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]):
                 except AirflowException as e:
                     self.log.exception("Error initializing bundle %s: %s", bundle.name, e)
                     continue
-            # TODO: AIP-66 test to make sure we get a fresh record from the db and it's not cached
             try:
                 bundle_state = self.get_bundle_state(bundle.name)
             except Exception:

airflow-core/tests/unit/dag_processing/test_manager.py

@@ -2712,6 +2712,34 @@ def test_get_bundle_state_returns_correct_state(self, session):
 
         assert state == BundleState(last_refreshed=refreshed_at, version="v1")
 
+    def test_get_bundle_state_reads_latest_database_values(self, session):
+        bundle_name = "test_fresh_state_bundle"
+        initial_refreshed_at = timezone.datetime(2024, 1, 15, 12, 0, 0)
+        refreshed_at = timezone.datetime(2024, 1, 16, 12, 0, 0)
+        model = DagBundleModel(name=bundle_name, version="old")
+        model.last_refreshed = initial_refreshed_at
+        session.add(model)
+        session.commit()
+
+        manager = DagFileProcessorManager(max_runs=1)
+        state = manager.get_bundle_state(bundle_name, session=session)
+
+        assert state == BundleState(last_refreshed=initial_refreshed_at, version="old")
+
+        with create_session(scoped=False) as update_session:
+            update_model = update_session.get(DagBundleModel, bundle_name)
+            assert update_model is not None
+            update_model.last_refreshed = refreshed_at
+            update_model.version = "fresh"
+
+        # End the read transaction started by the first get_bundle_state() call so we don't keep
+        # reading a stale snapshot on backends like SQLite.
+        session.commit()
+
+        state = manager.get_bundle_state(bundle_name, session=session)
+
+        assert state == BundleState(last_refreshed=refreshed_at, version="fresh")
+
     def test_get_bundle_state_null_fields(self, session):
         bundle_name = "test_null_state_bundle"
         session.add(DagBundleModel(name=bundle_name))