[airflow] branch main updated: Remove double collection of dags in `airflow dags reserialize` (#27030)

[ephraimanierobi]

This is an automated email from the ASF dual-hosted git repository.

ephraimanierobi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 36e2e43def Remove double collection of dags in `airflow dags 
reserialize`  (#27030)
36e2e43def is described below

commit 36e2e43def6a27d9bf2cab4d27d104414bea3f7f
Author: Ephraim Anierobi <splendidzig...@gmail.com>
AuthorDate: Thu Oct 13 19:29:00 2022 +0100

    Remove double collection of dags in `airflow dags reserialize`  (#27030)
    
    We explicitly call dagbag.collect_dags after instantiating DagBag in
    airflow dags reserialize code.
    
    The method collect_dags is called on instantiation
    of the DagBag so calling it again means more processing of the same dags.
    
    Here, we use a variable to achieve the same needed effect on reserialization
    
    Co-authored-by: Ash Berlin-Taylor <ash_git...@firemirror.com>
---
 airflow/models/dagbag.py    | 12 +++++++-----
 airflow/utils/db.py         |  4 ++--
 tests/models/test_dagbag.py | 13 +++++++++++++
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/airflow/models/dagbag.py b/airflow/models/dagbag.py
index cabc142f31..849498c1e4 100644
--- a/airflow/models/dagbag.py
+++ b/airflow/models/dagbag.py
@@ -98,6 +98,7 @@ class DagBag(LoggingMixin):
         read_dags_from_db: bool = False,
         store_serialized_dags: bool | None = None,
         load_op_links: bool = True,
+        collect_dags: bool = True,
     ):
         # Avoid circular import
         from airflow.models.dag import DAG
@@ -137,11 +138,12 @@ class DagBag(LoggingMixin):
 
         self.dagbag_import_error_tracebacks = conf.getboolean('core', 
'dagbag_import_error_tracebacks')
         self.dagbag_import_error_traceback_depth = conf.getint('core', 
'dagbag_import_error_traceback_depth')
-        self.collect_dags(
-            dag_folder=dag_folder,
-            include_examples=include_examples,
-            safe_mode=safe_mode,
-        )
+        if collect_dags:
+            self.collect_dags(
+                dag_folder=dag_folder,
+                include_examples=include_examples,
+                safe_mode=safe_mode,
+            )
         # Should the extra operator link be loaded via plugins?
         # This flag is set to False in Scheduler so that Extra Operator links 
are not loaded
         self.load_op_links = load_op_links
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index f796156f8f..d581925791 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -831,8 +831,8 @@ def reserialize_dags(*, session: Session = NEW_SESSION) -> 
None:
     from airflow.models.serialized_dag import SerializedDagModel
 
     session.query(SerializedDagModel).delete(synchronize_session=False)
-    dagbag = DagBag()
-    dagbag.collect_dags(only_if_updated=False, safe_mode=False)
+    dagbag = DagBag(collect_dags=False)
+    dagbag.collect_dags(only_if_updated=False)
     dagbag.sync_to_db(session=session)
 
 
diff --git a/tests/models/test_dagbag.py b/tests/models/test_dagbag.py
index ed20d41ef0..f55a31400d 100644
--- a/tests/models/test_dagbag.py
+++ b/tests/models/test_dagbag.py
@@ -1060,3 +1060,16 @@ class TestDagBag:
         dagbag = DagBag(dag_folder=dag_file, include_examples=False)
         assert len(dagbag.dag_ids) == 0
         assert "has no tags" in dagbag.import_errors[dag_file]
+
+    def test_dagbag_dag_collection(self):
+
+        dagbag = DagBag(dag_folder=TEST_DAGS_FOLDER, include_examples=False, 
collect_dags=False)
+        # since collect_dags is False, dagbag.dags should be empty
+        assert not dagbag.dags
+
+        dagbag.collect_dags()
+        assert dagbag.dags
+
+        # test that dagbag.dags is not empty if collect_dags is True
+        dagbag = DagBag(dag_folder=TEST_DAGS_FOLDER, include_examples=False)
+        assert dagbag.dags