Fix metadata timeout propagation

- propagate per-call read_timeout, connection_timeout, and timeout (operation deadline) options across query setup metadata calls (container read, query plan, /pkranges) in sync and async paths
- extend test coverage for timeout propagation

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: dibahlfi

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed per-call `read_timeout`, `connection_timeout`, and `timeout` (operation deadline) being dropped on the metadata calls a query makes before its first page.
 
 #### Other Changes
 

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -124,6 +124,11 @@ def build_options(kwargs: dict[str, Any]) -> dict[str, Any]:
         options[Constants.Kwargs.READ_TIMEOUT] = kwargs[Constants.Kwargs.READ_TIMEOUT]
     if Constants.Kwargs.TIMEOUT in kwargs:
         options[Constants.Kwargs.TIMEOUT] = kwargs[Constants.Kwargs.TIMEOUT]
+    # Copy (not pop) so connection_timeout stays in kwargs for the page fetch
+    # and is also placed in options, where the container read, partition-key
+    # ranges, and query plan calls read it.
+    if Constants.Kwargs.CONNECTION_TIMEOUT in kwargs:
+        options[Constants.Kwargs.CONNECTION_TIMEOUT] = kwargs[Constants.Kwargs.CONNECTION_TIMEOUT]
 
 
     options[Constants.OperationStartTime] = time.time()
@@ -1082,6 +1087,70 @@ def _build_properties_cache(properties: dict[str, Any], container_link: str) ->
         "partitionKey": properties.get("partitionKey", None), "container_link": container_link
     }
 
