diff --git a/airflow/www/utils.py b/airflow/www/utils.py index dbf326608c879..40312e02b547e 100644 --- a/airflow/www/utils.py +++ b/airflow/www/utils.py @@ -31,7 +31,6 @@ from airflow.configuration import conf from airflow.models.baseoperator import BaseOperator -from airflow.models.dag import DagBag, DagModel from airflow.operators.subdag_operator import SubDagOperator from airflow.utils import timezone from airflow.utils.code_utils import get_python_source @@ -425,19 +424,3 @@ def is_utcdatetime(self, col_name): (('is_utcdatetime', DateTimeField, AirflowDateTimePickerWidget),) + FieldConverter.conversion_table ) - - -def get_dag(orm_dag: DagModel, store_serialized_dags=False): - """Creates a dagbag to load and return a DAG. - - Calling it from UI should set store_serialized_dags = STORE_SERIALIZED_DAGS. - There may be a delay for scheduler to write serialized DAG into database, - loads from file in this case. - FIXME: remove it when webserver does not access to DAG folder in future. - """ - dag = DagBag( - dag_folder=orm_dag.fileloc, store_serialized_dags=store_serialized_dags).get_dag(orm_dag.dag_id) - if store_serialized_dags and dag is None: - dag = DagBag( - dag_folder=orm_dag.fileloc, store_serialized_dags=False).get_dag(orm_dag.dag_id) - return dag diff --git a/airflow/www/views.py b/airflow/www/views.py index d9b060c288fef..5c9684462a40b 100644 --- a/airflow/www/views.py +++ b/airflow/www/views.py @@ -77,7 +77,6 @@ from airflow.www.forms import ( ConnectionForm, DagRunForm, DateTimeForm, DateTimeWithNumRunsForm, DateTimeWithNumRunsWithDagRunsForm, ) -from airflow.www.utils import get_dag from airflow.www.widgets import AirflowModelListWidget PAGE_SIZE = conf.getint('webserver', 'page_size') @@ -575,9 +574,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 = get_dag(dag_orm, STORE_SERIALIZED_DAGS) + dag = dagbag.get_dag(dag_id) title = "DAG details" root = request.args.get('root', '') @@ -1057,7 +1054,7 @@ def trigger(self, session=None): conf=conf ) - dag = get_dag(dag_orm, STORE_SERIALIZED_DAGS) + dag = dagbag.get_dag(dag_id) dag.create_dagrun( run_id=run_id, execution_date=execution_date, diff --git a/tests/www/test_views.py b/tests/www/test_views.py index ab288b5f4d321..00e23a1dd48a8 100644 --- a/tests/www/test_views.py +++ b/tests/www/test_views.py @@ -532,6 +532,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.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 = f'{endpoint}?dag_id=example_bash_operator' + 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( @@ -2208,6 +2221,18 @@ def test_trigger_dag_form(self): self.assertEqual(resp.status_code, 200) self.check_content_in_response('Trigger DAG: {}'.format(test_dag_id), resp) + @mock.patch('airflow.www.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):