From 325734a166cee95ef4514bbbba3242e09bfe41ff Mon Sep 17 00:00:00 2001 From: JCoder01 Date: Sun, 2 Feb 2020 04:24:40 -0500 Subject: [PATCH] [AIRFLOW-5843] Add conf option to Add DAG Run view (#7281) --- airflow/www/forms.py | 9 ++++++++ airflow/www/validators.py | 22 +++++++++++++++++++ airflow/www/views.py | 2 +- tests/www/test_validators.py | 42 ++++++++++++++++++++++++++++++++++++ tests/www/test_views.py | 35 +++++++++++++++++++++++++++++- 5 files changed, 108 insertions(+), 2 deletions(-) diff --git a/airflow/www/forms.py b/airflow/www/forms.py index af75cccad5444..fbcb5649795dc 100644 --- a/airflow/www/forms.py +++ b/airflow/www/forms.py @@ -17,6 +17,8 @@ # specific language governing permissions and limitations # under the License. +import json + from flask_appbuilder.fieldwidgets import ( BS3PasswordFieldWidget, BS3TextAreaFieldWidget, BS3TextFieldWidget, DateTimePickerWidget, Select2Widget, ) @@ -30,6 +32,7 @@ from airflow.models import Connection from airflow.utils import timezone +from airflow.www.validators import ValidJson class DateTimeForm(FlaskForm): @@ -81,12 +84,18 @@ class DagRunForm(DynamicForm): widget=DateTimePickerWidget()) external_trigger = BooleanField( lazy_gettext('External Trigger')) + conf = TextAreaField( + lazy_gettext('Conf'), + validators=[ValidJson(), validators.Optional()], + widget=BS3TextAreaFieldWidget()) def populate_obj(self, item): # TODO: This is probably better done as a custom field type so we can # set TZ at parse time super().populate_obj(item) item.execution_date = timezone.make_aware(item.execution_date) + if item.conf: + item.conf = json.loads(item.conf) class ConnectionForm(DynamicForm): diff --git a/airflow/www/validators.py b/airflow/www/validators.py index ae20e450218ee..077894b5495d7 100644 --- a/airflow/www/validators.py +++ b/airflow/www/validators.py @@ -17,6 +17,8 @@ # specific language governing permissions and limitations # under the License. +import json + from wtforms.validators import EqualTo, ValidationError @@ -56,3 +58,23 @@ def __call__(self, form, field): message = message % d raise ValidationError(message) + + +class ValidJson(object): + """Validates data is valid JSON. + + :param message: + Error message to raise in case of a validation error. + """ + def __init__(self, message=None): + self.message = message + + def __call__(self, form, field): + if field.data: + try: + json.loads(field.data) + except Exception as ex: + message = self.message or 'JSON Validation Error: {}'.format(ex) + raise ValidationError( + message=field.gettext(message.format(field.data)) + ) diff --git a/airflow/www/views.py b/airflow/www/views.py index 3cb3a4db078b9..9f91686bf33ca 100644 --- a/airflow/www/views.py +++ b/airflow/www/views.py @@ -2434,7 +2434,7 @@ class DagRunModelView(AirflowModelView): base_permissions = ['can_list', 'can_add'] - add_columns = ['state', 'dag_id', 'execution_date', 'run_id', 'external_trigger'] + add_columns = ['state', 'dag_id', 'execution_date', 'run_id', 'external_trigger', 'conf'] list_columns = ['state', 'dag_id', 'execution_date', 'run_id', 'external_trigger'] search_columns = ['state', 'dag_id', 'execution_date', 'run_id', 'external_trigger'] diff --git a/tests/www/test_validators.py b/tests/www/test_validators.py index e60a30c180f9d..a166756d614ae 100644 --- a/tests/www/test_validators.py +++ b/tests/www/test_validators.py @@ -92,5 +92,47 @@ def test_validation_raises_custom_message(self): ) +class TestValidJson(unittest.TestCase): + + def setUp(self): + super().setUp() + self.form_field_mock = mock.MagicMock(data='{"valid":"True"}') + self.form_field_mock.gettext.side_effect = lambda msg: msg + self.form_mock = mock.MagicMock(spec_set=dict) + + def _validate(self, message=None): + + validator = validators.ValidJson(message=message) + + return validator(self.form_mock, self.form_field_mock) + + def test_form_field_is_none(self): + self.form_field_mock.data = None + + self.assertIsNone(self._validate()) + + def test_validation_pass(self): + self.assertIsNone(self._validate()) + + def test_validation_raises_default_message(self): + self.form_field_mock.data = '2017-05-04' + + self.assertRaisesRegex( + validators.ValidationError, + "JSON Validation Error:.*", + self._validate, + ) + + def test_validation_raises_custom_message(self): + self.form_field_mock.data = '2017-05-04' + + self.assertRaisesRegex( + validators.ValidationError, + "Invalid JSON", + self._validate, + message="Invalid JSON: {}", + ) + + if __name__ == '__main__': unittest.main() diff --git a/tests/www/test_views.py b/tests/www/test_views.py index b9f6f488fa753..8b6763c81bd4a 100644 --- a/tests/www/test_views.py +++ b/tests/www/test_views.py @@ -2130,7 +2130,7 @@ def test_create_dagrun(self): "state": "running", "dag_id": "example_bash_operator", "execution_date": "2018-07-06 05:04:03", - "run_id": "manual_abc", + "run_id": "test_create_dagrun", } resp = self.client.post('/dagrun/add', data=data, @@ -2141,6 +2141,39 @@ def test_create_dagrun(self): self.assertEqual(dr.execution_date, timezone.convert_to_utc(datetime(2018, 7, 6, 5, 4, 3))) + def test_create_dagrun_valid_conf(self): + conf_value = dict(Valid=True) + data = { + "state": "running", + "dag_id": "example_bash_operator", + "execution_date": "2018-07-06 05:05:03", + "run_id": "test_create_dagrun_valid_conf", + "conf": json.dumps(conf_value) + } + + resp = self.client.post('/dagrun/add', + data=data, + follow_redirects=True) + self.check_content_in_response('Added Row', resp) + dr = self.session.query(models.DagRun).one() + self.assertEqual(dr.conf, conf_value) + + def test_create_dagrun_invalid_conf(self): + data = { + "state": "running", + "dag_id": "example_bash_operator", + "execution_date": "2018-07-06 05:06:03", + "run_id": "test_create_dagrun_invalid_conf", + "conf": "INVALID: [JSON" + } + + resp = self.client.post('/dagrun/add', + data=data, + follow_redirects=True) + self.check_content_in_response('JSON Validation Error:', resp) + dr = self.session.query(models.DagRun).all() + self.assertFalse(dr) + class TestDecorators(TestBase): EXAMPLE_DAG_DEFAULT_DATE = dates.days_ago(2)