From 2309b7e825f5663b4952aa912ea030d5171188ff Mon Sep 17 00:00:00 2001 From: David Blain Date: Mon, 22 Apr 2024 13:18:14 +0200 Subject: [PATCH 1/3] refactor: Allow trust env parameter to be defined in extra options field of HTTP Connection --- airflow/providers/http/hooks/http.py | 4 ++++ tests/providers/http/hooks/test_http.py | 21 +++++++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/airflow/providers/http/hooks/http.py b/airflow/providers/http/hooks/http.py index 6c1d40f9b2662..76a2186151c91 100644 --- a/airflow/providers/http/hooks/http.py +++ b/airflow/providers/http/hooks/http.py @@ -132,6 +132,7 @@ def get_conn(self, headers: dict[Any, Any] | None = None) -> requests.Session: session.verify = extra.pop("verify", extra.pop("verify_ssl", True)) session.cert = extra.pop("cert", None) session.max_redirects = extra.pop("max_redirects", DEFAULT_REDIRECT_LIMIT) + session.trust_env = extra.pop("trust_env", True) try: session.headers.update(extra) @@ -425,6 +426,7 @@ def _process_extra_options_from_connection(cls, conn: Connection, extra_options: verify_ssl = extra.pop("verify", extra.pop("verify_ssl", None)) allow_redirects = extra.pop("allow_redirects", None) max_redirects = extra.pop("max_redirects", None) + trust_env = extra.pop("trust_env", None) if proxies is not None and "proxy" not in extra_options: extra_options["proxy"] = proxies @@ -436,6 +438,8 @@ def _process_extra_options_from_connection(cls, conn: Connection, extra_options: extra_options["allow_redirects"] = allow_redirects if max_redirects is not None and "max_redirects" not in extra_options: extra_options["max_redirects"] = max_redirects + if trust_env is not None and "trust_env" not in extra_options: + extra_options["trust_env"] = trust_env return extra def _retryable_error_async(self, exception: ClientResponseError) -> bool: diff --git a/tests/providers/http/hooks/test_http.py b/tests/providers/http/hooks/test_http.py index f82189a559713..7ede8cd0192e9 100644 --- a/tests/providers/http/hooks/test_http.py +++ b/tests/providers/http/hooks/test_http.py @@ -138,6 +138,7 @@ def test_hook_ignore_max_redirects_from_extra_field_as_header(self): assert conn.verify is True assert conn.cert is None assert conn.max_redirects == 3 + assert conn.trust_env is True def test_hook_ignore_proxies_from_extra_field_as_header(self): airflow_connection = get_airflow_connection_with_extra( @@ -154,6 +155,7 @@ def test_hook_ignore_proxies_from_extra_field_as_header(self): assert conn.verify is True assert conn.cert is None assert conn.max_redirects == DEFAULT_REDIRECT_LIMIT + assert conn.trust_env is True def test_hook_ignore_verify_from_extra_field_as_header(self): airflow_connection = get_airflow_connection_with_extra(extra={"bearer": "test", "verify": False}) @@ -168,6 +170,7 @@ def test_hook_ignore_verify_from_extra_field_as_header(self): assert conn.verify is False assert conn.cert is None assert conn.max_redirects == DEFAULT_REDIRECT_LIMIT + assert conn.trust_env is True def test_hook_ignore_cert_from_extra_field_as_header(self): airflow_connection = get_airflow_connection_with_extra(extra={"bearer": "test", "cert": "cert.crt"}) @@ -182,6 +185,22 @@ def test_hook_ignore_cert_from_extra_field_as_header(self): assert conn.verify is True assert conn.cert == "cert.crt" assert conn.max_redirects == DEFAULT_REDIRECT_LIMIT + assert conn.trust_env is True + + def test_hook_ignore_trust_env_from_extra_field_as_header(self): + airflow_connection = get_airflow_connection_with_extra(extra={"bearer": "test", "trust_env": False}) + with mock.patch("airflow.hooks.base.BaseHook.get_connection", side_effect=airflow_connection): + expected_conn = airflow_connection() + conn = self.get_hook.get_conn() + assert dict(conn.headers, **json.loads(expected_conn.extra)) != conn.headers + assert conn.headers.get("bearer") == "test" + assert conn.headers.get("cert") is None + assert conn.proxies == {} + assert conn.stream is False + assert conn.verify is True + assert conn.cert is None + assert conn.max_redirects == DEFAULT_REDIRECT_LIMIT + assert conn.trust_env is False @mock.patch("requests.Request") def test_hook_with_method_in_lowercase(self, mock_requests): @@ -616,6 +635,7 @@ async def test_async_request_uses_connection_extra_with_requests_parameters(self "verify": False, "allow_redirects": False, "max_redirects": 3, + "trust_env": False }, } ) @@ -633,6 +653,7 @@ async def test_async_request_uses_connection_extra_with_requests_parameters(self assert mocked_function.call_args.kwargs.get("verify_ssl") is False assert mocked_function.call_args.kwargs.get("allow_redirects") is False assert mocked_function.call_args.kwargs.get("max_redirects") == 3 + assert mocked_function.call_args.kwargs.get("trust_env") is False def test_process_extra_options_from_connection(self): extra_options = {} From 18c91c8a379f1e103224c10dfe94448e11f45b5b Mon Sep 17 00:00:00 2001 From: David Blain Date: Mon, 22 Apr 2024 13:54:03 +0200 Subject: [PATCH 2/3] refactor: Added missing comma --- tests/providers/http/hooks/test_http.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/http/hooks/test_http.py b/tests/providers/http/hooks/test_http.py index 7ede8cd0192e9..e3c7723ea9f16 100644 --- a/tests/providers/http/hooks/test_http.py +++ b/tests/providers/http/hooks/test_http.py @@ -635,7 +635,7 @@ async def test_async_request_uses_connection_extra_with_requests_parameters(self "verify": False, "allow_redirects": False, "max_redirects": 3, - "trust_env": False + "trust_env": False, }, } ) From c9c94a3f30c5cbd75151a78582f16f1b5b16f238 Mon Sep 17 00:00:00 2001 From: David Blain Date: Mon, 22 Apr 2024 19:13:26 +0200 Subject: [PATCH 3/3] refactor: Added special case when trust_env set to False --- tests/providers/http/hooks/test_http.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/providers/http/hooks/test_http.py b/tests/providers/http/hooks/test_http.py index e3c7723ea9f16..5b631df8a4d60 100644 --- a/tests/providers/http/hooks/test_http.py +++ b/tests/providers/http/hooks/test_http.py @@ -668,6 +668,7 @@ def test_process_extra_options_from_connection(self): "verify": False, "allow_redirects": False, "max_redirects": 3, + "trust_env": False, } )() @@ -679,6 +680,7 @@ def test_process_extra_options_from_connection(self): "verify_ssl": False, "allow_redirects": False, "max_redirects": 3, + "trust_env": False, } assert actual == {"bearer": "test"}