Skip to content

Add Cosmos DB dedicated gateway bypass cache and shard key support #44846

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Copilot wants to merge 9 commits into
base: main
Choose a base branch
from
Open

Add Cosmos DB dedicated gateway bypass cache and shard key support #44846

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
4c3e9d1
Initial plan
Copilot Jan 26, 2026
65f4491
Add bypass_integrated_cache and dedicated_gateway_shard_key parameters
Copilot Jan 26, 2026
87148ce
Add comprehensive tests for bypass_integrated_cache and dedicated_gat…
Copilot Jan 26, 2026
470329a
Complete async container query_items overloads updates
Copilot Jan 26, 2026
1a304ef
Address code review feedback
Copilot Jan 26, 2026
d404df8
Fix null checking and add missing async docstrings
Copilot Jan 26, 2026
7d04430
Final style and consistency fixes
Copilot Jan 26, 2026
59941d3
Merge branch 'main' into copilot/add-dedicated-gateway-features
simorenoh Feb 11, 2026
cf89cc9
Address code review feedback
Copilot Apr 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 18 additions & 1 deletion sdk/cosmos/azure-cosmos/azure/cosmos/_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -417,9 +417,15 @@ def GetHeaders( # pylint: disable=too-many-statements,too-many-branches
if options.get("populateQuotaInfo"):
headers[http_constants.HttpHeaders.PopulateQuotaInfo] = options["populateQuotaInfo"]

if options.get("maxIntegratedCacheStaleness"):
if "maxIntegratedCacheStaleness" in options:
headers[http_constants.HttpHeaders.DedicatedGatewayCacheStaleness] = options["maxIntegratedCacheStaleness"]

if "bypassIntegratedCache" in options:
headers[http_constants.HttpHeaders.DedicatedGatewayBypassCache] = "true" if options["bypassIntegratedCache"] else "false"

if options.get("dedicatedGatewayShardKey"):
headers[http_constants.HttpHeaders.DedicatedGatewayShardKey] = options["dedicatedGatewayShardKey"]

if options.get("autoUpgradePolicy"):
headers[http_constants.HttpHeaders.AutoscaleSettings] = options["autoUpgradePolicy"]

Expand Down Expand Up @@ -904,6 +910,17 @@ def validate_cache_staleness_value(max_integrated_cache_staleness: Any) -> None:
"integer greater than or equal to zero")


def validate_shard_key_value(shard_key: str) -> None:
"""Validate that shard key contains only alphanumeric characters and hyphens."""
if not isinstance(shard_key, str):
raise TypeError("Parameter 'dedicated_gateway_shard_key' must be a string")
if not shard_key:
raise ValueError("Parameter 'dedicated_gateway_shard_key' cannot be empty")
if not all(c.isalnum() or c == '-' for c in shard_key):
raise ValueError("Parameter 'dedicated_gateway_shard_key' can only contain "
"alphanumeric characters and hyphens")
Comment thread
simorenoh marked this conversation as resolved.


def _validate_resource(resource: Mapping[str, Any]) -> None:
id_: Optional[str] = resource.get("id")
if id_:
Expand Down
45 changes: 43 additions & 2 deletions sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@
from .. import _utils as utils
from .._availability_strategy_config import _validate_hedging_config
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, GenerateGuidId, validate_cache_staleness_value, validate_shard_key_value)
from .._change_feed.feed_range_internal import FeedRangeInternalEpk

