def test_read_dags(self):
"""DAGs can be read from database."""
example_dags = self._write_example_dags()
serialized_dags = SDM.read_all_dags()
self.assertTrue(len(example_dags) == len(serialized_dags))
for dag_id, dag in example_dags.items():
serialized_dag = serialized_dags[dag_id]
self.assertTrue(serialized_dag.dag_id == dag.dag_id)
self.assertTrue(set(serialized_dag.task_dict) == set(dag.task_dict))
def collect_dags_from_db(self):
"""Collects DAGs from database."""
from airflow.models.serialized_dag import SerializedDagModel
with Stats.timer('collect_db_dags'):
self.log.info("Filling up the DagBag from database")
# The dagbag contains all rows in serialized_dag table. Deleted DAGs are deleted
# from the table by the scheduler job.
self.dags = SerializedDagModel.read_all_dags()
# Adds subdags.
# DAG post-processing steps such as self.bag_dag and croniter are not needed as
# they are done by scheduler before serialization.
subdags = {}
for dag in self.dags.values():
for subdag in dag.subdags:
subdags[subdag.dag_id] = subdag
self.dags.update(subdags)