Skip to content

Commit

Permalink
fix Sync hook used as async hook
Browse files Browse the repository at this point in the history
Signed-off-by: Shi, Stone <[email protected]>
  • Loading branch information
shijiadong2022 committed Jul 27, 2024
1 parent 4f98519 commit 3ff1113
Showing 1 changed file with 13 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -731,7 +731,9 @@ def _instrument(self, **kwargs):
self._original_async_client = httpx.AsyncClient
request_hook = kwargs.get("request_hook")
response_hook = kwargs.get("response_hook")
async_request_hook = kwargs.get("async_request_hook", request_hook)
async_request_hook = kwargs.get(
"async_request_hook", self._wrap_async_request_hook(request_hook)
)
async_response_hook = kwargs.get("async_response_hook", response_hook)
if callable(request_hook):
_InstrumentedClient._request_hook = request_hook
Expand All @@ -749,6 +751,16 @@ def _instrument(self, **kwargs):
httpx.Client = httpx._api.Client = _InstrumentedClient
httpx.AsyncClient = _InstrumentedAsyncClient

# Wrap a given request hook function and ensure it is asynchronous
def _wrap_async_request_hook(self, request_hook_func):
if request_hook_func is None:
return None

async def async_request_hook(span, req):
return request_hook_func(span, req)

return async_request_hook

def _uninstrument(self, **kwargs):
httpx.Client = httpx._api.Client = self._original_client
httpx.AsyncClient = self._original_async_client
Expand Down

0 comments on commit 3ff1113

Please sign in to comment.