fix possible race condition when refreshing DAGs

airflow/models/dagbag.py

@@ -192,7 +192,8 @@ def get_dag(self, dag_id, session: Session = None):
 
             # If DAG is in the DagBag, check the following
             # 1. if time has come to check if DAG is updated (controlled by min_serialized_dag_fetch_secs)
-            # 2. check the last_updated column in SerializedDag table to see if Serialized DAG is updated
+            # 2. check the last_updated and hash columns in SerializedDag table to see if
+            # Serialized DAG is updated
             # 3. if (2) is yes, fetch the Serialized DAG.
             # 4. if (2) returns None (i.e. Serialized DAG is deleted), remove dag from dagbag
             # if it exists and return None.
@@ -201,18 +202,24 @@ def get_dag(self, dag_id, session: Session = None):
                 dag_id in self.dags_last_fetched
                 and timezone.utcnow() > self.dags_last_fetched[dag_id] + min_serialized_dag_fetch_secs
             ):
-                sd_last_updated_datetime = SerializedDagModel.get_last_updated_datetime(
-                    dag_id=dag_id,
-                    session=session,
+                sd_latest_version_and_updated_datetime = (
+                    SerializedDagModel.get_latest_version_hash_and_updated_datetime(
+                        dag_id=dag_id, session=session
+                    )
                 )
-                if not sd_last_updated_datetime:
+                if not sd_latest_version_and_updated_datetime:
                     self.log.warning("Serialized DAG %s no longer exists", dag_id)
                     del self.dags[dag_id]
                     del self.dags_last_fetched[dag_id]
                     del self.dags_hash[dag_id]
                     return None
 
-                if sd_last_updated_datetime > self.dags_last_fetched[dag_id]:
+                sd_latest_version, sd_last_updated_datetime = sd_latest_version_and_updated_datetime
+
+                if (
+                    sd_last_updated_datetime > self.dags_last_fetched[dag_id]
+                    or sd_latest_version != self.dags_hash[dag_id]
+                ):
                     self._add_dag_from_db(dag_id=dag_id, session=session)
 
             return self.dags.get(dag_id)

airflow/models/serialized_dag.py

@@ -359,6 +359,24 @@ def get_latest_version_hash(cls, dag_id: str, session: Session = NEW_SESSION) ->
         """
         return session.query(cls.dag_hash).filter(cls.dag_id == dag_id).scalar()
 
+    @classmethod
+    def get_latest_version_hash_and_updated_datetime(
+        cls,
+        dag_id: str,
+        *,
+        session: Session,
+    ) -> tuple[str, datetime] | None:
+        """
+        Get the latest DAG version for a given DAG ID, as well as the date when the Serialized DAG associated
+        to DAG was last updated in serialized_dag table.
+
+        :meta private:
+        :param dag_id: DAG ID
+        :param session: ORM Session
+        :return: A tuple of DAG Hash and last updated datetime, or None if the DAG is not found
+        """
+        return session.query(cls.dag_hash, cls.last_updated).filter(cls.dag_id == dag_id).one_or_none()
+
     @classmethod
     @provide_session
     def get_dag_dependencies(cls, session: Session = NEW_SESSION) -> dict[str, list[DagDependency]]:

tests/models/test_dagbag.py

@@ -980,6 +980,54 @@ def test_get_dag_with_dag_serialization(self):
         assert set(updated_ser_dag_1.tags) == {"example", "example2", "new_tag"}
         assert updated_ser_dag_1_update_time > ser_dag_1_update_time
 
+    @patch("airflow.models.dagbag.settings.MIN_SERIALIZED_DAG_UPDATE_INTERVAL", 5)
+    @patch("airflow.models.dagbag.settings.MIN_SERIALIZED_DAG_FETCH_INTERVAL", 5)
+    def test_get_dag_refresh_race_condition(self):
+        """
+        Test that DagBag.get_dag correctly refresh the Serialized DAG even if SerializedDagModel.last_updated
+        is before DagBag.dags_last_fetched.
+        """
+
+        # serialize the initial version of the DAG
+        with time_machine.travel((tz.datetime(2020, 1, 5, 0, 0, 0)), tick=False):
+            example_bash_op_dag = DagBag(include_examples=True).dags.get("example_bash_operator")
+            SerializedDagModel.write_dag(dag=example_bash_op_dag)
+
+        # deserialize the DAG
+        with time_machine.travel((tz.datetime(2020, 1, 5, 1, 0, 10)), tick=False):
+            dag_bag = DagBag(read_dags_from_db=True)
+
+            with assert_queries_count(2):
+                ser_dag = dag_bag.get_dag("example_bash_operator")
+
+            ser_dag_update_time = dag_bag.dags_last_fetched["example_bash_operator"]
+            assert ser_dag.tags == ["example", "example2"]
+            assert ser_dag_update_time == tz.datetime(2020, 1, 5, 1, 0, 10)
+
+            with create_session() as session:
+                assert SerializedDagModel.get_last_updated_datetime(
+                    dag_id="example_bash_operator",
+                    session=session,
+                ) == tz.datetime(2020, 1, 5, 0, 0, 0)
+
+        # Simulate a long-running serialization transaction
+        # Make a change in the DAG and write Serialized DAG to the DB
+        # Note the date *before* the deserialize step above, simulating a serialization happening
+        # long before the transaction is committed
+        with time_machine.travel((tz.datetime(2020, 1, 5, 1, 0, 0)), tick=False):
+            example_bash_op_dag.tags += ["new_tag"]
+            SerializedDagModel.write_dag(dag=example_bash_op_dag)
+
+        # Since min_serialized_dag_fetch_interval is passed verify that calling 'dag_bag.get_dag'
+        # fetches the Serialized DAG from DB
+        with time_machine.travel((tz.datetime(2020, 1, 5, 1, 0, 30)), tick=False):
+            with assert_queries_count(2):
+                updated_ser_dag = dag_bag.get_dag("example_bash_operator")
+                updated_ser_dag_update_time = dag_bag.dags_last_fetched["example_bash_operator"]
+
+        assert set(updated_ser_dag.tags) == {"example", "example2", "new_tag"}
+        assert updated_ser_dag_update_time > ser_dag_update_time
+
     def test_collect_dags_from_db(self):
         """DAGs are collected from Database"""
         db.clear_db_dags()