fix: Update custom connection field processing (#20883)

* fix: Update custom connection field processing Fixes issue where custom connectionfields are not updated because `extra` field is in form and has previous values, overriding custom field values. Adds portion of connection form tests to test functionality. (cherry picked from commit 44df142)
apache · Mar 22, 2022 · 6a86a5b · 6a86a5b
1 parent e92201d
commit 6a86a5b
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 11 deletions.
diff --git a/airflow/www/views.py b/airflow/www/views.py
@@ -3518,19 +3518,17 @@ def process_form(self, form, is_created):
         """Process form data."""
         conn_type = form.data['conn_type']
         conn_id = form.data["conn_id"]
-        extra = {
-            key: form.data[key]
-            for key in self.extra_fields
-            if key in form.data and key.startswith(f"extra__{conn_type}__")
-        }
 
-        # If parameters are added to the classic `Extra` field, include these values along with
-        # custom-field extras.
-        extra_conn_params = form.data.get("extra")
+        # The extra value is the combination of custom fields for this conn_type and the Extra field.
+        # The extra form field with all extra values (including custom fields) is in the form being processed
+        # so we start with those values, and override them with anything in the custom fields.
+        extra = {}
+
+        extra_field = form.data.get("extra")
 
-        if extra_conn_params:
+        if extra_field:
             try:
-                extra.update(json.loads(extra_conn_params))
+                extra.update(json.loads(extra_field))
             except (JSONDecodeError, TypeError):
                 flash(
                     Markup(
@@ -3539,11 +3537,19 @@ def process_form(self, form, is_created):
                         "<p>If connection parameters need to be added to <em>Extra</em>, "
                         "please make sure they are in the form of a single, valid JSON object.</p><br>"
                         "The following <em>Extra</em> parameters were <b>not</b> added to the connection:<br>"
-                        f"{extra_conn_params}",
+                        f"{extra_field}",
                     ),
                     category="error",
                 )
 
+        custom_fields = {
+            key: form.data[key]
+            for key in self.extra_fields
+            if key in form.data and key.startswith(f"extra__{conn_type}__")
+        }
+
+        extra.update(custom_fields)
+
         if extra.keys():
             form.extra.data = json.dumps(extra)
 

diff --git a/tests/www/views/test_views_connection.py b/tests/www/views/test_views_connection.py
@@ -107,6 +107,21 @@ def test_process_form_extras():
 
     assert json.loads(mock_form.extra.data) == {"extra__test3__custom_field": "custom_field_val3"}
 
+    # Testing parameters set in both extra and custom fields (cunnection updates).
+    mock_form = mock.Mock()
+    mock_form.data = {
+        "conn_type": "test4",
+        "conn_id": "extras_test4",
+        "extra": '{"extra__test4__custom_field": "custom_field_val3"}',
+        "extra__test4__custom_field": "custom_field_val4",
+    }
+
+    cmv = ConnectionModelView()
+    cmv.extra_fields = ["extra__test4__custom_field"]  # Custom field
+    cmv.process_form(form=mock_form, is_created=True)
+
+    assert json.loads(mock_form.extra.data) == {"extra__test4__custom_field": "custom_field_val4"}
+
 
 def test_duplicate_connection(admin_client):
     """Test Duplicate multiple connection with suffix"""