diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index cd6e48b542c3e..5f6157aa0244f 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/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: diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index ae28f3317fb79..b583e5a448ea9 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/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))