[airflow] branch main updated: Handle connection parameters added to Extra and custom fields (#17269)

[ur...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 1941f94  Handle connection parameters added to Extra and custom fields (#17269)
1941f94 is described below

commit 1941f9486e72b9c70654ea9aa285d566239f6ba1
Author: josh-fell <48...@users.noreply.github.com>
AuthorDate: Sat Jul 31 01:35:12 2021 -0400

    Handle connection parameters added to Extra and custom fields (#17269)
---
 airflow/www/views.py                     | 22 ++++++++++++++
 tests/www/views/test_views_connection.py | 52 ++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/airflow/www/views.py b/airflow/www/views.py
index 8309d99..3769c1a 100644
--- a/airflow/www/views.py
+++ b/airflow/www/views.py
@@ -3179,11 +3179,33 @@ class ConnectionModelView(AirflowModelView):
     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")
+
+        if extra_conn_params:
+            try:
+                extra.update(json.loads(extra_conn_params))
+            except (JSONDecodeError, TypeError):
+                flash(
+                    Markup(
+                        "<p>The <em>Extra</em> connection field contained an invalid value for Conn ID: "
+                        f"<q>{conn_id}</q>.</p>"
+                        "<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}",
+                    ),
+                    category="error",
+                )
+
         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
index 5576977..249bf2a 100644
--- a/tests/www/views/test_views_connection.py
+++ b/tests/www/views/test_views_connection.py
@@ -15,6 +15,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import json
 from unittest import mock
 
 import pytest
@@ -56,6 +57,57 @@ def test_prefill_form_null_extra():
     cmv.prefill_form(form=mock_form, pk=1)
 
 
+def test_process_form_extras():
+    """
+    Test the handling of connection parameters set with the classic `Extra` field as well as custom fields.
+    """
+
+    # Testing parameters set in both `Extra` and custom fields.
+    mock_form = mock.Mock()
+    mock_form.data = {
+        "conn_type": "test",
+        "conn_id": "extras_test",
+        "extra": '{"param1": "param1_val"}',
+        "extra__test__custom_field": "custom_field_val",
+    }
+
+    cmv = ConnectionModelView()
+    cmv.extra_fields = ["extra__test__custom_field"]  # Custom field
+    cmv.process_form(form=mock_form, is_created=True)
+
+    assert json.loads(mock_form.extra.data) == {
+        "extra__test__custom_field": "custom_field_val",
+        "param1": "param1_val",
+    }
+
+    # Testing parameters set in `Extra` field only.
+    mock_form = mock.Mock()
+    mock_form.data = {
+        "conn_type": "test2",
+        "conn_id": "extras_test2",
+        "extra": '{"param2": "param2_val"}',
+    }
+
+    cmv = ConnectionModelView()
+    cmv.process_form(form=mock_form, is_created=True)
+
+    assert json.loads(mock_form.extra.data) == {"param2": "param2_val"}
+
+    # Testing parameters set in custom fields only.
+    mock_form = mock.Mock()
+    mock_form.data = {
+        "conn_type": "test3",
+        "conn_id": "extras_test3",
+        "extra__test3__custom_field": "custom_field_val3",
+    }
+
+    cmv = ConnectionModelView()
+    cmv.extra_fields = ["extra__test3__custom_field"]  # Custom field
+    cmv.process_form(form=mock_form, is_created=True)
+
+    assert json.loads(mock_form.extra.data) == {"extra__test3__custom_field": "custom_field_val3"}
+
+
 def test_duplicate_connection(admin_client):
     """Test Duplicate multiple connection with suffix"""
     conn1 = Connection(