[AIRFLOW-2786] Gracefully handle Variable import errors (#3648)

Variables that are added through a file are not
checked as explicity as creating a Variable in the
web UI. This handles exceptions that could be caused
by improper keys or values.

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):

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')
 
 

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):

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())
 

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"))
 

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):

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"))
 

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("<img src='' onerror='alert(1);'>",
                          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)