Repository: PrefectHQ/prefect

Issue:
Add configuration for all Prefect API calls to be proxied
### First check

- [X] I added a descriptive title to this issue.
- [X] I used the GitHub search to find a similar request and didn't find it.
- [X] I searched the Prefect documentation for this feature.

### Prefect Version

2.x

### Describe the current behavior

httpx honors setting the HTTPS_PROXY env var which proxies all traffic, which means that the Prefect API calls will be proxied, but it also means that private network traffic will also get proxied.


### Describe the proposed behavior


Adding a configuration option in the Prefect configuration for a proxy you can pass it through to httpx and aiohttp internally, meaning that only traffic to Prefect API will use the proxy.

The libraries mentioned in the docs, httpx and aiohttp, both provide support for specifying proxies ([here](https://www.python-httpx.org/advanced/#http-proxying) and [here](https://docs.aiohttp.org/en/stable/client_advanced.html#proxy-support)).

A similar feature for a slack notification handler was implemented here https://github.com/PrefectHQ/prefect/issues/5236.

### Example Use

In CLI `prefect config set HTTPS_PROXY=https://proxy.quackme.co` for example

### Additional context

_No response_


Discussion:
I think that `PREFECT_CLIENT_PROXY_URL` makes sense to me. Seems simple to add this to the `OrionClient` and `CloudClient` startup if set. I'd use `all://` unless there's a specific reason we'd only want to target HTTPS?
That sounds great, I like `all://`

Resolving patch:
diff --git a/src/prefect/client/orion.py b/src/prefect/client/orion.py
--- a/src/prefect/client/orion.py
+++ b/src/prefect/client/orion.py
@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
 from uuid import UUID
 
+import httpcore
 import httpx
 import pendulum
 import pydantic
@@ -126,9 +127,7 @@ def __init__(
                     keepalive_expiry=25,
                 ),
             )
-            # See https://www.python-httpx.org/advanced/#custom-transports
-            # `retries` specifies the number of retries on connection errors
-            httpx_settings.setdefault("transport", httpx.AsyncHTTPTransport(retries=3))
+
             # See https://www.python-httpx.org/http2/
             # Enabling HTTP/2 support on the client does not necessarily mean that your requests
             # and responses will be transported over HTTP/2, since both the client and the server
@@ -161,6 +160,28 @@ def __init__(
         self._client = PrefectHttpxClient(
             **httpx_settings,
         )
+
+        # See https://www.python-httpx.org/advanced/#custom-transports
+        # 
+        # If we're using an HTTP/S client (not the ephemeral client), adjust the
+        # transport to add retries _after_ it is instantiated. If we alter the transport
+        # before instantiation, the transport will not be aware of proxies unless we
+        # reproduce all of the logic to make it so.
+        #
+        # Only alter the transport to set our default of 3 retries, don't modify any
+        # transport a user may have provided via httpx_settings.
+        #
+        # Making liberal use of getattr and isinstance checks here to avoid any
+        # surprises if the internals of httpx or httpcore change on us
+        if isinstance(api, str) and not httpx_settings.get("transport"):
+            transport_for_url = getattr(self._client, "_transport_for_url", None)
+            if callable(transport_for_url):
+                orion_transport = transport_for_url(httpx.URL(api))
+                if isinstance(orion_transport, httpx.AsyncHTTPTransport):
+                    pool = getattr(orion_transport, "_pool", None)
+                    if isinstance(pool, httpcore.AsyncConnectionPool):
+                        pool._retries = 3
+
         self.logger = get_logger("client")
 
     @property


Test patch:
diff --git a/tests/client/test_orion_client.py b/tests/client/test_orion_client.py
--- a/tests/client/test_orion_client.py
+++ b/tests/client/test_orion_client.py
@@ -1,12 +1,14 @@
 import datetime
+import os
 import random
 import threading
 from datetime import datetime, timedelta, timezone
-from typing import List
+from typing import Generator, List
 from unittest.mock import MagicMock
 from uuid import UUID, uuid4
 
 import anyio
+import httpcore
 import httpx
 import pendulum
 import pytest
@@ -33,6 +35,7 @@
 from prefect.orion.schemas.states import StateType
 from prefect.settings import (
     PREFECT_API_KEY,
+    PREFECT_API_URL,
     PREFECT_ORION_DATABASE_MIGRATE_ON_START,
     temporary_settings,
 )
@@ -56,6 +59,107 @@ def test_get_client_cache_uses_profile_settings(self):
             assert new_client is not client
 
 
+class TestClientProxyAwareness:
+    """Regression test for https://github.com/PrefectHQ/nebula/issues/2356, where
+    a customer reported that the Cloud client supported proxies, but the Orion client
+    did not.  This test suite is implementation-specific to httpx/httpcore, as there are
+    no other inexpensive ways to confirm both the proxy-awareness and preserving the
+    retry behavior without probing into the implementation details of the libraries."""
+
+    @pytest.fixture()
+    def remote_https_orion(self) -> Generator[httpx.URL, None, None]:
+        orion_url = "https://127.0.0.1:4242/"
+        with temporary_settings(updates={PREFECT_API_URL: orion_url}):
+            yield httpx.URL(orion_url)
+
+    def test_unproxied_remote_client_will_retry(self, remote_https_orion: httpx.URL):
+        """The original issue here was that we were overriding the `transport` in
+        order to set the retries to 3; this is what circumvented the proxy support.
+        This test (and those below) should confirm that we are setting the retries on
+        the transport's pool in all cases."""
+        httpx_client = get_client()._client
+        assert isinstance(httpx_client, httpx.AsyncClient)
+
+        transport_for_orion = httpx_client._transport_for_url(remote_https_orion)
+        assert isinstance(transport_for_orion, httpx.AsyncHTTPTransport)
+
+        pool = transport_for_orion._pool
+        assert isinstance(pool, httpcore.AsyncConnectionPool)
+        assert pool._retries == 3  # set in prefect.client.orion.get_client()
+
+    def test_users_can_still_provide_transport(self, remote_https_orion: httpx.URL):
+        """If users want to supply an alternative transport, they still can and
+        we will not alter it"""
+        httpx_settings = {"transport": httpx.AsyncHTTPTransport(retries=11)}
+        httpx_client = get_client(httpx_settings)._client
+        assert isinstance(httpx_client, httpx.AsyncClient)
+
+        transport_for_orion = httpx_client._transport_for_url(remote_https_orion)
+        assert isinstance(transport_for_orion, httpx.AsyncHTTPTransport)
+
+        pool = transport_for_orion._pool
+        assert isinstance(pool, httpcore.AsyncConnectionPool)
+        assert pool._retries == 11  # not overridden by get_client() in this case
+
+    @pytest.fixture
+    def https_proxy(self) -> Generator[httpcore.URL, None, None]:
+        original = os.environ.get("HTTPS_PROXY")
+        try:
+            os.environ["HTTPS_PROXY"] = "https://127.0.0.1:6666"
+            yield httpcore.URL(os.environ["HTTPS_PROXY"])
+        finally:
+            if original is None:
+                del os.environ["HTTPS_PROXY"]
+            else:
+                os.environ["HTTPS_PROXY"] = original
+
+    async def test_client_is_aware_of_https_proxy(
+        self, remote_https_orion: httpx.URL, https_proxy: httpcore.URL
+    ):
+        httpx_client = get_client()._client
+        assert isinstance(httpx_client, httpx.AsyncClient)
+
+        transport_for_orion = httpx_client._transport_for_url(remote_https_orion)
+        assert isinstance(transport_for_orion, httpx.AsyncHTTPTransport)
+
+        pool = transport_for_orion._pool
+        assert isinstance(pool, httpcore.AsyncHTTPProxy)
+        assert pool._proxy_url == https_proxy
+        assert pool._retries == 3  # set in prefect.client.orion.get_client()
+
+    @pytest.fixture()
+    def remote_http_orion(self) -> Generator[httpx.URL, None, None]:
+        orion_url = "http://127.0.0.1:4242/"
+        with temporary_settings(updates={PREFECT_API_URL: orion_url}):
+            yield httpx.URL(orion_url)
+
+    @pytest.fixture
+    def http_proxy(self) -> Generator[httpcore.URL, None, None]:
+        original = os.environ.get("HTTP_PROXY")
+        try:
+            os.environ["HTTP_PROXY"] = "http://127.0.0.1:6666"
+            yield httpcore.URL(os.environ["HTTP_PROXY"])
+        finally:
+            if original is None:
+                del os.environ["HTTP_PROXY"]
+            else:
+                os.environ["HTTP_PROXY"] = original
+
+    async def test_client_is_aware_of_http_proxy(
+        self, remote_http_orion: httpx.URL, http_proxy: httpcore.URL
+    ):
+        httpx_client = get_client()._client
+        assert isinstance(httpx_client, httpx.AsyncClient)
+
+        transport_for_orion = httpx_client._transport_for_url(remote_http_orion)
+        assert isinstance(transport_for_orion, httpx.AsyncHTTPTransport)
+
+        pool = transport_for_orion._pool
+        assert isinstance(pool, httpcore.AsyncHTTPProxy)
+        assert pool._proxy_url == http_proxy
+        assert pool._retries == 3  # set in prefect.client.orion.get_client()
+
+
 class TestInjectClient:
     @staticmethod
     @inject_client