+# The per-call timeout keys a caller can set on a single request. Listed once
+# here so format_pk_range_options and the hybrid-search fetch forward the same
+# set.
+_PER_CALL_TIMEOUT_OPTION_KEYS: Tuple[str, ...] = (
+    Constants.Kwargs.READ_TIMEOUT,
+    Constants.Kwargs.CONNECTION_TIMEOUT,
+    Constants.Kwargs.TIMEOUT,
+)
+
+# The operation deadline is checked as elapsed = now - OperationStartTime, and
+# OperationStartTime defaults to the current time when it is missing. So timeout
+# and OperationStartTime must be carried together onto the metadata setup calls;
+# otherwise a setup call measures the deadline from its own start instead of the
+# operation's start. This adds OperationStartTime to the three timeout keys above.
+_PER_CALL_DEADLINE_OPTION_KEYS: Tuple[str, ...] = _PER_CALL_TIMEOUT_OPTION_KEYS + (
+    Constants.OperationStartTime,
+)
+
+
+def _carry_per_call_timeout_options(source: Mapping[str, Any], destination: dict[str, Any]) -> None:
+    """Copy the per-call timeouts and the operation start time from source into destination.
+
+    Copies read_timeout, connection_timeout, timeout, and OperationStartTime. Only
+    keys present in source are copied, so a timeout the caller did not set stays
+    absent and the request uses the client default instead of None.
+
+    :param source: The request options to read the timeouts from.
+    :type source: ~collections.abc.Mapping[str, typing.Any]
+    :param destination: The options dict to copy the timeouts into.
+    :type destination: dict[str, typing.Any]
+    :return: None
+    :rtype: None
+    """
+    for key in _PER_CALL_DEADLINE_OPTION_KEYS:
+        if key in source:
+            destination[key] = source[key]
+
+
+def _copy_per_call_timeouts_to_kwargs(
+        options: Optional[Mapping[str, Any]],
+        kwargs: dict[str, Any]
+) -> None:
+    """Copy the per-call timeouts and the operation start time from options into kwargs.
+
+    Moves read_timeout, connection_timeout, timeout, and OperationStartTime from
+    the request options into the kwargs the request layer reads. A value is copied
+    only when it is set (not None), so an unset timeout falls back to the client
+    default instead of None; setdefault keeps any value already in kwargs.
+
+    :param options: The request options to read the timeouts from (may be None or empty).
+    :type options: ~collections.abc.Mapping[str, typing.Any] or None
+    :param kwargs: The kwargs dict to copy the timeouts into; mutated in place.
+    :type kwargs: dict[str, typing.Any]
+    :return: None
+    :rtype: None
+    """
+    if not options:
+        return
+    for key in _PER_CALL_DEADLINE_OPTION_KEYS:
+        value = options.get(key)
+        if value is not None:
+            kwargs.setdefault(key, value)
+
+
 def format_pk_range_options(query_options: Mapping[str, Any]) -> dict[str, Any]:
     """Formats the partition key range options to be used internally from the query ones.
     :param dict query_options: The query options being used.
@@ -1094,4 +1163,7 @@ def format_pk_range_options(query_options: Mapping[str, Any]) -> dict[str, Any]:
             pk_range_options[Constants.ContainerRID] = query_options[Constants.ContainerRID]
         if "excludedLocations" in query_options:
             pk_range_options["excludedLocations"] = query_options["excludedLocations"]
+        # Keep the per-call timeouts so the partition-key ranges fetch uses them
+        # instead of the client default.
+        _carry_per_call_timeout_options(query_options, pk_range_options)
     return pk_range_options

sdk/cosmos/azure-cosmos/azure/cosmos/_constants.py

@@ -116,15 +116,17 @@ class _Constants:
     class Kwargs:
         """Keyword arguments used in the azure-cosmos package"""
 
+        # Whether to retry write operations if they fail. Used either at client level or request level.
         RETRY_WRITE: Literal["retry_write"] = "retry_write"
-        """Whether to retry write operations if they fail. Used either at client level or request level."""
         EXCLUDED_LOCATIONS: Literal["excludedLocations"] = "excludedLocations"
+        # Availability strategy config. Used either at client level or request level.
         AVAILABILITY_STRATEGY: Literal["availabilityStrategy"] = "availabilityStrategy"
-        """Availability strategy config. Used either at client level or request level"""
+        # Socket read timeout in seconds. Used either at client level or request level.
         READ_TIMEOUT: Literal["read_timeout"] = "read_timeout"
-        """Socket read timeout in seconds. Used either at client level or request level."""
+        # Absolute timeout in seconds for the combined HTTP request and response processing.
         TIMEOUT: Literal["timeout"] = "timeout"
-        """Absolute timeout in seconds for the combined HTTP request and response processing."""
+        # Socket connect (handshake) timeout in seconds. Used either at client level or request level.
+        CONNECTION_TIMEOUT: Literal["connection_timeout"] = "connection_timeout"
 
     class UserAgentFeatureFlags(IntEnum):
         """

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -3244,18 +3244,11 @@ def __QueryFeed(  # pylint: disable=too-many-locals, too-many-statements, too-ma
         """
         if options is None:
             options = {}
-        read_timeout = options.get("read_timeout")
-        if read_timeout is not None:
-            # we currently have a gap where kwargs are not getting passed correctly down the pipeline. In order to make
-            # absolute time out work, we are passing read_timeout via kwargs as a temporary fix
-            kwargs.setdefault("read_timeout", read_timeout)
-
-        operation_start_time = options.get(Constants.OperationStartTime)
-        if operation_start_time is not None:
-            kwargs.setdefault(Constants.OperationStartTime, operation_start_time)
-        timeout = options.get("timeout")
-        if timeout is not None:
-            kwargs.setdefault("timeout", timeout)
+        # Copy the per-call timeouts and the operation start time out of options into
+        # kwargs, where _Request reads them. A value is copied only when set, so
+        # an unset timeout falls back to the client/policy default instead of
+        # None; setdefault keeps any explicit kwarg the caller already placed.
+        base._copy_per_call_timeouts_to_kwargs(options, kwargs)
 
         # Execution context injects this via request options; keep kwargs fallback
         # for compatibility with call paths that still thread internal values there.

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/execution_dispatcher.py

@@ -35,6 +35,7 @@
 from azure.cosmos.exceptions import CosmosHttpResponseError
 from azure.cosmos.http_constants import StatusCodes
 from ..._constants import _Constants as Constants
+from ... import _base
 
 # pylint: disable=protected-access
 
@@ -67,11 +68,17 @@ def __init__(self, client, resource_link, query, options, fetch_function,
     async def _create_execution_context_with_query_plan(self):
         self._fetched_query_plan = True
         query_to_use = self._query if self._query is not None else "Select * from root r"
+        # read_timeout is forwarded as-is (None when the caller did not set it) to
+        # keep its existing behavior. It is set before the helper, so the helper's
+        # setdefault leaves it unchanged and only adds connection_timeout, timeout,
+        # and OperationStartTime when the caller set them.
+        query_plan_kwargs = {"read_timeout": self._options.get('read_timeout')}
+        _base._copy_per_call_timeouts_to_kwargs(self._options, query_plan_kwargs)
         query_plan = await self._client._GetQueryPlanThroughGateway(
             query_to_use,
             self._resource_link,
             self._options.get('excludedLocations'),
-            read_timeout=self._options.get('read_timeout')
+            **query_plan_kwargs
         )
         query_execution_info = _PartitionedQueryExecutionInfo(query_plan)
         qe_info = getattr(query_execution_info, "_query_execution_info", None)

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/hybrid_search_aggregator.py

@@ -10,6 +10,7 @@
     _FULL_TEXT_SCORE_SCOPE_KEY, _FULL_TEXT_SCORE_SCOPE_LOCAL, _FULL_TEXT_SCORE_SCOPE_DEFAULT
 from azure.cosmos._routing import routing_range
 from azure.cosmos import exceptions
+from azure.cosmos import _base
 from ..._constants import _Constants as Constants
 
 # pylint: disable=protected-access
@@ -297,6 +298,9 @@ async def _get_target_partition_key_range(self, target_all_ranges):
             feed_options = {}
             if Constants.ContainerRID in self._options:
                 feed_options[Constants.ContainerRID] = self._options[Constants.ContainerRID]
+            # This path calls _ReadPartitionKeyRanges directly and skips
+            # format_pk_range_options, so copy the per-call timeouts here too.
+            _base._carry_per_call_timeout_options(self._options, feed_options)
             return [item async for item in self._client._ReadPartitionKeyRanges(
                 collection_link=self._resource_link, feed_options=feed_options)]
         query_ranges = self._partitioned_query_ex_info.get_query_ranges()

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/execution_dispatcher.py

@@ -34,6 +34,7 @@
 from azure.cosmos.documents import _DistinctType
 from azure.cosmos.http_constants import StatusCodes, SubStatusCodes
 from .._constants import _Constants as Constants
+from .. import _base
 
 # pylint: disable=protected-access
 
@@ -97,11 +98,17 @@ def __init__(self, client, resource_link, query, options, fetch_function, respon
     def _create_execution_context_with_query_plan(self):
         self._fetched_query_plan = True
         query_to_use = self._query if self._query is not None else "Select * from root r"
+        # read_timeout is forwarded as-is (None when the caller did not set it) to
+        # keep its existing behavior. It is set before the helper, so the helper's
+        # setdefault leaves it unchanged and only adds connection_timeout, timeout,
+        # and OperationStartTime when the caller set them.
+        query_plan_kwargs = {"read_timeout": self._options.get('read_timeout')}
+        _base._copy_per_call_timeouts_to_kwargs(self._options, query_plan_kwargs)
         query_plan = self._client._GetQueryPlanThroughGateway(
             query_to_use,
             self._resource_link,
             self._options.get('excludedLocations'),
-            read_timeout=self._options.get('read_timeout')
+            **query_plan_kwargs
         )
         query_execution_info = _PartitionedQueryExecutionInfo(query_plan)
         qe_info = getattr(query_execution_info, "_query_execution_info", None)

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/hybrid_search_aggregator.py

@@ -8,6 +8,7 @@
 from azure.cosmos._execution_context import document_producer
 from azure.cosmos._routing import routing_range
 from azure.cosmos import exceptions
+from azure.cosmos import _base
 from .._constants import _Constants as Constants
 
 # pylint: disable=protected-access
@@ -454,6 +455,9 @@ def _get_target_partition_key_range(self, target_all_ranges):
             feed_options = {}
             if Constants.ContainerRID in self._options:
                 feed_options[Constants.ContainerRID] = self._options[Constants.ContainerRID]
+            # This path calls _ReadPartitionKeyRanges directly and skips
+            # format_pk_range_options, so copy the per-call timeouts here too.
+            _base._carry_per_call_timeout_options(self._options, feed_options)
             return list(self._client._ReadPartitionKeyRanges(
                 collection_link=self._resource_link, feed_options=feed_options))
         query_ranges = self._partitioned_query_ex_info.get_query_ranges()

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py

@@ -39,7 +39,8 @@
 from .. import _utils as utils
 from .._availability_strategy_config import _validate_request_hedging_strategy
 from .._base import (_build_properties_cache, _deserialize_throughput, _replace_throughput,
-                     build_options as _build_options, GenerateGuidId, validate_cache_staleness_value)
+                     build_options as _build_options, _copy_per_call_timeouts_to_kwargs,
+                     GenerateGuidId, validate_cache_staleness_value)
 from .._change_feed.feed_range_internal import FeedRangeInternalEpk
 
 from .._cosmos_responses import CosmosDict, CosmosList, CosmosAsyncItemPaged
@@ -102,13 +103,9 @@ async def _get_properties_with_options(self, options: Optional[dict[str, Any]] =
         if options:
             if "excludedLocations" in options:
                 kwargs['excluded_locations'] = options['excludedLocations']
-            if Constants.OperationStartTime in options:
-                kwargs[Constants.OperationStartTime] = options[Constants.OperationStartTime]
-            if Constants.Kwargs.TIMEOUT in options:
-                kwargs[Constants.Kwargs.TIMEOUT] = options[Constants.Kwargs.TIMEOUT]
-            if Constants.Kwargs.READ_TIMEOUT in options:
-                kwargs[Constants.Kwargs.READ_TIMEOUT] = options[Constants.Kwargs.READ_TIMEOUT]
-
+            # Forward the per-call timeouts and the operation start time so the
+            # container read honors them instead of the client/policy default.
+            _copy_per_call_timeouts_to_kwargs(options, kwargs)
         return await self._get_properties(**kwargs)
 
     async def _get_properties(self, **kwargs: Any) -> dict[str, Any]:

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py

@@ -3038,20 +3038,11 @@ async def __QueryFeed(  # pylint: disable=too-many-branches,too-many-statements,
         if options is None:
             options = {}
 
-        read_timeout = options.get("read_timeout")
-        if read_timeout is not None:
-            # we currently have a gap where kwargs are not getting passed correctly down the pipeline. In order to make
-            # absolute time out work, we are passing read_timeout via kwargs as a temporary fix
-            kwargs.setdefault("read_timeout", read_timeout)
-
-        operation_start_time = options.get(Constants.OperationStartTime)
-        if operation_start_time is not None:
-            # we need to set operation_state in kwargs as thats where it is looked at while sending the request
-            kwargs.setdefault(Constants.OperationStartTime, operation_start_time)
-        timeout = options.get("timeout")
-        if timeout is not None:
-            # we need to set operation_state in kwargs as that's where it is looked at while sending the request
-            kwargs.setdefault("timeout", timeout)
+        # Copy the per-call timeouts and the operation start time out of options into
+        # kwargs, where _Request reads them. A value is copied only when set, so
+        # an unset timeout falls back to the client/policy default instead of
+        # None; setdefault keeps any explicit kwarg the caller already placed.
+        base._copy_per_call_timeouts_to_kwargs(options, kwargs)
 
         # The capture dict can arrive via two upstream paths:
         #   1. The query execution context puts it into ``options`` (the