from .._cosmos_responses import CosmosDict, CosmosList
Expand Down Expand Up @@ -321,6 +321,8 @@ async def read_item(
session_token: Optional[str] = None,
initial_headers: Optional[dict[str, str]] = None,
max_integrated_cache_staleness_in_ms: Optional[int] = None,
bypass_integrated_cache: Optional[bool] = None,
dedicated_gateway_shard_key: Optional[str] = None,
priority: Optional[Literal["High", "Low"]] = None,
throughput_bucket: Optional[int] = None,
availability_strategy_config: Optional[dict[str, Any]] = _Unset,
Expand All @@ -342,6 +344,13 @@ async def read_item(
:keyword int max_integrated_cache_staleness_in_ms: The max cache staleness for the integrated cache in
milliseconds. For accounts configured to use the integrated cache, using Session or Eventual consistency,
responses are guaranteed to be no staler than this value.
:keyword bool bypass_integrated_cache: If set to True, the read will be served by the backend and won't be
cached in the dedicated gateway. If set to False or not provided, the integrated cache will be used if
configured. This option only applies to accounts configured with dedicated gateway.
:keyword str dedicated_gateway_shard_key: The shard key to use for the dedicated gateway request. Specifying
the shard key will help route the request to an instance that has cached data for this shard or bypass
the SQLx cache if the correlated instance is down. If not specified, the request will fall back to
randomly selecting an instance. This option only applies to accounts configured with dedicated gateway.
:keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each
request. Once the user has reached their provisioned throughput, low priority requests are throttled
before high priority requests start getting throttled. Feature must first be enabled at the account level.
Expand Down Expand Up @@ -386,6 +395,11 @@ async def read_item(
if max_integrated_cache_staleness_in_ms is not None:
validate_cache_staleness_value(max_integrated_cache_staleness_in_ms)
request_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms
if bypass_integrated_cache is not None:
request_options["bypassIntegratedCache"] = bypass_integrated_cache
if dedicated_gateway_shard_key is not None:
validate_shard_key_value(dedicated_gateway_shard_key)
request_options["dedicatedGatewayShardKey"] = dedicated_gateway_shard_key
await self._get_properties_with_options(request_options)
request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]

Expand All @@ -399,6 +413,8 @@ def read_all_items(
session_token: Optional[str] = None,
initial_headers: Optional[dict[str, str]] = None,
max_integrated_cache_staleness_in_ms: Optional[int] = None,
bypass_integrated_cache: Optional[bool] = None,
dedicated_gateway_shard_key: Optional[str] = None,
Comment on lines 413 to +417
Copy link

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

aio.ContainerProxy.read_all_items adds new parameters bypass_integrated_cache and dedicated_gateway_shard_key, but the method docstring below the signature does not document these keywords. Please update the docstring so async users discover the new options via help()/generated docs.

Copilot uses AI. Check for mistakes.
Copy link
Copy Markdown
Member

@simorenoh simorenoh Apr 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@copilot apply changes based on this feedback

Copy link
Copy Markdown
Contributor Author

Copilot AI Apr 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added missing docstrings for bypass_integrated_cache and dedicated_gateway_shard_key parameters in async read_all_items method. Also addressed other review feedback:

  • Fixed maxIntegratedCacheStaleness null check to support value 0
  • Added type validation for shard key
  • Updated runtime query_items docstring
  • Added comprehensive tests

Commit: cf89cc9

priority: Optional[Literal["High", "Low"]] = None,
throughput_bucket: Optional[int] = None,
availability_strategy_config: Optional[dict[str, Any]] = _Unset,
Expand All @@ -414,6 +430,13 @@ def read_all_items(
:keyword int max_integrated_cache_staleness_in_ms: The max cache staleness for the integrated cache in
milliseconds. For accounts configured to use the integrated cache, using Session or Eventual consistency,
responses are guaranteed to be no staler than this value.
:keyword bool bypass_integrated_cache: If set to True, the read will be served by the backend and won't be
cached in the dedicated gateway. If set to False or not provided, the integrated cache will be used if
configured. This option only applies to accounts configured with dedicated gateway.
:keyword str dedicated_gateway_shard_key: The shard key to use for the dedicated gateway request. Specifying
the shard key will help route the request to an instance that has cached data for this shard or bypass
the SQLx cache if the correlated instance is down. If not specified, the request will fall back to
randomly selecting an instance. This option only applies to accounts configured with dedicated gateway.
:keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for each
request. Once the user has reached their provisioned throughput, low priority requests are throttled
before high priority requests start getting throttled. Feature must first be enabled at the account level.
Expand Down Expand Up @@ -442,9 +465,14 @@ def read_all_items(
feed_options = _build_options(kwargs)
if max_item_count is not None:
feed_options["maxItemCount"] = max_item_count
if max_integrated_cache_staleness_in_ms:
if max_integrated_cache_staleness_in_ms is not None:
validate_cache_staleness_value(max_integrated_cache_staleness_in_ms)
feed_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms
if bypass_integrated_cache is not None:
feed_options["bypassIntegratedCache"] = bypass_integrated_cache
if dedicated_gateway_shard_key is not None:
validate_shard_key_value(dedicated_gateway_shard_key)
feed_options["dedicatedGatewayShardKey"] = dedicated_gateway_shard_key
response_hook = kwargs.pop("response_hook", None)
if response_hook and hasattr(response_hook, "clear"):
response_hook.clear()
Expand Down Expand Up @@ -537,6 +565,8 @@ def query_items(
enable_scan_in_query: Optional[bool] = None,
initial_headers: Optional[dict[str, str]] = None,
max_integrated_cache_staleness_in_ms: Optional[int] = None,
bypass_integrated_cache: Optional[bool] = None,
dedicated_gateway_shard_key: Optional[str] = None,
max_item_count: Optional[int] = None,
parameters: Optional[list[dict[str, object]]] = None,
partition_key: PartitionKeyType,
Expand Down Expand Up @@ -624,6 +654,8 @@ def query_items(
feed_range: dict[str, Any],
initial_headers: Optional[dict[str, str]] = None,
max_integrated_cache_staleness_in_ms: Optional[int] = None,
bypass_integrated_cache: Optional[bool] = None,
dedicated_gateway_shard_key: Optional[str] = None,
max_item_count: Optional[int] = None,
parameters: Optional[list[dict[str, object]]] = None,
populate_index_metrics: Optional[bool] = None,
Expand Down Expand Up @@ -706,6 +738,8 @@ def query_items(
enable_scan_in_query: Optional[bool] = None,
initial_headers: Optional[dict[str, str]] = None,
max_integrated_cache_staleness_in_ms: Optional[int] = None,
bypass_integrated_cache: Optional[bool] = None,
dedicated_gateway_shard_key: Optional[str] = None,
max_item_count: Optional[int] = None,
parameters: Optional[list[dict[str, object]]] = None,
populate_index_metrics: Optional[bool] = None,
Expand Down Expand Up @@ -868,6 +902,13 @@ def query_items(
max_integrated_cache_staleness_in_ms = kwargs.pop("max_integrated_cache_staleness_in_ms")
validate_cache_staleness_value(max_integrated_cache_staleness_in_ms)
feed_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms
if utils.valid_key_value_exist(kwargs, "bypass_integrated_cache"):
bypass_integrated_cache = kwargs.pop("bypass_integrated_cache")
feed_options["bypassIntegratedCache"] = bypass_integrated_cache
if utils.valid_key_value_exist(kwargs, "dedicated_gateway_shard_key"):
dedicated_gateway_shard_key = kwargs.pop("dedicated_gateway_shard_key")
validate_shard_key_value(dedicated_gateway_shard_key)
feed_options["dedicatedGatewayShardKey"] = dedicated_gateway_shard_key
Comment on lines +905 to +911
Copy link

Copilot AI Feb 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Async dedicated gateway headers are now supported (bypass_integrated_cache / dedicated_gateway_shard_key), but there are no tests covering that the async query_items path propagates these options into request headers. Adding coverage in the async headers tests would help ensure parity with the sync behavior.

Copilot uses AI. Check for mistakes.
Copy link
Copy Markdown
Member

@simorenoh simorenoh Apr 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@copilot apply changes based on this feedback

Copy link
Copy Markdown
Contributor Author

Copilot AI Apr 13, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added async query_items test in test_headers_async.py that verifies bypass_integrated_cache and dedicated_gateway_shard_key headers are correctly propagated. The test uses async iteration and follows the same pattern as sync tests. Commit: cf89cc9

if utils.valid_key_value_exist(kwargs, "continuation_token_limit"):
feed_options["responseContinuationTokenLimitInKb"] = kwargs.pop("continuation_token_limit")

Expand Down
Loading
Loading