opensearch-project / opensearch-py

Python Client for OpenSearch
https://opensearch.org/docs/latest/clients/python/
Apache License 2.0
337 stars 168 forks source link

[FEATURE] Support for aiohttp Trace Config in AsyncHttpConnection for Enhanced Tracing with X-Ray #681

Open angelcabo opened 6 months ago

angelcabo commented 6 months ago

Is your feature request related to a problem?

I'm encountering difficulties integrating X-Ray tracing with aiohttp requests made through the AsyncHttpConnection class. The core of the issue is that AsyncHttpConnection does not currently seem to allow passing trace_config information to the aiohttp.ClientSession, which is necessary for implementing detailed tracing with AWS X-Ray in asynchronous environments.

What solution would you like?

Would it be possible for AsyncHttpConnection to either support passing trace_configs as part of its initialization parameters or take an option to initialize the trace hooks? Apologies if this is already supported somehow and I missed it. If that's the case, let me know a way I can hook this up without sub-classing AsyncHttpConnection myself.

What alternatives have you considered?

As a workaround, I subclassed AsyncHttpConnection to manually create an aiohttp.ClientSession with the required trace_configs parameter. While this approach works, it requires additional maintenance and familiarity with the internals of both OpenSearch and aiohttp.

from aws_xray_sdk.ext.aiohttp.client import (
        aws_xray_trace_config,  # pragma: no cover
    )

class CustomAsyncConnection(AsyncHttpConnection):
    async def _create_aiohttp_session(self) -> Any:
        if self.loop is None:
            self.loop = get_running_loop()
        self.session = aiohttp.ClientSession(
            headers=self.headers,
            skip_auto_headers=("accept", "accept-encoding"),
            auto_decompress=True,
            loop=self.loop,
            cookie_jar=aiohttp.DummyCookieJar(),
            response_class=OpenSearchClientResponse,
            connector=aiohttp.TCPConnector(limit=self._limit, use_dns_cache=True, ssl=self._ssl_context),
            trace_configs=[aws_xray_trace_config()],
        )
    test_client = AsyncOpenSearch(
        hosts=[settings.SEARCH_CLUSTER_URL],
        connection_class=CustomAsyncConnection,
    )
    info = await test_client.info()
    print(info)
    await test_client.close()
dblock commented 6 months ago

Seems reasonable. There's also a related issue on wanting to introduce metrics, #678. Feel free to contribute!