Use existing DagBag for 'dag_details' & trigger Endpoints (#8501)

apache · Jun 29, 2020 · 39459c9 · 39459c9
1 parent a1be32b
commit 39459c9
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git a/airflow/www_rbac/views.py b/airflow/www_rbac/views.py
@@ -562,9 +562,7 @@ def code(self, session=None):
     @provide_session
     def dag_details(self, session=None):
         dag_id = request.args.get('dag_id')
-        dag_orm = DagModel.get_dagmodel(dag_id, session=session)
-        # FIXME: items needed for this view should move to the database
-        dag = dag_orm.get_dag(STORE_SERIALIZED_DAGS)
+        dag = dagbag.get_dag(dag_id)
         title = "DAG details"
         root = request.args.get('root', '')
 

diff --git a/tests/www_rbac/test_views.py b/tests/www_rbac/test_views.py
@@ -565,6 +565,19 @@ def test_dag_details(self):
         resp = self.client.get(url, follow_redirects=True)
         self.check_content_in_response('DAG details', resp)
 
+    @parameterized.expand(["graph", "tree", "dag_details"])
+    @mock.patch('airflow.www_rbac.views.dagbag.get_dag')
+    def test_view_uses_existing_dagbag(self, endpoint, mock_get_dag):
+        """
+        Test that Graph, Tree & Dag Details View uses the DagBag already created in views.py
+        instead of creating a new one.
+        """
+        mock_get_dag.return_value = DAG(dag_id='example_bash_operator')
+        url = '{}?dag_id=example_bash_operator'.format(endpoint)
+        resp = self.client.get(url, follow_redirects=True)
+        mock_get_dag.assert_called_once_with('example_bash_operator')
+        self.check_content_in_response('example_bash_operator', resp)
+
     def test_dag_details_trigger_origin_tree_view(self):
         dag = self.dagbag.dags['test_tree_view']
         dag.create_dagrun(
@@ -2207,6 +2220,18 @@ def test_trigger_serialized_dag(self, mock_os_isfile, mock_dagrun):
         self.check_content_in_response(
             'Triggered example_bash_operator, it should start any moment now.', response)
 
+    @mock.patch('airflow.www_rbac.views.dagbag.get_dag')
+    def test_trigger_endpoint_uses_existing_dagbag(self, mock_get_dag):
+        """
+        Test that Trigger Endpoint uses the DagBag already created in views.py
+        instead of creating a new one.
+        """
+        mock_get_dag.return_value = DAG(dag_id='example_bash_operator')
+        url = 'trigger?dag_id=example_bash_operator'
+        resp = self.client.post(url, data={}, follow_redirects=True)
+        mock_get_dag.assert_called_once_with('example_bash_operator')
+        self.check_content_in_response('example_bash_operator', resp)
+
 
 class TestExtraLinks(TestBase):
     def setUp(self):