Merge remote-tracking branch 'origin/main' into trio-support

Author: pquentin

.buildkite/Dockerfile

@@ -1,4 +1,4 @@
-ARG PYTHON_VERSION=3.13
+ARG PYTHON_VERSION=3.14
 FROM python:${PYTHON_VERSION}
 
 # Default UID/GID to 1000

.buildkite/pipeline.yml

@@ -16,18 +16,19 @@ steps:
           - "3.11"
           - "3.12"
           - "3.13"
+          - "3.14"
         connection:
           - "urllib3"
           - "requests"
         nox_session:
           - "test"
       adjustments:
         - with:
-            python: "3.9"
+            python: "3.10"
             connection: "urllib3"
             nox_session: "test_otel"
         - with:
-            python: "3.13"
+            python: "3.14"
             connection: "urllib3"
             nox_session: "test_otel"
     command: ./.buildkite/run-tests

.buildkite/run-tests

@@ -7,7 +7,7 @@
 # Default environment variables
 export STACK_VERSION="${STACK_VERSION:=8.0.0-SNAPSHOT}"
 export TEST_SUITE="${TEST_SUITE:=platinum}"
-export PYTHON_VERSION="${PYTHON_VERSION:=3.13}"
+export PYTHON_VERSION="${PYTHON_VERSION:=3.14}"
 export PYTHON_CONNECTION_CLASS="${PYTHON_CONNECTION_CLASS:=urllib3}"
 
 script_path=$(dirname $(realpath $0))

.github/workflows/ci.yml

@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout Repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
       - name: Set up Python 3.x
         uses: actions/setup-python@v5
         with:
@@ -23,7 +23,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout Repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
       - name: Set up Python 3.x
         uses: actions/setup-python@v5
         with:
@@ -38,7 +38,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"]
         nox-session: [""]
         runs-on: ["ubuntu-latest"]
 
@@ -47,7 +47,7 @@ jobs:
     continue-on-error: false
     steps:
       - name: Checkout Repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
       - name: Set Up Python - ${{ matrix.python-version }}
         uses: actions/setup-python@v5
         with:

elasticsearch/_async/helpers.py

@@ -35,12 +35,16 @@
 
 import sniffio
 
