diff --git a/src/harness_sdk/instrumentation/httpx/__init__.py b/src/harness_sdk/instrumentation/httpx/__init__.py index e40ef04..c046de3 100644 --- a/src/harness_sdk/instrumentation/httpx/__init__.py +++ b/src/harness_sdk/instrumentation/httpx/__init__.py @@ -6,11 +6,8 @@ from harness_sdk.plugins.control import get_control_registry from harness_sdk.instrumentation import BaseInstrumentorWrapper from harness_sdk.instrumentation.httpx.utils import ( - decode_response_body_for_capture, headers_from_httpx, read_request_body, - read_response_body, - read_response_body_async, url_from_request_info, ) @@ -26,6 +23,7 @@ def __init__(self): logger.debug('Entering HTTPXClientInstrumentorWrapper.__init__().') HTTPXClientInstrumentor.__init__(self) BaseInstrumentorWrapper.__init__(self) + self._process_response_body = False def _instrument(self, **kwargs) -> None: '''Enable instrumentation with request/response hooks.''' @@ -46,22 +44,14 @@ def _process_request(self, span, request_info): def _process_response(self, span, response_info): headers = headers_from_httpx(response_info.headers) - body = read_response_body(response_info.stream) - body = decode_response_body_for_capture(headers, body) - self.generic_response_handler(headers, body, span) - - async def _process_response_async(self, span, response_info): - headers = headers_from_httpx(response_info.headers) - body = await read_response_body_async(response_info.stream) - body = decode_response_body_for_capture(headers, body) - self.generic_response_handler(headers, body, span) + self.generic_response_handler(headers, None, span) def request_hook(self, span, request_info): '''Capture sync client request data and run evaluation.''' self._process_request(span, request_info) def response_hook(self, span, request_info, response_info): # pylint: disable=unused-argument - '''Capture sync client response data.''' + '''Capture sync client response headers.''' self._process_response(span, response_info) async def async_request_hook(self, span, request_info): @@ -69,5 +59,5 @@ async def async_request_hook(self, span, request_info): self._process_request(span, request_info) async def async_response_hook(self, span, request_info, response_info): # pylint: disable=unused-argument - '''Capture async client response data.''' - await self._process_response_async(span, response_info) + '''Capture async client response headers.''' + self._process_response(span, response_info) diff --git a/src/harness_sdk/instrumentation/httpx/utils.py b/src/harness_sdk/instrumentation/httpx/utils.py index 2ea4df3..eab643f 100644 --- a/src/harness_sdk/instrumentation/httpx/utils.py +++ b/src/harness_sdk/instrumentation/httpx/utils.py @@ -1,6 +1,4 @@ '''Helpers for extracting httpx request/response data for tracing.''' -import gzip -import zlib from harness_sdk.custom_logger import get_custom_logger @@ -19,65 +17,6 @@ def url_from_request_info(request_info): return str(request_info.url) -def _content_encoding_chain(headers): - '''Return Content-Encoding tokens in wire order (outer coding first).''' - if not headers: - return [] - lower = {str(k).lower(): v for k, v in headers.items()} - raw = lower.get('content-encoding') or '' - return [t.strip().lower() for t in raw.split(',') if t.strip()] - - -def _apply_content_encoding_decodings(original, chain): - '''Apply decodings in reverse wire order; return ``original`` on failure.''' - out = original - for encoding in reversed(chain): - try: - if encoding in ('gzip', 'x-gzip'): - out = gzip.decompress(out) - elif encoding == 'deflate': - try: - out = zlib.decompress(out, -zlib.MAX_WBITS) - except zlib.error: - out = zlib.decompress(out) - elif encoding == 'br': - try: - import brotli # pylint: disable=import-outside-toplevel - except ImportError: - logger.debug('brotli not installed; skipping br decompression for capture') - return original - out = brotli.decompress(out) - elif encoding in ('identity', 'compress'): - continue - else: - logger.debug('Unsupported content-encoding for capture: %s', encoding) - return original - except Exception: # pylint: disable=broad-except - logger.debug('Response body decompression failed (%s)', encoding, exc_info=True) - return original - return out - - -def decode_response_body_for_capture(headers, body): - '''Decode compressed response bytes for span capture (gzip/deflate/br). - - Transport hooks often see the on-the-wire body while ``Content-Encoding`` - still names the compression. httpx may also surface compressed bytes when - decoding is deferred. If decoding fails or an encoding is unsupported, - the original ``body`` is returned. - ''' - if body in (None, b''): - return body - if not isinstance(body, (bytes, bytearray)): - return body - - chain = _content_encoding_chain(headers) - if not chain: - return body - - return _apply_content_encoding_decodings(bytes(body), chain) - - def _body_from_byte_stream(stream): internal = getattr(stream, "_stream", None) if isinstance(internal, (bytes, bytearray)): @@ -99,62 +38,3 @@ def read_request_body(stream): except Exception: # pylint: disable=broad-except logger.debug("Unable to read httpx request stream body", exc_info=True) return None - - -def read_response_body(stream): - '''Read response body bytes and restore the stream for downstream consumers.''' - if stream is None: - return None - - body = _body_from_byte_stream(stream) - if body is not None: - return body - - httpcore_stream = getattr(stream, "_httpcore_stream", None) - if httpcore_stream is not None: - try: - chunks = list(httpcore_stream) - body = b"".join(chunks) - stream._httpcore_stream = chunks # pylint: disable=protected-access - return body - except Exception: # pylint: disable=broad-except - logger.debug("Unable to read httpx response stream body", exc_info=True) - return None - - return None - - -async def read_response_body_async(stream): - '''Read async response body bytes and restore the stream for downstream consumers.''' - if stream is None: - return None - - body = _body_from_byte_stream(stream) - if body is not None: - return body - - httpcore_stream = getattr(stream, "_httpcore_stream", None) - if httpcore_stream is not None: - try: - chunks = [] - async for part in httpcore_stream: - chunks.append(part) - body = b"".join(chunks) - stream._httpcore_stream = _ReplayAsyncStream(chunks) # pylint: disable=protected-access - return body - except Exception: # pylint: disable=broad-except - logger.debug("Unable to read httpx async response stream body", exc_info=True) - return None - - return None - - -class _ReplayAsyncStream: # pylint: disable=too-few-public-methods - '''Minimal async iterable used to replay captured response chunks.''' - - def __init__(self, chunks): - self._chunks = chunks - - async def __aiter__(self): - for chunk in self._chunks: - yield chunk diff --git a/test/instrumentation/httpx_client/httpx_integration_test.py b/test/instrumentation/httpx_client/httpx_integration_test.py index 87b07ff..a6cdfa3 100644 --- a/test/instrumentation/httpx_client/httpx_integration_test.py +++ b/test/instrumentation/httpx_client/httpx_integration_test.py @@ -9,11 +9,12 @@ from test.instrumentation.flask.app import FlaskServer -def _client_span(spans): +def _client_span(spans, url): for span in spans: - if span.kind == SpanKind.CLIENT: - return json.loads(span.to_json()) - raise AssertionError('No client span found') + span_data = json.loads(span.to_json()) + if span.kind == SpanKind.CLIENT and span_data['attributes'].get('http.url') == url: + return span_data + raise AssertionError(f'No client span found for {url}') def test_httpx_client_get(agent, exporter): @@ -38,11 +39,11 @@ def api_example(): spans = exporter.get_finished_spans() assert spans - client_span = _client_span(spans) + client_span = _client_span(spans, url) assert client_span['attributes']['http.method'] == 'GET' assert client_span['attributes']['http.url'] == url - assert client_span['attributes']['http.response.body'] == '{ "a": "a", "xyz": "xyz" }' + assert 'http.response.body' not in client_span['attributes'] assert client_span['attributes']['http.status_code'] == 200 assert client_span['attributes']['http.response.header.tester3'] == 'tester3' finally: @@ -71,15 +72,16 @@ def api_example(): spans = exporter.get_finished_spans() assert spans - client_span = _client_span(spans) + client_span = _client_span(spans, url) assert client_span['kind'] == "SpanKind.CLIENT" assert client_span['attributes']['http.method'] == 'POST' assert client_span['attributes']['http.url'] == url assert client_span['attributes']['http.request.header.content-type'] == 'application/json' assert client_span['attributes']['http.request.body'] == '{"test":"body"}' - assert client_span['attributes']['http.response.body'] == '{ "a": "a", "xyz": "xyz" }' + assert 'http.response.body' not in client_span['attributes'] assert client_span['attributes']['http.status_code'] == 200 + assert client_span['attributes']['http.response.header.tester3'] == 'tester3' finally: server.shutdown() @@ -114,7 +116,7 @@ def api_example(): span_list = exporter.get_finished_spans() assert span_list - httpx_span = _client_span(span_list) + httpx_span = _client_span(span_list, url) assert httpx_span['attributes']['http.method'] == 'POST' assert httpx_span['attributes']['http.url'] == url @@ -122,6 +124,6 @@ def api_example(): assert httpx_span['attributes']['http.request.header.tester2'] == 'tester2' assert httpx_span['attributes']['http.request.body'] == '{ "a":"b", "c": "d" }' assert httpx_span['attributes']['http.response.header.content-type'] == 'application/json' - assert httpx_span['attributes']['http.response.body'] == '{ "a": "a", "xyz": "xyz" }' + assert 'http.response.body' not in httpx_span['attributes'] assert httpx_span['attributes']['http.status_code'] == 200 server.shutdown() diff --git a/test/instrumentation/httpx_client/test_httpx_hooks.py b/test/instrumentation/httpx_client/test_httpx_hooks.py new file mode 100644 index 0000000..0c6205e --- /dev/null +++ b/test/instrumentation/httpx_client/test_httpx_hooks.py @@ -0,0 +1,71 @@ +from types import SimpleNamespace + +import pytest + +from harness_sdk.instrumentation.httpx import HTTPXClientInstrumentorWrapper + + +class _SyncStream: + def __init__(self): + self.iterations = 0 + + def __iter__(self): + self.iterations += 1 + yield b"response" + + +class _AsyncStream: + def __init__(self): + self.iterations = 0 + + async def __aiter__(self): + self.iterations += 1 + yield b"response" + + +def _wrapper_with_response_capture(): + wrapper = HTTPXClientInstrumentorWrapper() + calls = [] + wrapper.generic_response_handler = ( + lambda headers, body, span: calls.append((headers, body, span)) + ) + return wrapper, calls + + +def test_response_hook_does_not_consume_or_replace_stream(): + wrapper, calls = _wrapper_with_response_capture() + span = object() + transport_stream = _SyncStream() + response_stream = SimpleNamespace(_httpcore_stream=transport_stream) + response_info = SimpleNamespace( + status_code=200, + headers={"content-type": "text/plain"}, + stream=response_stream, + extensions={}, + ) + + wrapper.response_hook(span, None, response_info) + + assert calls == [({"content-type": "text/plain"}, None, span)] + assert transport_stream.iterations == 0 + assert response_stream._httpcore_stream is transport_stream + + +@pytest.mark.asyncio +async def test_async_response_hook_does_not_consume_or_replace_stream(): + wrapper, calls = _wrapper_with_response_capture() + span = object() + transport_stream = _AsyncStream() + response_stream = SimpleNamespace(_httpcore_stream=transport_stream) + response_info = SimpleNamespace( + status_code=200, + headers={"content-type": "text/plain"}, + stream=response_stream, + extensions={}, + ) + + await wrapper.async_response_hook(span, None, response_info) + + assert calls == [({"content-type": "text/plain"}, None, span)] + assert transport_stream.iterations == 0 + assert response_stream._httpcore_stream is transport_stream diff --git a/test/instrumentation/httpx_client/test_httpx_utils.py b/test/instrumentation/httpx_client/test_httpx_utils.py deleted file mode 100644 index 3996fed..0000000 --- a/test/instrumentation/httpx_client/test_httpx_utils.py +++ /dev/null @@ -1,36 +0,0 @@ -import gzip -import zlib - -import pytest - -from harness_sdk.instrumentation.httpx.utils import decode_response_body_for_capture - - -def test_decode_gzip_response_body(): - payload = b'{"hello":"world"}' - compressed = gzip.compress(payload) - headers = {'Content-Encoding': 'gzip', 'content-type': 'application/json'} - assert decode_response_body_for_capture(headers, compressed) == payload - - -def test_decode_deflate_response_body(): - payload = b'plain text body' - compressed = zlib.compress(payload) - headers = {'Content-Encoding': 'deflate'} - assert decode_response_body_for_capture(headers, compressed) == payload - - -def test_no_content_encoding_returns_raw(): - payload = b'not compressed' - assert decode_response_body_for_capture({'content-type': 'text/plain'}, payload) == payload - - -def test_bad_gzip_returns_original(): - payload = b'not gzip at all' - headers = {'Content-Encoding': 'gzip'} - assert decode_response_body_for_capture(headers, payload) == payload - - -@pytest.mark.parametrize('empty', [None, b'']) -def test_decode_empty_body(empty): - assert decode_response_body_for_capture({'Content-Encoding': 'gzip'}, empty) == empty