diff --git a/airflow/www/utils.py b/airflow/www/utils.py index 1bbc0936b3977..0c4f4b05d6128 100644 --- a/airflow/www/utils.py +++ b/airflow/www/utils.py @@ -56,8 +56,13 @@ def should_hide_value_for_key(key_name): - return any(s in key_name.lower() for s in DEFAULT_SENSITIVE_VARIABLE_FIELDS) \ - and configuration.conf.getboolean('admin', 'hide_sensitive_variable_fields') + # It is possible via importing variables from file that a key is empty. + if key_name: + config_set = configuration.conf.getboolean('admin', + 'hide_sensitive_variable_fields') + field_comp = any(s in key_name.lower() for s in DEFAULT_SENSITIVE_VARIABLE_FIELDS) + return config_set and field_comp + return False class LoginMixin(object): diff --git a/airflow/www/views.py b/airflow/www/views.py index 3e41d2d02a1b9..0c0dcff801a59 100644 --- a/airflow/www/views.py +++ b/airflow/www/views.py @@ -2012,9 +2012,20 @@ def varimport(self): except Exception as e: flash("Missing file or syntax error: {}.".format(e)) else: + suc_count = fail_count = 0 for k, v in d.items(): - models.Variable.set(k, v, serialize_json=isinstance(v, dict)) - flash("{} variable(s) successfully updated.".format(len(d))) + try: + models.Variable.set(k, v, serialize_json=isinstance(v, dict)) + except Exception as e: + logging.info('Variable import failed: {}'.format(repr(e))) + fail_count += 1 + else: + suc_count += 1 + flash("{} variable(s) successfully updated.".format(suc_count), 'info') + if fail_count: + flash( + "{} variables(s) failed to be updated.".format(fail_count), 'error') + return redirect('/admin/variable') diff --git a/airflow/www_rbac/utils.py b/airflow/www_rbac/utils.py index 7bbdada5556f2..a0e9258eae30e 100644 --- a/airflow/www_rbac/utils.py +++ b/airflow/www_rbac/utils.py @@ -54,8 +54,13 @@ def should_hide_value_for_key(key_name): - return any(s in key_name.lower() for s in DEFAULT_SENSITIVE_VARIABLE_FIELDS) \ - and configuration.getboolean('admin', 'hide_sensitive_variable_fields') + # It is possible via importing variables from file that a key is empty. + if key_name: + config_set = configuration.conf.getboolean('admin', + 'hide_sensitive_variable_fields') + field_comp = any(s in key_name.lower() for s in DEFAULT_SENSITIVE_VARIABLE_FIELDS) + return config_set and field_comp + return False def get_params(**kwargs): diff --git a/airflow/www_rbac/views.py b/airflow/www_rbac/views.py index a9947ae096be3..629f488fc7218 100644 --- a/airflow/www_rbac/views.py +++ b/airflow/www_rbac/views.py @@ -2053,9 +2053,18 @@ def varimport(self): except Exception: flash("Missing file or syntax error.") else: + suc_count = fail_count = 0 for k, v in d.items(): - models.Variable.set(k, v, serialize_json=isinstance(v, dict)) - flash("{} variable(s) successfully updated.".format(len(d))) + try: + models.Variable.set(k, v, serialize_json=isinstance(v, dict)) + except Exception as e: + logging.info('Variable import failed: {}'.format(repr(e))) + fail_count += 1 + else: + suc_count += 1 + flash("{} variable(s) successfully updated.".format(suc_count), 'info') + if fail_count: + flash("{} variables(s) failed to be updated.".format(fail_count), 'error') self.update_redirect() return redirect(self.get_redirect()) diff --git a/tests/www/test_utils.py b/tests/www/test_utils.py index 891298c0a9d8e..9034b8b5fd28e 100644 --- a/tests/www/test_utils.py +++ b/tests/www/test_utils.py @@ -32,6 +32,10 @@ class UtilsTest(unittest.TestCase): def setUp(self): super(UtilsTest, self).setUp() + def test_empty_variable_should_not_be_hidden(self): + self.assertFalse(utils.should_hide_value_for_key("")) + self.assertFalse(utils.should_hide_value_for_key(None)) + def test_normal_variable_should_not_be_hidden(self): self.assertFalse(utils.should_hide_value_for_key("key")) diff --git a/tests/www/test_views.py b/tests/www/test_views.py index bd385fc5696c6..fb7a2f3411870 100644 --- a/tests/www/test_views.py +++ b/tests/www/test_views.py @@ -20,6 +20,7 @@ import io import copy import logging.config +import mock import os import shutil import tempfile @@ -448,6 +449,28 @@ def tearDownClass(cls): session.close() super(TestVarImportView, cls).tearDownClass() + def test_import_variable_fail(self): + with mock.patch('airflow.models.Variable.set') as set_mock: + set_mock.side_effect = UnicodeEncodeError + content = '{"fail_key": "fail_val"}' + + try: + # python 3+ + bytes_content = io.BytesIO(bytes(content, encoding='utf-8')) + except TypeError: + # python 2.7 + bytes_content = io.BytesIO(bytes(content)) + response = self.app.post( + self.IMPORT_ENDPOINT, + data={'file': (bytes_content, 'test.json')}, + follow_redirects=True + ) + self.assertEqual(response.status_code, 200) + session = Session() + db_dict = {x.key: x.get_val() for x in session.query(models.Variable).all()} + session.close() + self.assertNotIn('fail_key', db_dict) + def test_import_variables(self): content = ('{"str_key": "str_value", "int_key": 60,' '"list_key": [1, 2], "dict_key": {"k_a": 2, "k_b": 3}}') @@ -463,21 +486,24 @@ def test_import_variables(self): follow_redirects=True ) self.assertEqual(response.status_code, 200) - body = response.data.decode('utf-8') - self.assertIn('str_key', body) - self.assertIn('int_key', body) - self.assertIn('list_key', body) - self.assertIn('dict_key', body) - self.assertIn('str_value', body) - self.assertIn('60', body) - self.assertIn('[1, 2]', body) - # As dicts are not ordered, we may get any of the following cases. - case_a_dict = '{"k_a": 2, "k_b": 3}' - case_b_dict = '{"k_b": 3, "k_a": 2}' + session = Session() + # Extract values from Variable + db_dict = {x.key: x.get_val() for x in session.query(models.Variable).all()} + session.close() + self.assertIn('str_key', db_dict) + self.assertIn('int_key', db_dict) + self.assertIn('list_key', db_dict) + self.assertIn('dict_key', db_dict) + self.assertEquals('str_value', db_dict['str_key']) + self.assertEquals('60', db_dict['int_key']) + self.assertEquals('[1, 2]', db_dict['list_key']) + + case_a_dict = '{"k_a": 2, "k_b": 3}' + case_b_dict = '{"k_b": 3, "k_a": 2}' try: - self.assertIn(case_a_dict, body) + self.assertEquals(case_a_dict, db_dict['dict_key']) except AssertionError: - self.assertIn(case_b_dict, body) + self.assertEquals(case_b_dict, db_dict['dict_key']) class TestMountPoint(unittest.TestCase): diff --git a/tests/www_rbac/test_utils.py b/tests/www_rbac/test_utils.py index 05114881ddecf..1879ba082693d 100644 --- a/tests/www_rbac/test_utils.py +++ b/tests/www_rbac/test_utils.py @@ -28,6 +28,10 @@ class UtilsTest(unittest.TestCase): def setUp(self): super(UtilsTest, self).setUp() + def test_empty_variable_should_not_be_hidden(self): + self.assertFalse(utils.should_hide_value_for_key("")) + self.assertFalse(utils.should_hide_value_for_key(None)) + def test_normal_variable_should_not_be_hidden(self): self.assertFalse(utils.should_hide_value_for_key("key")) diff --git a/tests/www_rbac/test_views.py b/tests/www_rbac/test_views.py index b4a2f0214248d..f6a350150405b 100644 --- a/tests/www_rbac/test_views.py +++ b/tests/www_rbac/test_views.py @@ -21,6 +21,7 @@ import io import json import logging.config +import mock import os import shutil import sys @@ -172,6 +173,25 @@ def test_xss_prevention(self): self.assertNotIn("", resp.data.decode("utf-8")) + def test_import_variables(self): + content = '{"str_key": "str_value"}' + + with mock.patch('airflow.models.Variable.set') as set_mock: + set_mock.side_effect = UnicodeEncodeError + self.assertEqual(self.session.query(models.Variable).count(), 0) + + try: + # python 3+ + bytes_content = io.BytesIO(bytes(content, encoding='utf-8')) + except TypeError: + # python 2.7 + bytes_content = io.BytesIO(bytes(content)) + + resp = self.client.post('/variable/varimport', + data={'file': (bytes_content, 'test.json')}, + follow_redirects=True) + self.check_content_in_response('1 variable(s) failed to be updated.', resp) + def test_import_variables(self): self.assertEqual(self.session.query(models.Variable).count(), 0)