+from ..compat import safe_task
 from ..exceptions import ApiError, NotFoundError, TransportError
 from ..helpers.actions import (
     _TYPE_BULK_ACTION,
     _TYPE_BULK_ACTION_BODY,
     _TYPE_BULK_ACTION_HEADER,
     _TYPE_BULK_ACTION_HEADER_AND_BODY,
+    _TYPE_BULK_ACTION_HEADER_WITH_META_AND_BODY,
+    _TYPE_BULK_ACTION_WITH_META,
+    BulkMeta,
     _ActionChunker,
     _process_bulk_chunk_error,
     _process_bulk_chunk_success,
@@ -65,9 +69,10 @@ async def _sleep(seconds: float) -> None:
 
 
 async def _chunk_actions(
-    actions: AsyncIterable[_TYPE_BULK_ACTION_HEADER_AND_BODY],
+    actions: AsyncIterable[_TYPE_BULK_ACTION_HEADER_WITH_META_AND_BODY],
     chunk_size: int,
     max_chunk_bytes: int,
+    flush_after_seconds: Optional[float],
     serializer: Serializer,
 ) -> AsyncIterable[
     Tuple[
@@ -87,10 +92,42 @@ async def _chunk_actions(
     chunker = _ActionChunker(
         chunk_size=chunk_size, max_chunk_bytes=max_chunk_bytes, serializer=serializer
     )
-    async for action, data in actions:
-        ret = chunker.feed(action, data)
-        if ret:
-            yield ret
+
+    if not flush_after_seconds:
+        async for action, data in actions:
+            ret = chunker.feed(action, data)
+            if ret:
+                yield ret
+    else:
+        item_queue: asyncio.Queue[_TYPE_BULK_ACTION_HEADER_WITH_META_AND_BODY] = (
+            asyncio.Queue()
+        )
+
+        async def get_items() -> None:
+            try:
+                async for item in actions:
+                    await item_queue.put(item)
+            finally:
+                await item_queue.put((BulkMeta.done, None))
+
+        async with safe_task(get_items()):
+            timeout: Optional[float] = flush_after_seconds
+            while True:
+                try:
+                    action, data = await asyncio.wait_for(
+                        item_queue.get(), timeout=timeout
+                    )
+                    timeout = flush_after_seconds
+                except asyncio.TimeoutError:
+                    action, data = BulkMeta.flush, None
+                    timeout = None
+
+                if action is BulkMeta.done:
+                    break
+                ret = chunker.feed(action, data)
+                if ret:
+                    yield ret
+
     ret = chunker.flush()
     if ret:
         yield ret
@@ -170,9 +207,13 @@ async def azip(
 
 async def async_streaming_bulk(
     client: AsyncElasticsearch,
-    actions: Union[Iterable[_TYPE_BULK_ACTION], AsyncIterable[_TYPE_BULK_ACTION]],
+    actions: Union[
+        Iterable[_TYPE_BULK_ACTION_WITH_META],
+        AsyncIterable[_TYPE_BULK_ACTION_WITH_META],
+    ],
     chunk_size: int = 500,
     max_chunk_bytes: int = 100 * 1024 * 1024,
+    flush_after_seconds: Optional[float] = None,
     raise_on_error: bool = True,
     expand_action_callback: Callable[
         [_TYPE_BULK_ACTION], _TYPE_BULK_ACTION_HEADER_AND_BODY
@@ -205,6 +246,9 @@ async def async_streaming_bulk(
     :arg actions: iterable or async iterable containing the actions to be executed
     :arg chunk_size: number of docs in one chunk sent to es (default: 500)
     :arg max_chunk_bytes: the maximum size of the request in bytes (default: 100MB)
+    :arg flush_after_seconds: time in seconds after which a chunk is written even
+        if hasn't reached `chunk_size` or `max_chunk_bytes`. Set to 0 to not use a
+        timeout-based flush. (default: 0)
     :arg raise_on_error: raise ``BulkIndexError`` containing errors (as `.errors`)
         from the execution of the last chunk when some occur. By default we raise.
     :arg raise_on_exception: if ``False`` then don't propagate exceptions from
@@ -231,9 +275,14 @@ async def async_streaming_bulk(
     if isinstance(retry_on_status, int):
         retry_on_status = (retry_on_status,)
 
-    async def map_actions() -> AsyncIterable[_TYPE_BULK_ACTION_HEADER_AND_BODY]:
+    async def map_actions() -> (
+        AsyncIterable[_TYPE_BULK_ACTION_HEADER_WITH_META_AND_BODY]
+    ):
         async for item in aiter(actions):
-            yield expand_action_callback(item)
+            if isinstance(item, BulkMeta):
+                yield item, None
+            else:
+                yield expand_action_callback(item)
 
     serializer = client.transport.serializers.get_serializer("application/json")
 
@@ -245,7 +294,7 @@ async def map_actions() -> AsyncIterable[_TYPE_BULK_ACTION_HEADER_AND_BODY]:
     ]
     bulk_actions: List[bytes]
     async for bulk_data, bulk_actions in _chunk_actions(
-        map_actions(), chunk_size, max_chunk_bytes, serializer
+        map_actions(), chunk_size, max_chunk_bytes, flush_after_seconds, serializer
     ):
         for attempt in range(max_retries + 1):
             to_retry: List[bytes] = []

elasticsearch/compat.py

@@ -15,11 +15,14 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
+import asyncio
 import inspect
 import os
 import sys
+from contextlib import asynccontextmanager, contextmanager
 from pathlib import Path
-from typing import Tuple, Type, Union
+from threading import Thread
+from typing import Any, AsyncIterator, Callable, Coroutine, Iterator, Tuple, Type, Union
 
 string_types: Tuple[Type[str], Type[bytes]] = (str, bytes)
 
@@ -76,9 +79,48 @@ def warn_stacklevel() -> int:
     return 0
 
 
+@contextmanager
+def safe_thread(
+    target: Callable[..., Any], *args: Any, **kwargs: Any
+) -> Iterator[Thread]:
+    """Run a thread within a context manager block.
+
+    The thread is automatically joined when the block ends. If the thread raised
+    an exception, it is raised in the caller's context.
+    """
+    captured_exception = None
+
+    def run() -> None:
+        try:
+            target(*args, **kwargs)
+        except BaseException as exc:
+            nonlocal captured_exception
+            captured_exception = exc
+
+    thread = Thread(target=run)
+    thread.start()
+    yield thread
+    thread.join()
+    if captured_exception:
+        raise captured_exception
+
+
+@asynccontextmanager
+async def safe_task(coro: Coroutine[Any, Any, Any]) -> AsyncIterator[asyncio.Task[Any]]:
+    """Run a background task within a context manager block.
+
+    The task is awaited when the block ends.
+    """
+    task = asyncio.create_task(coro)
+    yield task
+    await task
+
+
 __all__ = [
     "string_types",
     "to_str",
     "to_bytes",
     "warn_stacklevel",
+    "safe_thread",
+    "safe_task",
 ]

elasticsearch/dsl/document_base.py

@@ -34,6 +34,11 @@
     overload,
 )
 
+try:
+    import annotationlib
+except ImportError:
+    annotationlib = None
+
 try:
     from types import UnionType
 except ImportError:
@@ -332,6 +337,16 @@ def __init__(self, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]):
         #     # ignore attributes
         #     field10: ClassVar[string] = "a regular class variable"
         annotations = attrs.get("__annotations__", {})
+        if not annotations and annotationlib:
+            # Python 3.14+ uses annotationlib
+            annotate = annotationlib.get_annotate_from_class_namespace(attrs)
+            if annotate:
+                annotations = (
+                    annotationlib.call_annotate_function(
+                        annotate, format=annotationlib.Format.VALUE
+                    )
+                    or {}
+                )
         fields = {n for n in attrs if isinstance(attrs[n], Field)}
         fields.update(annotations.keys())
         field_defaults = {}