Support Trio with httpx

docs/reference/async.md

@@ -50,6 +50,27 @@ All APIs that are available under the sync client are also available under the a
 
 See also the [Using OpenTelemetry](/reference/opentelemetry.md) page.
 
+## Trio support
+
+If you prefer using Trio instead of asyncio to take advantage of its better structured concurrency support, you can use the HTTPX async node which supports Trio out of the box.
+
+```python
+import trio
+from elasticsearch import AsyncElasticsearch
+
+client = AsyncElasticsearch(
+    "https://...",
+    api_key="...",
+    node_class="httpxasync")
+
+async def main():
+    resp = await client.info()
+    print(resp.body)
+
+trio.run(main)
+```
+
+The one limitation of Trio support is that it does not currently support node sniffing, which was not implemented with structured concurrency in mind.
 
 ## Frequently Asked Questions [_frequently_asked_questions]
 

docs/reference/dsl_how_to_guides.md

@@ -1555,6 +1555,12 @@ The DSL module supports async/await with [asyncio](https://docs.python.org/3/lib
 $ python -m pip install "elasticsearch[async]"
 ```
 
+The DSL module also supports [Trio](https://trio.readthedocs.io/en/stable/) when using the Async HTTPX client. You do need to install Trio and HTTPX separately:
+
+```bash
+$ python -m pip install "elasticsearch trio httpx"
+```
+
 ### Connections [_connections]
 
 Use the `async_connections` module to manage your asynchronous connections.
@@ -1565,6 +1571,14 @@ from elasticsearch.dsl import async_connections
 async_connections.create_connection(hosts=['localhost'], timeout=20)
 ```
 
+If you're using Trio, you need to explicitly request the Async HTTP client:
+
+```python
+from elasticsearch.dsl import async_connections
+
+async_connections.create_connection(hosts=['localhost'], node_class="httpxasync")
+```
+
 All the options available in the `connections` module can be used with `async_connections`.
 
 #### How to avoid *Unclosed client session / connector* warnings on exit [_how_to_avoid_unclosed_client_session_connector_warnings_on_exit]
@@ -1576,8 +1590,6 @@ es = async_connections.get_connection()
 await es.close()
 ```
 
-
-
 ### Search DSL [_search_dsl]
 
 Use the `AsyncSearch` class to perform asynchronous searches.

elasticsearch/_async/helpers.py

@@ -33,6 +33,8 @@
     Union,
 )
 
+import sniffio
+
 from ..exceptions import ApiError, NotFoundError, TransportError
 from ..helpers.actions import (
     _TYPE_BULK_ACTION,
@@ -53,6 +55,15 @@
 T = TypeVar("T")
 
 
+async def _sleep(seconds: float) -> None:
+    if sniffio.current_async_library() == "trio":
+        import trio
+
+        await trio.sleep(seconds)
+    else:
+        await asyncio.sleep(seconds)
+
+
 async def _chunk_actions(
     actions: AsyncIterable[_TYPE_BULK_ACTION_HEADER_AND_BODY],
     chunk_size: int,
@@ -245,9 +256,7 @@ async def map_actions() -> AsyncIterable[_TYPE_BULK_ACTION_HEADER_AND_BODY]:
                 ]
             ] = []
             if attempt:
-                await asyncio.sleep(
-                    min(max_backoff, initial_backoff * 2 ** (attempt - 1))
-                )
+                await _sleep(min(max_backoff, initial_backoff * 2 ** (attempt - 1)))
 
             try:
                 data: Union[

pyproject.toml

@@ -41,11 +41,17 @@ keywords = [
 ]
 dynamic = ["version"]
 dependencies = [
-    "elastic-transport>=9.1.0,<10",
+    # TODO revert before merging/releasing
+    "elastic-transport @ git+https://github.com/pquentin/elastic-transport-python.git@trio-support",
     "python-dateutil",
     "typing-extensions",
+    "sniffio",
 ]
 
+# TODO revert before merging/releasing
+[tool.hatch.metadata]
+allow-direct-references = true
+
 [project.optional-dependencies]
 async = ["aiohttp>=3,<4"]
 requests = ["requests>=2.4.0, !=2.32.2, <3.0.0"]
@@ -56,6 +62,7 @@ vectorstore_mmr = ["numpy>=1", "simsimd>=3"]
 dev = [
     "requests>=2, <3",
     "aiohttp",
+    "httpx",
     "pytest",
     "pytest-cov",
     "pytest-mock",
@@ -78,6 +85,8 @@ dev = [
     "mapbox-vector-tile",
     "jinja2",
     "tqdm",
+    "trio",
+    "anyio",
     "mypy",
     "pyright",
     "types-python-dateutil",

test_elasticsearch/test_async/test_server/conftest.py

@@ -16,27 +16,27 @@
 #  under the License.
 
 import pytest
-import pytest_asyncio
+import sniffio
 
 import elasticsearch
 
 from ...utils import CA_CERTS, wipe_cluster
 
-pytestmark = pytest.mark.asyncio
 
-
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 async def async_client_factory(elasticsearch_url):
-
-    if not hasattr(elasticsearch, "AsyncElasticsearch"):
-        pytest.skip("test requires 'AsyncElasticsearch' and aiohttp to be installed")
-
+    print("async!", elasticsearch_url)
+    kwargs = {}
+    if sniffio.current_async_library() == "trio":
+        kwargs["node_class"] = "httpxasync"
     # Unfortunately the asyncio client needs to be rebuilt every
     # test execution due to how pytest-asyncio manages
     # event loops (one per test!)
     client = None
     try:
-        client = elasticsearch.AsyncElasticsearch(elasticsearch_url, ca_certs=CA_CERTS)
+        client = elasticsearch.AsyncElasticsearch(
+            elasticsearch_url, ca_certs=CA_CERTS, **kwargs
+        )
         yield client
     finally:
         if client:

test_elasticsearch/test_async/test_server/test_clients.py

@@ -18,7 +18,7 @@
 
 import pytest
 
-pytestmark = pytest.mark.asyncio
+pytestmark = pytest.mark.anyio
 
 
 @pytest.mark.parametrize("kwargs", [{"body": {"text": "привет"}}, {"text": "привет"}])

test_elasticsearch/test_async/test_server/test_helpers.py

@@ -15,20 +15,19 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-import asyncio
 import logging
 from datetime import datetime, timedelta, timezone
 from unittest.mock import MagicMock, call, patch
 
+import anyio
 import pytest
-import pytest_asyncio
 from elastic_transport import ApiResponseMeta, ObjectApiResponse
 
 from elasticsearch import helpers
 from elasticsearch.exceptions import ApiError
 from elasticsearch.helpers import ScanError
 
-pytestmark = [pytest.mark.asyncio]
+pytestmark = pytest.mark.anyio
 
 
 class AsyncMock(MagicMock):
@@ -92,7 +91,7 @@ async def test_all_documents_get_inserted(self, async_client):
     async def test_documents_data_types(self, async_client):
         async def async_gen():
             for x in range(100):
-                await asyncio.sleep(0)
+                await anyio.sleep(0)
                 yield {"answer": x, "_id": x}
 
         def sync_gen():
@@ -491,7 +490,7 @@ def __await__(self):
         return self().__await__()
 
 
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 async def scan_teardown(async_client):
     yield
     await async_client.clear_scroll(scroll_id="_all")
@@ -915,7 +914,7 @@ async def test_scan_from_keyword_is_aliased(async_client, scan_kwargs):
         assert "from" not in search_mock.call_args[1]
 
 
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 async def reindex_setup(async_client):
     bulk = []
     for x in range(100):
@@ -993,7 +992,7 @@ async def test_all_documents_get_moved(self, async_client, reindex_setup):
         )["_source"]
 
 
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 async def parent_reindex_setup(async_client):
     body = {
         "settings": {"number_of_shards": 1, "number_of_replicas": 0},
@@ -1054,7 +1053,7 @@ async def test_children_are_reindexed_correctly(
         } == q
 
 
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 async def reindex_data_stream_setup(async_client):
     dt = datetime.now(tz=timezone.utc)
     bulk = []

test_elasticsearch/test_async/test_server/test_mapbox_vector_tile.py

@@ -16,14 +16,13 @@
 #  under the License.
 
 import pytest
-import pytest_asyncio
 
 from elasticsearch import RequestError
 
-pytestmark = pytest.mark.asyncio
+pytestmark = pytest.mark.anyio
 
 
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 async def mvt_setup(async_client):
     await async_client.indices.create(
         index="museums",

test_elasticsearch/test_async/test_server/test_rest_api_spec.py

@@ -25,7 +25,6 @@
 import warnings
 
 import pytest
-import pytest_asyncio
 
 from elasticsearch import ElasticsearchWarning, RequestError
 
@@ -39,6 +38,8 @@
 )
 from ...utils import parse_version
 
+# We're not using `pytest.mark.anyio` here because it would run the test suite twice,
+# which does not work as it does not fully clean up after itself.
 pytestmark = pytest.mark.asyncio
 
 XPACK_FEATURES = None
@@ -240,7 +241,7 @@ async def _feature_enabled(self, name):
         return name in XPACK_FEATURES
 
 
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 def async_runner(async_client_factory):
     return AsyncYamlRunner(async_client_factory)