diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 038185919c0c..7a428a20fb56 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -2,13 +2,18 @@ > [!IMPORTANT] > We strongly recommend that customers use at least version 4.16.3 of `azure-cosmos`. -### 4.16.5 (Unreleased) +### 4.17.0 (Unreleased) #### Features Added +* Added the `enable_compact_utf8_item_writes` client option. Set it to `True` to reduce item write request sizes by + serializing valid Unicode as compact UTF-8 for create, upsert, replace, patch, and transactional batch operations. + See [PR 48914](https://github.com/Azure/azure-sdk-for-python/pull/48914). #### Breaking Changes #### Bugs Fixed +* Fixed sync and async item PATCH requests to explicitly send the registered `application/json-patch+json` + content type. See [PR 48914](https://github.com/Azure/azure-sdk-for-python/pull/48914). #### Other Changes diff --git a/sdk/cosmos/azure-cosmos/README.md b/sdk/cosmos/azure-cosmos/README.md index f7755f7b6fd9..e8703adece6e 100644 --- a/sdk/cosmos/azure-cosmos/README.md +++ b/sdk/cosmos/azure-cosmos/README.md @@ -81,6 +81,30 @@ KEY = os.environ['ACCOUNT_KEY'] client = CosmosClient(URL, credential=KEY) ``` +### Compact UTF-8 item writes + +By default, the client escapes non-ASCII characters in JSON item bodies, so each such character is sent as a +`\uXXXX` escape sequence. A CJK character that occupies 3 bytes as UTF-8 therefore occupies 6 bytes on the wire, and +a 4-byte emoji occupies 12. For Unicode-heavy items this expansion increases the size of the request sent over the +network and can push a request past the 2 MiB request size limit. Applications writing such items can opt in to +compact UTF-8 serialization when creating the client: + +```python +client = CosmosClient( + URL, + credential=KEY, + enable_compact_utf8_item_writes=True, +) +``` + +The option is disabled by default and applies to create, upsert, replace, patch, and transactional batch item +bodies. Queries, control-plane requests, and responses are unchanged. Semantic request headers - including the +partition-key header - are also unchanged; body-derived headers such as `Content-Length` necessarily reflect the +compact byte count and are recalculated accordingly. Both representations describe the +same JSON document, so the values stored in the service and returned on reads are identical - only the encoding of +the outgoing request body differs. See the [synchronous][sample_compact_utf8_item_writes] and +[asynchronous][sample_compact_utf8_item_writes_async] samples for complete examples. + ### AAD Authentication You can also authenticate a client utilizing your service principal's AAD credentials and the azure identity package. @@ -1268,6 +1292,8 @@ For more extensive documentation on the Cosmos DB service, see the [Azure Cosmos [ref_database]: https://aka.ms/azsdk-python-cosmos-ref-database [ref_httpfailure]: https://aka.ms/azsdk-python-cosmos-ref-http-failure [sample_database_mgmt]: https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples/database_management.py +[sample_compact_utf8_item_writes]: https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples/compact_utf8_item_writes.py +[sample_compact_utf8_item_writes_async]: https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples/compact_utf8_item_writes_async.py [sample_document_mgmt]: https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples/document_management.py [sample_document_mgmt_async]: https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples/document_management_async.py [sample_examples_misc]: https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples/examples.py diff --git a/sdk/cosmos/azure-cosmos/api.md b/sdk/cosmos/azure-cosmos/api.md index 541a7aef7487..f8afab1473d5 100644 --- a/sdk/cosmos/azure-cosmos/api.md +++ b/sdk/cosmos/azure-cosmos/api.md @@ -3914,6 +3914,7 @@ namespace azure.cosmos.scripts consistency_level: Optional[str] = None, availability_strategy: Union[bool, dict[str, Any]] = False, availability_strategy_executor: Optional[ThreadPoolExecutor] = None, + enable_compact_utf8_item_writes: bool = False, **kwargs: Any ) -> None: ... @@ -4708,6 +4709,7 @@ namespace azure.cosmos.user consistency_level: Optional[str] = None, availability_strategy: Union[bool, dict[str, Any]] = False, availability_strategy_executor: Optional[ThreadPoolExecutor] = None, + enable_compact_utf8_item_writes: bool = False, **kwargs: Any ) -> None: ... diff --git a/sdk/cosmos/azure-cosmos/api.metadata.yml b/sdk/cosmos/azure-cosmos/api.metadata.yml index fd25fd0f7de8..53c04dd49bf8 100644 --- a/sdk/cosmos/azure-cosmos/api.metadata.yml +++ b/sdk/cosmos/azure-cosmos/api.metadata.yml @@ -1,3 +1,4 @@ -apiMdSha256: 1538a79b2c2da38fb83fd37ccea945314bf2f34218acda3d8f15e2f0268f3040 -parserVersion: 0.3.30 -pythonVersion: 3.13.14 +apiMdSha256: 54ea9d4c0c0f08c381ee7e69a7468ed9c08e5fb55c7af0fc1b9f6d0d6f770a86 +packageVersion: 4.17.0 +parserVersion: 0.3.31 +pythonVersion: 3.10.21 diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py index 297e62d69d47..69804f8d2489 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py @@ -394,7 +394,12 @@ def GetHeaders( # pylint: disable=too-many-statements,too-many-branches authorization = urllib_quote(authorization, "-_.!~*'()") headers[http_constants.HttpHeaders.Authorization] = authorization - if verb in ("post", "put"): + if verb == "patch": + if not headers.get(http_constants.HttpHeaders.ContentType): + headers[http_constants.HttpHeaders.ContentType] = ( + _runtime_constants.MediaTypes.JsonPatch + ) + elif verb in ("post", "put"): if not headers.get(http_constants.HttpHeaders.ContentType): headers[http_constants.HttpHeaders.ContentType] = _runtime_constants.MediaTypes.Json diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index f55238a92983..54347c770b6f 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -134,7 +134,7 @@ class _QueryCompatibilityMode: _DefaultStringHashPrecision = 3 _DefaultStringRangePrecision = -1 - def __init__( # pylint: disable=too-many-statements + def __init__( # pylint: disable=too-many-statements, too-many-locals self, url_connection: str, auth: CredentialDict, @@ -142,6 +142,7 @@ def __init__( # pylint: disable=too-many-statements consistency_level: Optional[str] = None, availability_strategy: Union[bool, dict[str, Any]] = False, availability_strategy_executor: Optional[ThreadPoolExecutor] = None, + enable_compact_utf8_item_writes: bool = False, **kwargs: Any ) -> None: """ @@ -160,10 +161,15 @@ def __init__( # pylint: disable=too-many-statements The availability strategy configuration for routing requests across regions. :param concurrent.futures.ThreadPoolExecutor availability_strategy_executor: The thread pool executor for handling availability strategy requests. + :param bool enable_compact_utf8_item_writes: + Whether item write bodies should use compact UTF-8 serialization. :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for the client. 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. """ + self._enable_compact_utf8_item_writes = _utils._validate_enable_compact_utf8_item_writes( + enable_compact_utf8_item_writes + ) self.client_id = str(uuid.uuid4()) self.url_connection = url_connection self.availability_strategy: Union[CrossRegionHedgingStrategy, None] =\ diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_runtime_constants.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_runtime_constants.py index 249702ebb014..55125016f6c4 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_runtime_constants.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_runtime_constants.py @@ -35,6 +35,7 @@ class MediaTypes(object): ImagePng = "image/png" JavaScript = "application/x-javascript" Json = "application/json" + JsonPatch = "application/json-patch+json" OctetStream = "application/octet-stream" QueryJson = "application/query+json" SQL = "application/sql" diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_synchronized_request.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_synchronized_request.py index 1a7e24e5feba..cd1ee789f8c6 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_synchronized_request.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_synchronized_request.py @@ -21,6 +21,7 @@ """Synchronized request in the Azure Cosmos database service. """ +# cspell:ignore surrogatepass import copy import json import time @@ -37,6 +38,17 @@ from ._request_object import RequestObject from .documents import _OperationType +_ITEM_BODY_WRITE_OPERATIONS = frozenset(( + _OperationType.Create, + _OperationType.Upsert, + _OperationType.Replace, + _OperationType.Patch, + # Batch membership is necessary but not sufficient: a batch made up only of + # read and delete operations carries no item body. _should_escape_non_ascii_ + # in_request_body refines this entry with _batch_contains_item_body below. + _OperationType.Batch, +)) + # cspell:ignore ppaf def _is_readable_stream(obj): """Checks whether obj is a file-like readable stream. @@ -50,26 +62,104 @@ def _is_readable_stream(obj): return False -def _request_body_from_data(data): - """Gets request body from data. +def _request_body_from_data(data, ensure_ascii=True): + """Convert supported request data into an HTTP body. + + Dictionaries, lists, and tuples are serialized as compact JSON. Other + supported body types are returned unchanged. + + When ``ensure_ascii`` is False the serialized body is returned as UTF-8 + ``bytes``. The escaped default and caller-supplied pre-serialized strings + keep returning ``str``, preserving existing behavior. - When `data` is dict and list into unicode string; otherwise return `data` - without making any change. + Returning bytes is required for correctness, not just convenience. The sync + Requests transport forwards the body to ``requests`` unchanged, and the + supported dependency range still permits urllib3 1.x, whose connection path + hands ``str`` bodies to ``http.client``. ``http.client`` encodes them as + Latin-1, so a compact body would be sent with the wrong encoding: text that + has no Latin-1 representation (any CJK or emoji) raises ``UnicodeEncodeError`` + before the request is sent, and text that does have one is written in Latin-1 + rather than UTF-8, so 'é' goes out as the single byte 0xE9 instead of + 0xC3 0xA9 and the service receives different characters than the caller + supplied. Handing the transport the encoded bytes removes every subsequent + re-encoding step. :param Union[str, unicode, file-like stream object, dict, list, None] data: - :returns: the json dump data. - :rtype: Union[str, unicode, file-like stream object, None] + :param bool ensure_ascii: Whether non-ASCII characters should be escaped. + :returns: The serialized or unchanged request body. + :rtype: Union[str, bytes, file-like stream object, None] """ if data is None or isinstance(data, str) or _is_readable_stream(data): return data if isinstance(data, (dict, list, tuple)): - json_dumped = json.dumps(data, separators=(",", ":")) - - return json_dumped + if ensure_ascii: + return json.dumps(data, separators=(",", ":")) + json_dumped = json.dumps(data, separators=(",", ":"), ensure_ascii=False) + try: + # Encode once; callers derive Content-Length directly from these bytes. + encoded_body = json_dumped.encode("utf-8") + except UnicodeEncodeError: + # Rare path for text originating from UTF-16 APIs. Combine adjacent + # high/low surrogate pairs into their Unicode scalar while preserving + # unpaired surrogates, then escape only those remaining invalid code + # units as valid JSON \uXXXX sequences. + normalized_body = json_dumped.encode( + "utf-16-le", "surrogatepass" + ).decode("utf-16-le", "surrogatepass") + encoded_body = normalized_body.encode("utf-8", "backslashreplace") + # Send exactly these bytes, so no transport re-encodes them. + return encoded_body return None +def _batch_contains_item_body(batch_operations): + """Return whether a formatted transactional batch contains an item body. + + Read and delete batch operations contain only an item ID. Create, upsert, + replace, and patch operations contain a resourceBody. + + :param list[dict[str, object]] batch_operations: The formatted batch operations. + :returns: Whether at least one operation contains an item body. + :rtype: bool + """ + return ( + isinstance(batch_operations, (list, tuple)) + and any( + isinstance(operation, dict) and "resourceBody" in operation + for operation in batch_operations + ) + ) + + +def _should_escape_non_ascii_in_request_body(client, request_params, request_data): + """Decide whether a request body must keep non-ASCII characters escaped. + + Compact UTF-8 is only used when the client opted in and the request is one + of the item write operations the option is scoped to. A transactional batch + uses compact UTF-8 only when at least one operation contains an item body; + read/delete-only batches keep the escaped form. Every other request, + including control-plane bodies and queries, also keeps the escaped form. + + :param object client: the client connection issuing the request. + :param ~azure.cosmos._request_object.RequestObject request_params: the request parameters. + :param object request_data: The body data that will be serialized. + :returns: whether non-ASCII characters should be escaped in the body. + :rtype: bool + """ + # getattr keeps the safe (escaped) default for any caller that supplies a + # client object without the option, e.g. custom or legacy connections. + return ( + not getattr(client, "_enable_compact_utf8_item_writes", False) + or request_params.resource_type != http_constants.ResourceType.Document + or request_params.operation_type not in _ITEM_BODY_WRITE_OPERATIONS + or ( + request_params.operation_type == _OperationType.Batch + and not _batch_contains_item_body(request_data) + ) + ) + + def _Request(global_endpoint_manager, request_params, connection_policy, pipeline_client, request, **kwargs): # pylint: disable=too-many-statements """Makes one http request using the requests module. @@ -285,12 +375,21 @@ def SynchronizedRequest( :return: tuple of (result, headers) :rtype: tuple of (dict dict) """ - request.data = _request_body_from_data(request_data) - if request.data and isinstance(request.data, str): + request.data = _request_body_from_data( + request_data, + ensure_ascii=_should_escape_non_ascii_in_request_body(client, request_params, request_data) + ) + if isinstance(request.data, (bytes, bytearray)): + # Compact UTF-8 bodies reach the transport as their final bytes, so + # Content-Length cannot drift from the wire body. + request.headers[http_constants.HttpHeaders.ContentLength] = len(request.data) + elif request.data and isinstance(request.data, str): # Use UTF-8 byte length, not str length (code-point count), so the # header matches the bytes the transport actually writes for any # non-ASCII payload. - request.headers[http_constants.HttpHeaders.ContentLength] = len(request.data.encode("utf-8")) + request.headers[http_constants.HttpHeaders.ContentLength] = len( + request.data.encode("utf-8") + ) elif request.data is None: request.headers[http_constants.HttpHeaders.ContentLength] = 0 diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py index 12ff8cd118e6..8c62105fad98 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_utils.py @@ -35,6 +35,21 @@ # cspell:ignore ppcb # pylint: disable=protected-access + +def _validate_enable_compact_utf8_item_writes(value: Any) -> bool: + """Validate the compact UTF-8 item-write setting. + + :param value: The setting to validate. + :type value: Any + :returns: The validated setting. + :rtype: bool + :raises TypeError: If the setting is not a bool. + """ + if not isinstance(value, bool): + raise TypeError("enable_compact_utf8_item_writes must be a bool.") + return value + + def get_user_agent(suffix: Optional[str] = None) -> str: os_name = safe_user_agent_header(platform.platform()) python_version = safe_user_agent_header(platform.python_version()) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_version.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_version.py index cca0a5dbc553..9676fd1ff1fc 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_version.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_version.py @@ -19,4 +19,4 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -VERSION = "4.16.5" +VERSION = "4.17.0" diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py index ce7d5a44536c..3f65b5a4dbd4 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py @@ -36,7 +36,11 @@ from .._constants import _Constants from .._request_object import RequestObject from .._response_decoding import decode_response_body_for_status -from .._synchronized_request import _request_body_from_data, _replace_url_prefix +from .._synchronized_request import ( + _request_body_from_data, + _replace_url_prefix, + _should_escape_non_ascii_in_request_body, +) from ..documents import _OperationType # cspell:ignore ppaf @@ -238,12 +242,21 @@ async def AsynchronousRequest( :return: tuple of (result, headers) :rtype: tuple of (dict dict) """ - request.data = _request_body_from_data(request_data) - if request.data and isinstance(request.data, str): + request.data = _request_body_from_data( + request_data, + ensure_ascii=_should_escape_non_ascii_in_request_body(client, request_params, request_data) + ) + if isinstance(request.data, (bytes, bytearray)): + # Compact UTF-8 bodies reach the transport as their final bytes, so + # Content-Length cannot drift from the wire body. + request.headers[http_constants.HttpHeaders.ContentLength] = len(request.data) + elif request.data and isinstance(request.data, str): # Use UTF-8 byte length, not str length (code-point count), so the # header matches the bytes the transport actually writes for any # non-ASCII payload. - request.headers[http_constants.HttpHeaders.ContentLength] = len(request.data.encode("utf-8")) + request.headers[http_constants.HttpHeaders.ContentLength] = len( + request.data.encode("utf-8") + ) elif request.data is None: request.headers[http_constants.HttpHeaders.ContentLength] = 0 diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py index 522a3b65ba9b..129a4440e198 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py @@ -192,6 +192,9 @@ class CosmosClient: # pylint: disable=client-accepts-api-version-keyword Default value is False (hedging disabled). :paramtype availability_strategy: Union[bool, dict[str, Any]] :keyword int availability_strategy_max_concurrency: The max concurrency for parallel requests. + :keyword bool enable_compact_utf8_item_writes: + Use compact UTF-8 when serializing item bodies for create, upsert, replace, patch, and transactional batch + operations. The default is False. .. admonition:: Example: @@ -216,6 +219,7 @@ def __init__( ) -> None: """Instantiate a new CosmosClient.""" auth = _build_auth(credential) + enable_compact_utf8_item_writes = kwargs.pop("enable_compact_utf8_item_writes", False) connection_policy = _build_connection_policy(kwargs) self.client_connection = CosmosClientConnection( url_connection=url, @@ -224,6 +228,7 @@ def __init__( connection_policy=connection_policy, availability_strategy=availability_strategy, availability_strategy_max_concurrency=availability_strategy_max_concurrency, + enable_compact_utf8_item_writes=enable_compact_utf8_item_writes, **kwargs ) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py index 92e9c107d390..3a766edefe70 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py @@ -139,6 +139,7 @@ def __init__( # pylint: disable=too-many-statements consistency_level: Optional[str] = None, availability_strategy: Union[bool, dict[str, Any]] = False, availability_strategy_max_concurrency: Optional[int] = None, + enable_compact_utf8_item_writes: bool = False, **kwargs: Any ) -> None: """ @@ -153,10 +154,15 @@ def __init__( # pylint: disable=too-many-statements The connection policy for the client. :param documents.ConsistencyLevel consistency_level: The default consistency policy for client operations. + :param bool enable_compact_utf8_item_writes: + Whether item write bodies should use compact UTF-8 serialization. :keyword Literal["High", "Low"] priority: Priority based execution allows users to set a priority for the client. 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. """ + self._enable_compact_utf8_item_writes = _utils._validate_enable_compact_utf8_item_writes( + enable_compact_utf8_item_writes + ) self.client_id = str(uuid.uuid4()) self.url_connection = url_connection self.availability_strategy: Union[CrossRegionHedgingStrategy, None] =\ diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py b/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py index 08f32f7b4c93..00901a6bc463 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py @@ -215,6 +215,9 @@ class CosmosClient: # pylint: disable=client-accepts-api-version-keyword :paramtype availability_strategy: Union[bool, dict[str, Any]] :keyword ~concurrent.futures.thread.ThreadPoolExecutor availability_strategy_executor: Optional ThreadPoolExecutor for handling concurrent operations. + :keyword bool enable_compact_utf8_item_writes: + Use compact UTF-8 when serializing item bodies for create, upsert, replace, patch, and transactional batch + operations. The default is False. .. admonition:: Example: @@ -237,6 +240,7 @@ def __init__( """ auth = _build_auth(credential) + enable_compact_utf8_item_writes = kwargs.pop("enable_compact_utf8_item_writes", False) connection_policy = _build_connection_policy(kwargs) self.client_connection = CosmosClientConnection( url_connection=url, @@ -245,6 +249,7 @@ def __init__( connection_policy=connection_policy, availability_strategy=kwargs.pop("availability_strategy", False), availability_strategy_executor=kwargs.pop("availability_strategy_executor", None), + enable_compact_utf8_item_writes=enable_compact_utf8_item_writes, **kwargs ) diff --git a/sdk/cosmos/azure-cosmos/samples/README.md b/sdk/cosmos/azure-cosmos/samples/README.md index 6e66a4437eba..2b52a74d92a4 100644 --- a/sdk/cosmos/azure-cosmos/samples/README.md +++ b/sdk/cosmos/azure-cosmos/samples/README.md @@ -57,6 +57,11 @@ The following are code samples that show common scenario operations with the Azu * [diagnostics_filter_sample.py](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples/diagnostics_filter_sample.py) - Example using logging filters for diagnostics filtering +* [compact_utf8_item_writes.py](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples/compact_utf8_item_writes.py) - Example demonstrating the `enable_compact_utf8_item_writes` client option: + * Sending item write bodies as compact UTF-8 instead of `\uXXXX` escape sequences + * Comparing the resulting request body size for Unicode-heavy items + * Async version: [compact_utf8_item_writes_async.py](https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/cosmos/azure-cosmos/samples/compact_utf8_item_writes_async.py) + ## Prerequisites * Python 3.8+ * You must have an [Azure subscription](https://azure.microsoft.com/free/) and an diff --git a/sdk/cosmos/azure-cosmos/samples/compact_utf8_item_writes.py b/sdk/cosmos/azure-cosmos/samples/compact_utf8_item_writes.py new file mode 100644 index 000000000000..5f3d5f7ed489 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/samples/compact_utf8_item_writes.py @@ -0,0 +1,136 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- +import json +import uuid + +import azure.cosmos.cosmos_client as cosmos_client +import azure.cosmos.exceptions as exceptions +from azure.cosmos.partition_key import PartitionKey + +import config + +# ---------------------------------------------------------------------------------------------- +# Prerequisites - +# +# 1. An Azure Cosmos account - +# https://azure.microsoft.com/documentation/articles/documentdb-create-account/ +# +# 2. Microsoft Azure Cosmos PyPi package - +# https://pypi.python.org/pypi/azure-cosmos/ +# ---------------------------------------------------------------------------------------------- +# Sample - demonstrates the `enable_compact_utf8_item_writes` client option. +# +# By default, the SDK serializes item bodies with `ensure_ascii=True`, so every non-ASCII +# character is expanded into a `\uXXXX` escape sequence. A single CJK character occupies 3 bytes +# as raw UTF-8 but 6 bytes once escaped, and a 4-byte emoji occupies 12 bytes as an escaped +# surrogate pair. For Unicode-heavy items that expansion increases the size of the request sent +# over the network and can push a request past the 2 MiB request size limit. +# +# Setting `enable_compact_utf8_item_writes=True` sends the body as compact UTF-8 instead. The +# option is: +# * opt-in - the default (False) keeps the existing escaped wire format exactly as-is. +# * scoped to item writes - create, upsert, replace, patch, and transactional batch. +# Queries, control-plane bodies, and the partition-key header are unaffected. +# +# Both representations describe the same JSON document, so the stored item and the values +# returned on reads are identical. Only the bytes on the wire change. +# ---------------------------------------------------------------------------------------------- +# Note - +# +# Running this sample will create the configured Database if it does not already exist, and +# leaves that Database in place. It creates and then deletes a Container within it. Each time a +# Container is created the account will be billed for 1 hour of usage based on the provisioned +# throughput (RU/s) of that account. +# ---------------------------------------------------------------------------------------------- + +HOST = config.settings['host'] +MASTER_KEY = config.settings['master_key'] +DATABASE_ID = config.settings['database_id'] + +# A dedicated container for this sample. The shared container id from config is created by other +# samples with a '/id' partition key path, so this sample uses its own container to stay +# independent of the order in which the samples are run. The id is made unique per run and the +# container is created with `create_container` (not `create_container_if_not_exists`) so this +# sample can never take over - and then delete - a container it did not create. +CONTAINER_ID_PREFIX = 'compact-utf8-item-writes-' +PARTITION_KEY_VALUE = 'compact-utf8-sample' + + +def show_wire_size_difference(item): + """Print the request body size with and without ASCII escape expansion.""" + escaped = json.dumps(item, separators=(',', ':')).encode('utf-8') + compact = json.dumps(item, separators=(',', ':'), ensure_ascii=False).encode('utf-8') + print('Escaped body (default): {0} bytes'.format(len(escaped))) + print('Compact UTF-8 body (opt-in): {0} bytes'.format(len(compact))) + print('Saved: {0} bytes'.format(len(escaped) - len(compact))) + + +def write_compact_utf8_item(container): + """Write a Unicode-heavy item without ASCII escape expansion.""" + item_id = 'compact-utf8-' + str(uuid.uuid4()) + item = { + 'id': item_id, + 'pk': PARTITION_KEY_VALUE, + # Repeated so the size difference is easy to see in the output. + 'content': 'Customer text: 日本語 🎉' * 200, # cspell:disable-line + } + + show_wire_size_difference(item) + + created_item = container.create_item(item) + print('Created compact UTF-8 item: {0}'.format(created_item['id'])) + + # The value stored in the service is the same either way - only the + # serialization of the outgoing request changed. + read_item = container.read_item(item_id, partition_key=PARTITION_KEY_VALUE) + print('Round trip preserved content: {0}'.format( + read_item['content'] == item['content'])) + + container.delete_item(item_id, partition_key=PARTITION_KEY_VALUE) + + +def run_sample(): + """Run the compact UTF-8 item write sample.""" + container_id = CONTAINER_ID_PREFIX + str(uuid.uuid4()) + # The option is set once at client construction and applies to every item + # write issued by this client. It is disabled by default. + with cosmos_client.CosmosClient( + HOST, + {'masterKey': MASTER_KEY}, + enable_compact_utf8_item_writes=True, + ) as client: + db = None + container = None + try: + # setup database for this sample + db = client.create_database_if_not_exists(id=DATABASE_ID) + # setup container for this sample - a uniquely named container created by this run only, + # so the cleanup below can never delete a pre-existing container or its data + container = db.create_container( + id=container_id, + partition_key=PartitionKey(path='/pk', kind='Hash'), + ) + + write_compact_utf8_item(container) + + except exceptions.CosmosHttpResponseError as e: + print('\nrun_sample has caught an error. {0}'.format(e.message)) + + finally: + # cleanup container after sample, even if the sample failed part way through + if db is not None and container is not None: + try: + db.delete_container(container) + + except exceptions.CosmosHttpResponseError as cleanup_error: + print('\nFailed to delete the sample container. {0}'.format( + cleanup_error.message)) + + print("\nrun_sample done") + + +if __name__ == '__main__': + run_sample() diff --git a/sdk/cosmos/azure-cosmos/samples/compact_utf8_item_writes_async.py b/sdk/cosmos/azure-cosmos/samples/compact_utf8_item_writes_async.py new file mode 100644 index 000000000000..bf253a2000b0 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/samples/compact_utf8_item_writes_async.py @@ -0,0 +1,138 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See LICENSE.txt in the project root for +# license information. +# ------------------------------------------------------------------------- +import asyncio +import json +import uuid + +import azure.cosmos.exceptions as exceptions +from azure.cosmos.aio import CosmosClient +from azure.cosmos.partition_key import PartitionKey + +import config + +# ---------------------------------------------------------------------------------------------- +# Prerequisites - +# +# 1. An Azure Cosmos account - +# https://azure.microsoft.com/documentation/articles/documentdb-create-account/ +# +# 2. Microsoft Azure Cosmos PyPi package - +# https://pypi.python.org/pypi/azure-cosmos/ +# ---------------------------------------------------------------------------------------------- +# Sample - demonstrates the `enable_compact_utf8_item_writes` client option with the async +# client. +# +# By default, the SDK serializes item bodies with `ensure_ascii=True`, so every non-ASCII +# character is expanded into a `\uXXXX` escape sequence. A single CJK character occupies 3 bytes +# as raw UTF-8 but 6 bytes once escaped, and a 4-byte emoji occupies 12 bytes as an escaped +# surrogate pair. For Unicode-heavy items that expansion increases the size of the request sent +# over the network and can push a request past the 2 MiB request size limit. +# +# Setting `enable_compact_utf8_item_writes=True` sends the body as compact UTF-8 instead. The +# option is: +# * opt-in - the default (False) keeps the existing escaped wire format exactly as-is. +# * scoped to item writes - create, upsert, replace, patch, and transactional batch. +# Queries, control-plane bodies, and the partition-key header are unaffected. +# +# Both representations describe the same JSON document, so the stored item and the values +# returned on reads are identical. Only the bytes on the wire change. +# ---------------------------------------------------------------------------------------------- +# Note - +# +# Running this sample will create the configured Database if it does not already exist, and +# leaves that Database in place. It creates and then deletes a Container within it. Each time a +# Container is created the account will be billed for 1 hour of usage based on the provisioned +# throughput (RU/s) of that account. +# ---------------------------------------------------------------------------------------------- + +HOST = config.settings['host'] +MASTER_KEY = config.settings['master_key'] +DATABASE_ID = config.settings['database_id'] + +# A dedicated container for this sample. The shared container id from config is created by other +# samples with a '/id' partition key path, so this sample uses its own container to stay +# independent of the order in which the samples are run. The id is made unique per run and the +# container is created with `create_container` (not `create_container_if_not_exists`) so this +# sample can never take over - and then delete - a container it did not create. +CONTAINER_ID_PREFIX = 'compact-utf8-item-writes-async-' +PARTITION_KEY_VALUE = 'compact-utf8-sample-async' + + +def show_wire_size_difference(item): + """Print the request body size with and without ASCII escape expansion.""" + escaped = json.dumps(item, separators=(',', ':')).encode('utf-8') + compact = json.dumps(item, separators=(',', ':'), ensure_ascii=False).encode('utf-8') + print('Escaped body (default): {0} bytes'.format(len(escaped))) + print('Compact UTF-8 body (opt-in): {0} bytes'.format(len(compact))) + print('Saved: {0} bytes'.format(len(escaped) - len(compact))) + + +async def write_compact_utf8_item(container): + """Write a Unicode-heavy item without ASCII escape expansion.""" + item_id = 'compact-utf8-async-' + str(uuid.uuid4()) + item = { + 'id': item_id, + 'pk': PARTITION_KEY_VALUE, + # Repeated so the size difference is easy to see in the output. + 'content': 'Customer text: 日本語 🎉' * 200, # cspell:disable-line + } + + show_wire_size_difference(item) + + created_item = await container.create_item(item) + print('Created compact UTF-8 item: {0}'.format(created_item['id'])) + + # The value stored in the service is the same either way - only the + # serialization of the outgoing request changed. + read_item = await container.read_item(item_id, partition_key=PARTITION_KEY_VALUE) + print('Round trip preserved content: {0}'.format( + read_item['content'] == item['content'])) + + await container.delete_item(item_id, partition_key=PARTITION_KEY_VALUE) + + +async def run_sample(): + """Run the compact UTF-8 item write sample.""" + container_id = CONTAINER_ID_PREFIX + str(uuid.uuid4()) + # The option is set once at client construction and applies to every item + # write issued by this client. It is disabled by default. + async with CosmosClient( + HOST, + {'masterKey': MASTER_KEY}, + enable_compact_utf8_item_writes=True, + ) as client: + db = None + container = None + try: + # setup database for this sample + db = await client.create_database_if_not_exists(id=DATABASE_ID) + # setup container for this sample - a uniquely named container created by this run + # only, so the cleanup below can never delete a pre-existing container or its data + container = await db.create_container( + id=container_id, + partition_key=PartitionKey(path='/pk', kind='Hash'), + ) + + await write_compact_utf8_item(container) + + except exceptions.CosmosHttpResponseError as e: + print('\nrun_sample has caught an error. {0}'.format(e.message)) + + finally: + # cleanup container after sample, even if the sample failed part way through + if db is not None and container is not None: + try: + await db.delete_container(container) + + except exceptions.CosmosHttpResponseError as cleanup_error: + print('\nFailed to delete the sample container. {0}'.format( + cleanup_error.message)) + + print("\nrun_sample done") + + +if __name__ == '__main__': + asyncio.run(run_sample()) diff --git a/sdk/cosmos/azure-cosmos/tests/test_aad.py b/sdk/cosmos/azure-cosmos/tests/test_aad.py index 437ad59aa519..31d733c13507 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_aad.py +++ b/sdk/cosmos/azure-cosmos/tests/test_aad.py @@ -6,6 +6,7 @@ import os import time import unittest +import uuid from io import StringIO import pytest @@ -119,6 +120,46 @@ def test_aad_credentials(self): print("Query result: " + str(query_results[0])) self.container.delete_item(item='Item_0', partition_key='pk') + @_skip_on_non_emulator + def test_compact_utf8_item_write_with_aad(self): + """Verify compact UTF-8 item writes with a token credential.""" + document = { + 'id': 'aad-compact-utf8-' + str(uuid.uuid4()), + 'pk': 'pk', + 'content': '日本🎉', + } + captured = {} + + def capture_body(request): + captured['body'] = request.http_request.body + + with cosmos_client.CosmosClient( + self.host, + self.credential, + enable_compact_utf8_item_writes=True, + ) as client: + container = client.get_database_client( + self.configs.TEST_DATABASE_ID + ).get_container_client( + self.configs.TEST_SINGLE_PARTITION_CONTAINER_ID + ) + created = None + try: + created = container.create_item( + document, + raw_request_hook=capture_body, + ) + + self.assertEqual(created['content'], document['content']) + # Compact bodies reach the transport as UTF-8 bytes, not str. + self.assertIsInstance(captured['body'], bytes) + decoded_body = captured['body'].decode('utf-8') + self.assertIn('日本🎉', decoded_body) + self.assertNotIn('\\u65e5', decoded_body) + finally: + if created is not None: + container.delete_item(document['id'], partition_key='pk') + def _run_with_scope_capture(self, credential_cls, action, *args, **kwargs): scopes_captured = [] original_get_token = credential_cls.get_token diff --git a/sdk/cosmos/azure-cosmos/tests/test_aad_async.py b/sdk/cosmos/azure-cosmos/tests/test_aad_async.py index 20fec64a666b..f8c67dda5336 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_aad_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_aad_async.py @@ -6,6 +6,7 @@ import time import os import unittest +import uuid from io import StringIO import pytest @@ -118,6 +119,46 @@ async def test_aad_credentials_async(self): print("Query result: " + str(query_results[0])) await self.container.delete_item(item='Item_0', partition_key='pk') + @_skip_scope_tests_on_non_emulator + async def test_compact_utf8_item_write_with_aad_async(self): + """Verify compact UTF-8 item writes with an async token credential.""" + document = { + 'id': 'aad-compact-utf8-async-' + str(uuid.uuid4()), + 'pk': 'pk', + 'content': '日本🎉', + } + captured = {} + + def capture_body(request): + captured['body'] = request.http_request.body + + async with CosmosClient( + self.host, + self.credential, + enable_compact_utf8_item_writes=True, + ) as client: + container = client.get_database_client( + self.configs.TEST_DATABASE_ID + ).get_container_client( + self.configs.TEST_SINGLE_PARTITION_CONTAINER_ID + ) + created = None + try: + created = await container.create_item( + document, + raw_request_hook=capture_body, + ) + + self.assertEqual(created['content'], document['content']) + # Compact bodies reach the transport as UTF-8 bytes, not str. + self.assertIsInstance(captured['body'], bytes) + decoded_body = captured['body'].decode('utf-8') + self.assertIn('日本🎉', decoded_body) + self.assertNotIn('\\u65e5', decoded_body) + finally: + if created is not None: + await container.delete_item(document['id'], partition_key='pk') + async def _run_with_scope_capture_async(self, credential_cls, action): scopes_captured = [] diff --git a/sdk/cosmos/azure-cosmos/tests/test_compact_utf8_resilience.py b/sdk/cosmos/azure-cosmos/tests/test_compact_utf8_resilience.py new file mode 100644 index 000000000000..b35ee5b92b9d --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_compact_utf8_resilience.py @@ -0,0 +1,290 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +"""Regression tests for indirect consumers of compact UTF-8 byte bodies.""" +# pylint: disable=invalid-name,missing-function-docstring,protected-access,too-few-public-methods + +import asyncio +from threading import Event +from types import SimpleNamespace +import unittest +from unittest import mock + +import pytest + +from azure.core.pipeline.transport import HttpRequest + +from azure.cosmos import ( + _container_recreate_retry_policy, + _retry_utility, + _synchronized_request, + exceptions, +) +from azure.cosmos._availability_strategy_config import CrossRegionHedgingStrategy +from azure.cosmos._availability_strategy_handler import CrossRegionHedgingHandler +from azure.cosmos._request_object import RequestObject +from azure.cosmos.aio import ( + _asynchronous_availability_strategy_handler, + _retry_utility_async, +) +from azure.cosmos.documents import ConnectionPolicy, _OperationType +from azure.cosmos.http_constants import HttpHeaders, ResourceType, StatusCodes + + +def _compact_request(): + body = _synchronized_request._request_body_from_data( + {"id": "item", "pk": "日本", "text": "café 🎉"}, + ensure_ascii=False, + ) + request = HttpRequest("POST", "https://example.test/dbs/db/colls/container/docs") + request.body = body + request.headers[HttpHeaders.ContentLength] = str(len(body)) + return request + + +def _retry_client(): + client = mock.MagicMock() + client.connection_policy = ConnectionPolicy() + client._container_properties_cache = {} + client.last_response_headers = {} + client.session = None + client._enable_diagnostics_logging = False + return client + + +def _retry_endpoint_manager(async_client=False): + manager = mock.MagicMock() + manager.is_per_partition_automatic_failover_applicable.return_value = False + manager.is_circuit_breaker_applicable.return_value = False + manager.can_use_multiple_write_locations.return_value = False + manager.resolve_service_endpoint_for_partition.return_value = "https://example.test" + manager.location_cache.read_regional_routing_contexts = [] + manager.location_cache._get_applicable_write_regional_routing_contexts.return_value = ["write"] + manager.location_cache._get_applicable_read_regional_routing_contexts.return_value = ["read"] + if async_client: + manager.record_success = mock.AsyncMock() + manager.record_failure = mock.AsyncMock() + manager.record_ppcb_success = mock.AsyncMock() + manager.record_ppcb_failure = mock.AsyncMock() + return manager + + +def _retry_params(): + params = RequestObject(ResourceType.Document, _OperationType.Create, {}) + params.retry_write = 1 + return params + + +def _throttled_exception(): + response = exceptions._InternalCosmosException( + StatusCodes.TOO_MANY_REQUESTS, + {HttpHeaders.RetryAfterInMilliseconds: "0"}, + ) + return exceptions.CosmosHttpResponseError( + status_code=StatusCodes.TOO_MANY_REQUESTS, + response=response, + ) + + +class _PartitionKeyClient: + @staticmethod + def _AddPartitionKey(_container_link, body, _options): + return {"partitionKey": body["pk"]} + + +class _AsyncPartitionKeyClient: + @staticmethod + async def _AddPartitionKey(_container_link, body, _options): + return {"partitionKey": body["pk"]} + + +def _container_recreate_policy(request): + request.headers[HttpHeaders.IntendedCollectionRID] = "old-rid" + return _container_recreate_retry_policy.ContainerRecreateRetryPolicy( + _PartitionKeyClient(), + {"old-rid": {"container_link": "dbs/db/colls/container"}}, + request, + ) + + +@pytest.mark.cosmosEmulator +class TestCompactUtf8Resilience(unittest.TestCase): + """Sync retries, container recreation, and hedging preserve byte bodies.""" + + def test_transient_retry_reuses_exact_bytes_and_content_length(self): + """A throttled write must retry with the original UTF-8 bytes and + their matching Content-Length instead of serializing the body again.""" + request = _compact_request() + attempts = [] + + def execute(_manager, _params, _policy, _pipeline, attempted_request, **_kwargs): + attempts.append( + ( + attempted_request.body, + attempted_request.headers[HttpHeaders.ContentLength], + ) + ) + if len(attempts) == 1: + raise _throttled_exception() + return {}, {} + + _retry_utility.Execute( + _retry_client(), + _retry_endpoint_manager(), + execute, + _retry_params(), + ConnectionPolicy(), + object(), + request, + ) + + self.assertEqual(len(attempts), 2) + self.assertTrue(all(body is request.body for body, _ in attempts)) + self.assertTrue( + all(int(length) == len(request.body) for _, length in attempts) + ) + + def test_container_recreation_extracts_non_ascii_partition_key_from_bytes(self): + """Container-recreation recovery must parse a compact byte body and + recover its non-ASCII partition key without changing its value.""" + request = _compact_request() + policy = _container_recreate_policy(request) + container = {"partitionKey": {"kind": "Hash", "paths": ["/pk"]}} + + partition_key = policy._extract_partition_key( + _PartitionKeyClient(), + container, + request.body, + ) + + self.assertEqual(partition_key, '["\\u65e5\\u672c"]') + + def test_hedged_write_clone_preserves_bytes_and_content_length(self): + """The request copy used for a hedged write must retain the compact + byte body and its matching Content-Length.""" + request = _compact_request() + params = _retry_params() + params.availability_strategy = CrossRegionHedgingStrategy( + {"threshold_ms": 1, "threshold_steps_ms": 1} + ) + captured = {} + + def execute(cloned_params, cloned_request): + captured["params"] = cloned_params + captured["request"] = cloned_request + return {}, {} + + CrossRegionHedgingHandler().execute_single_request_with_delay( + request_params=params, + request=request, + execute_request_fn=execute, + location_index=1, + available_locations=["primary", "secondary"], + complete_status=Event(), + first_request_params_holder=SimpleNamespace(request_params=None), + ) + + cloned_request = captured["request"] + self.assertTrue(captured["params"].is_hedging_request) + self.assertIsNot(cloned_request, request) + self.assertEqual(cloned_request.body, request.body) + self.assertIsInstance(cloned_request.body, bytes) + self.assertEqual( + int(cloned_request.headers[HttpHeaders.ContentLength]), + len(cloned_request.body), + ) + + +@pytest.mark.cosmosEmulator +class TestCompactUtf8ResilienceAsync(unittest.IsolatedAsyncioTestCase): + """Async twins protect the independent retry and hedging implementations.""" + + async def test_transient_retry_reuses_exact_bytes_and_content_length(self): + """An async throttled write must retry with the original UTF-8 bytes + and their matching Content-Length.""" + request = _compact_request() + attempts = [] + + async def execute(_manager, _params, _policy, _pipeline, attempted_request, **_kwargs): + attempts.append( + ( + attempted_request.body, + attempted_request.headers[HttpHeaders.ContentLength], + ) + ) + if len(attempts) == 1: + raise _throttled_exception() + return {}, {} + + await _retry_utility_async.ExecuteAsync( + _retry_client(), + _retry_endpoint_manager(async_client=True), + execute, + _retry_params(), + ConnectionPolicy(), + object(), + request, + ) + + self.assertEqual(len(attempts), 2) + self.assertTrue(all(body is request.body for body, _ in attempts)) + self.assertTrue( + all(int(length) == len(request.body) for _, length in attempts) + ) + + async def test_container_recreation_extracts_non_ascii_partition_key_from_bytes(self): + """Async container-recreation recovery must parse a compact byte body + and recover its non-ASCII partition key.""" + request = _compact_request() + policy = _container_recreate_policy(request) + container = {"partitionKey": {"kind": "Hash", "paths": ["/pk"]}} + + partition_key = await policy._extract_partition_key_async( + _AsyncPartitionKeyClient(), + container, + request.body, + ) + + self.assertEqual(partition_key, '["\\u65e5\\u672c"]') + + async def test_hedged_write_clone_preserves_bytes_and_content_length(self): + """The async hedging copy must retain the compact byte body and its + matching Content-Length.""" + request = _compact_request() + params = _retry_params() + params.availability_strategy = CrossRegionHedgingStrategy( + {"threshold_ms": 1, "threshold_steps_ms": 1} + ) + captured = {} + + async def execute(cloned_params, cloned_request): + captured["params"] = cloned_params + captured["request"] = cloned_request + return {}, {} + + handler = ( + _asynchronous_availability_strategy_handler.CrossRegionAsyncHedgingHandler() + ) + await handler.execute_single_request_with_delay( + request_params=params, + request=request, + execute_request_fn=execute, + location_index=1, + available_locations=["primary", "secondary"], + complete_status=asyncio.Event(), + first_request_params_holder=SimpleNamespace(request_params=None), + ) + + cloned_request = captured["request"] + self.assertTrue(captured["params"].is_hedging_request) + self.assertIsNot(cloned_request, request) + self.assertEqual(cloned_request.body, request.body) + self.assertIsInstance(cloned_request.body, bytes) + self.assertEqual( + int(cloned_request.headers[HttpHeaders.ContentLength]), + len(cloned_request.body), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/tests/test_content_length_encoding.py b/sdk/cosmos/azure-cosmos/tests/test_content_length_encoding.py index 80a2352d3472..7e75bbc9a2b0 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_content_length_encoding.py +++ b/sdk/cosmos/azure-cosmos/tests/test_content_length_encoding.py @@ -6,6 +6,8 @@ import unittest from unittest import mock +import pytest + from azure.cosmos import _synchronized_request, http_constants from azure.cosmos.aio import _asynchronous_request from azure.cosmos.documents import _OperationType @@ -25,6 +27,8 @@ class _DummyRequestParams: + """Stand-in for RequestObject with hedging disabled and a plain item + create, so requests reach the mocked executor directly.""" def __init__(self): self.availability_strategy = None self.is_hedging_request = False @@ -34,22 +38,36 @@ def __init__(self): class _DummyGlobalEndpointManager: + """Endpoint manager stub reporting per-partition failover as off.""" @staticmethod def is_per_partition_automatic_failover_enabled(): return False class _DummyRequest: + """Minimal HttpRequest stand-in capturing body and headers.""" def __init__(self): self.headers = {} self.data = None +class _DummyClient: + """Client stub with the compact UTF-8 option off, so these tests cover + the default serialization path.""" + _enable_compact_utf8_item_writes = False + + +# These tests need no emulator, but the Cosmos CI lane selects tests with +# "-m cosmosEmulator" (see eng/pipelines/templates/stages/cosmos-sdk-client.yml), +# so an unmarked test is silently deselected and would never run. +@pytest.mark.cosmosEmulator class TestContentLengthWiringSync(unittest.TestCase): """Checks the sync request path sets Content-Length to the byte count of the body.""" def _capture_outgoing_request(self, request_data): + """Run one request through the sync path with the retry executor + mocked out, returning the body and Content-Length that would be sent.""" params = _DummyRequestParams() manager = _DummyGlobalEndpointManager() request = _DummyRequest() @@ -66,7 +84,7 @@ def _fake_execute(*args, **kwargs): _synchronized_request._retry_utility, "Execute", side_effect=_fake_execute ): _synchronized_request.SynchronizedRequest( - client=object(), + client=_DummyClient(), request_params=params, global_endpoint_manager=manager, connection_policy=object(), @@ -113,10 +131,12 @@ def test_bytearray_body_is_coerced_to_none_and_content_length_zero(self): self.assertEqual(captured["content_length"], 0) +@pytest.mark.cosmosEmulator class TestContentLengthWiringAsync(unittest.IsolatedAsyncioTestCase): """Async version of the sync class above. Same checks.""" async def _capture_outgoing_request(self, request_data): + """Async twin of the sync capture helper.""" params = _DummyRequestParams() manager = _DummyGlobalEndpointManager() request = _DummyRequest() @@ -135,7 +155,7 @@ async def _fake_execute_async(*args, **kwargs): side_effect=_fake_execute_async, ): await _asynchronous_request.AsynchronousRequest( - client=object(), + client=_DummyClient(), request_params=params, global_endpoint_manager=manager, connection_policy=object(), @@ -146,6 +166,8 @@ async def _fake_execute_async(*args, **kwargs): return captured async def test_str_bodies_set_utf8_byte_content_length(self): + """Async twin: Content-Length is the UTF-8 byte count, and for + multi-byte payloads it differs from the character count.""" for label, payload in _STR_PAYLOADS: with self.subTest(payload=label): captured = await self._capture_outgoing_request(payload) @@ -157,6 +179,7 @@ async def test_str_bodies_set_utf8_byte_content_length(self): self.assertNotEqual(captured["content_length"], len(body)) async def test_none_body_sets_content_length_zero(self): + """Async twin: a request with no body still gets Content-Length 0.""" captured = await self._capture_outgoing_request(None) self.assertEqual(captured["content_length"], 0) @@ -176,4 +199,3 @@ async def test_bytearray_body_is_coerced_to_none_and_content_length_zero(self): if __name__ == "__main__": unittest.main() - diff --git a/sdk/cosmos/azure-cosmos/tests/test_encoding.py b/sdk/cosmos/azure-cosmos/tests/test_encoding.py index 0fd15e0b5e58..6603ed839330 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_encoding.py +++ b/sdk/cosmos/azure-cosmos/tests/test_encoding.py @@ -2,6 +2,7 @@ # The MIT License (MIT) # Copyright (c) Microsoft Corporation. All rights reserved. +import json import unittest import uuid @@ -9,12 +10,20 @@ import azure.cosmos.cosmos_client as cosmos_client import test_config -from azure.cosmos import DatabaseProxy, ContainerProxy +from azure.cosmos import DatabaseProxy, ContainerProxy, exceptions @pytest.mark.cosmosEmulator +@pytest.mark.cosmosLong +@pytest.mark.cosmosAADLong class TestEncoding(unittest.TestCase): - """Test to ensure escaping of non-ascii characters from partition key""" + """Test to ensure escaping of non-ascii characters from partition key. + + Marked for the emulator lane and both live lanes: the compact UTF-8 + item-write coverage below asserts on service behavior (2 MiB request + limit, non-ASCII round trips) that the emulator only approximates, so + these must also run against a live account. + """ host = test_config.TestConfig.host masterKey = test_config.TestConfig.masterKey @@ -26,6 +35,14 @@ class TestEncoding(unittest.TestCase): created_container: ContainerProxy = None key_container: ContainerProxy = None + @staticmethod + def _large_cjk_document(id_prefix): + return { + 'id': id_prefix + str(uuid.uuid4()), + 'pk': 'pk', + 'content': '日' * 400000, + } + @classmethod def setUpClass(cls): if (cls.masterKey == '[YOUR_KEY_HERE]' or @@ -63,6 +80,11 @@ def test_create_document_with_line_separator_para_seperator_next_line_unicodes(s read_doc = self.created_container.read_item(item=created_doc['id'], partition_key='pk') self.assertEqual(read_doc['unicode_content'], test_string) + @pytest.mark.skipif( + test_config.TestConfig.data_auth_mode == 'aad', + reason="Stored-procedure creation is a control-plane operation that needs key auth; " + "the account under test may have local auth disabled.", + ) def test_create_stored_procedure_with_line_separator_para_seperator_next_line_unicodes(self): # scripts.create_stored_procedure and scripts.get_stored_procedure are control-plane. # operations that will return 403 under AAD Data Contributor role. This test uses key_container @@ -125,6 +147,190 @@ def test_round_trip_emoji_document_via_query(self): self.assertEqual(len(results), 1) self.assertEqual(results[0]['multibyte_content'], emoji_payload) + def test_compact_utf8_item_writes_through_full_sdk_stack(self): + """Exercise every item-write operation covered by the client option.""" + captured = {} + + def capture_body(request): + captured['body'] = request.http_request.body + + def assert_compact_body(expected_content): + # Compact bodies are handed to the transport as UTF-8 bytes, not str, + # so that no transport layer can re-encode them (urllib3 1.x would + # otherwise Latin-1 encode a str body via http.client). + body = captured['body'] + self.assertIsInstance(body, bytes) + decoded = body.decode('utf-8') + self.assertIn(expected_content, decoded) + self.assertNotIn('\\u65e5', decoded) + + with test_config.TestConfig.create_data_client( + enable_compact_utf8_item_writes=True + ) as client: + container = client.get_database_client( + test_config.TestConfig.TEST_DATABASE_ID + ).get_container_client( + test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID + ) + doc_id = 'utf8-writes-' + str(uuid.uuid4()) + + # Keep the non-ASCII partition key: these operations must send an + # escaped partition-key header alongside a raw UTF-8 item body. + created = container.create_item({ + 'id': doc_id, + 'pk': '日本', + 'content': 'create café 日本 🎉', + }, raw_request_hook=capture_body) + self.assertEqual(created['content'], 'create café 日本 🎉') + + # Round-tripping alone cannot detect a regression to escaped + # output, since escaped JSON round-trips identically. Assert on + # the bytes actually put on the wire. + assert_compact_body(created['content']) + + created['content'] = 'upsert مرحبا 日本 🚀' # cspell:disable-line + upserted = container.upsert_item( + created, + raw_request_hook=capture_body, + ) + self.assertEqual(upserted['content'], created['content']) + assert_compact_body(upserted['content']) + + upserted['content'] = 'replace नमस्ते 日本 🌍' # cspell:disable-line + replaced = container.replace_item( + doc_id, + upserted, + raw_request_hook=capture_body, + ) + self.assertEqual(replaced['content'], upserted['content']) + assert_compact_body(replaced['content']) + + patched = container.patch_item( + doc_id, + partition_key='日本', + patch_operations=[ + { + 'op': 'set', + 'path': '/content', + 'value': 'patch שלום 日本 🎊', # cspell:disable-line + }, + ], + raw_request_hook=capture_body, + ) + self.assertEqual(patched['content'], 'patch שלום 日本 🎊') # cspell:disable-line + assert_compact_body(patched['content']) + + # A patch filter predicate travels in the same request body as the + # patch operations, so a non-ASCII predicate is affected by the + # option too and must be sent compact and accepted by the service. + conditional = container.patch_item( + doc_id, + partition_key='日本', + patch_operations=[ + { + 'op': 'set', + 'path': '/content', + 'value': 'conditional patch 日本 ✅', + }, + ], + filter_predicate="FROM c WHERE c.pk = '日本'", + raw_request_hook=capture_body, + ) + self.assertEqual(conditional['content'], 'conditional patch 日本 ✅') + assert_compact_body("FROM c WHERE c.pk = '日本'") + + batch_id = 'utf8-batch-日本-' + str(uuid.uuid4()) + batch_result = container.execute_item_batch( + [( + 'create', + ({'id': batch_id, 'pk': '日本', 'content': 'batch ไทย 日本 🎉'},), # cspell:disable-line + )], + partition_key='日本', + raw_request_hook=capture_body, + ) + self.assertEqual(batch_result[0]['statusCode'], 201) + assert_compact_body('batch ไทย 日本 🎉') # cspell:disable-line + + read_batch_result = container.execute_item_batch( + [('read', (batch_id,))], + partition_key='日本', + raw_request_hook=capture_body, + ) + self.assertEqual(read_batch_result[0]['statusCode'], 200) + self.assertEqual(read_batch_result[0]['resourceBody']['content'], 'batch ไทย 日本 🎉') # cspell:disable-line + self.assertNotIn('日本', captured['body']) + self.assertIn('\\u65e5', captured['body']) + + # A delete-only batch is the other body-free shape: like read, the + # operation carries just an id, so the item-write option must leave + # it escaped. The service must still accept that escaped form. + delete_batch_result = container.execute_item_batch( + [('delete', (batch_id,))], + partition_key='日本', + raw_request_hook=capture_body, + ) + self.assertEqual(delete_batch_result[0]['statusCode'], 204) + self.assertNotIn('日本', captured['body']) + self.assertIn('\\u65e5', captured['body']) + + queried = list(container.query_items( + query="SELECT * FROM c WHERE c.id = @id", + parameters=[{'name': '@id', 'value': doc_id}], + partition_key='日本', + )) + self.assertEqual(len(queried), 1) + # The conditional patch is the last write to touch this item. + self.assertEqual(queried[0]['content'], conditional['content']) + + def test_default_ascii_escaping_rejects_large_cjk_item(self): + """Verify the backend rejects the escaped request body because it exceeds 2 MiB.""" + document = self._large_cjk_document('utf8-large-default-') + escaped_body = json.dumps(document, separators=(",", ":")).encode("utf-8") + self.assertGreater(len(escaped_body), 2 * 1024 * 1024) + + with test_config.TestConfig.create_data_client() as client: + container = client.get_database_client( + test_config.TestConfig.TEST_DATABASE_ID + ).get_container_client( + test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID + ) + + with self.assertRaises(exceptions.CosmosHttpResponseError) as context: + container.create_item(document) + + self.assertEqual(context.exception.status_code, 413) + + def test_compact_utf8_large_cjk_item_stays_under_wire_limit(self): + """Write the same oversized-when-escaped item using compact UTF-8.""" + document = self._large_cjk_document('utf8-large-compact-') + compact_body = json.dumps( + document, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + self.assertLess(len(compact_body), 2 * 1024 * 1024) + + with test_config.TestConfig.create_data_client( + enable_compact_utf8_item_writes=True + ) as client: + container = client.get_database_client( + test_config.TestConfig.TEST_DATABASE_ID + ).get_container_client( + test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID + ) + created = container.create_item(document) + try: + self.assertEqual(created['content'], document['content']) + finally: + # This item is roughly 1.2 MiB and lives in a shared container that is reused by + # the emulator, key-live, and AAD-live lanes. Delete it in a finally block so + # repeated CI runs cannot accumulate large indexed documents, and so an assertion + # failure above does not leak one either. + try: + container.delete_item(document['id'], partition_key=document['pk']) + except exceptions.CosmosResourceNotFoundError: + pass + if __name__ == "__main__": unittest.main() diff --git a/sdk/cosmos/azure-cosmos/tests/test_encoding_async.py b/sdk/cosmos/azure-cosmos/tests/test_encoding_async.py index ba6df138af3b..831d1047d127 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_encoding_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_encoding_async.py @@ -7,22 +7,38 @@ using the async client. The stored-procedure control-plane check remains in the sync file because it relies on key-auth script operations. """ +import json import unittest import uuid import pytest import test_config +from azure.cosmos import exceptions @pytest.mark.cosmosEmulator +@pytest.mark.cosmosLong +@pytest.mark.cosmosAADLong class TestEncodingAsync(unittest.IsolatedAsyncioTestCase): - """Async round-trips for non-ASCII document content.""" + """Async round-trips for non-ASCII document content. + + Marked for the emulator lane and both live lanes so the compact UTF-8 + item-write coverage runs against a real account as well. + """ host = test_config.TestConfig.host masterKey = test_config.TestConfig.masterKey connectionPolicy = test_config.TestConfig.connectionPolicy + @staticmethod + def _large_cjk_document(id_prefix): + return { + 'id': id_prefix + str(uuid.uuid4()), + 'pk': 'pk', + 'content': '日' * 400000, + } + @classmethod def setUpClass(cls): if (cls.masterKey == '[YOUR_KEY_HERE]' @@ -122,6 +138,186 @@ async def test_round_trip_emoji_document_via_query_async(self): self.assertEqual(len(results), 1) self.assertEqual(results[0]['multibyte_content'], emoji_payload) + async def test_compact_utf8_item_writes_through_full_sdk_stack_async(self): + """Exercise every item-write operation covered by the client option.""" + captured = {} + + def capture_body(request): + captured['body'] = request.http_request.body + + def assert_compact_body(expected_content): + # Compact bodies are handed to the transport as UTF-8 bytes, not str, + # so that no transport layer can re-encode them (urllib3 1.x would + # otherwise Latin-1 encode a str body via http.client). + body = captured['body'] + self.assertIsInstance(body, bytes) + decoded = body.decode('utf-8') + self.assertIn(expected_content, decoded) + self.assertNotIn('\\u65e5', decoded) + + async with test_config.TestConfig.create_data_client_async( + enable_compact_utf8_item_writes=True + ) as client: + container = client.get_database_client( + test_config.TestConfig.TEST_DATABASE_ID + ).get_container_client( + test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID + ) + doc_id = 'utf8-writes-async-' + str(uuid.uuid4()) + + # Keep the non-ASCII partition key: these operations must send an + # escaped partition-key header alongside a raw UTF-8 item body. + created = await container.create_item({ + 'id': doc_id, + 'pk': '日本', + 'content': 'create café 日本 🎉', + }, raw_request_hook=capture_body) + self.assertEqual(created['content'], 'create café 日本 🎉') + + # Round-tripping alone cannot detect a regression to escaped + # output, since escaped JSON round-trips identically. Assert on + # the bytes actually put on the wire. + assert_compact_body(created['content']) + + created['content'] = 'upsert مرحبا 日本 🚀' # cspell:disable-line + upserted = await container.upsert_item( + created, + raw_request_hook=capture_body, + ) + self.assertEqual(upserted['content'], created['content']) + assert_compact_body(upserted['content']) + + upserted['content'] = 'replace नमस्ते 日本 🌍' # cspell:disable-line + replaced = await container.replace_item( + doc_id, + upserted, + raw_request_hook=capture_body, + ) + self.assertEqual(replaced['content'], upserted['content']) + assert_compact_body(replaced['content']) + + patched = await container.patch_item( + doc_id, + partition_key='日本', + patch_operations=[ + { + 'op': 'set', + 'path': '/content', + 'value': 'patch שלום 日本 🎊', # cspell:disable-line + }, + ], + raw_request_hook=capture_body, + ) + self.assertEqual(patched['content'], 'patch שלום 日本 🎊') # cspell:disable-line + assert_compact_body(patched['content']) + + # A patch filter predicate travels in the same request body as the + # patch operations, so a non-ASCII predicate is affected by the + # option too and must be sent compact and accepted by the service. + conditional = await container.patch_item( + doc_id, + partition_key='日本', + patch_operations=[ + { + 'op': 'set', + 'path': '/content', + 'value': 'conditional patch 日本 ✅', + }, + ], + filter_predicate="FROM c WHERE c.pk = '日本'", + raw_request_hook=capture_body, + ) + self.assertEqual(conditional['content'], 'conditional patch 日本 ✅') + assert_compact_body("FROM c WHERE c.pk = '日本'") + + batch_id = 'utf8-batch-async-日本-' + str(uuid.uuid4()) + batch_result = await container.execute_item_batch( + [( + 'create', + ({'id': batch_id, 'pk': '日本', 'content': 'batch ไทย 日本 🎉'},), # cspell:disable-line + )], + partition_key='日本', + raw_request_hook=capture_body, + ) + self.assertEqual(batch_result[0]['statusCode'], 201) + assert_compact_body('batch ไทย 日本 🎉') # cspell:disable-line + + read_batch_result = await container.execute_item_batch( + [('read', (batch_id,))], + partition_key='日本', + raw_request_hook=capture_body, + ) + self.assertEqual(read_batch_result[0]['statusCode'], 200) + self.assertEqual(read_batch_result[0]['resourceBody']['content'], 'batch ไทย 日本 🎉') # cspell:disable-line + self.assertNotIn('日本', captured['body']) + self.assertIn('\\u65e5', captured['body']) + + # A delete-only batch is the other body-free shape: like read, the + # operation carries just an id, so the item-write option must leave + # it escaped. The service must still accept that escaped form. + delete_batch_result = await container.execute_item_batch( + [('delete', (batch_id,))], + partition_key='日本', + raw_request_hook=capture_body, + ) + self.assertEqual(delete_batch_result[0]['statusCode'], 204) + self.assertNotIn('日本', captured['body']) + self.assertIn('\\u65e5', captured['body']) + + queried = [] + async for item in container.query_items( + query="SELECT * FROM c WHERE c.id = @id", + parameters=[{'name': '@id', 'value': doc_id}], + partition_key='日本'): + queried.append(item) + self.assertEqual(len(queried), 1) + # The conditional patch is the last write to touch this item. + self.assertEqual(queried[0]['content'], conditional['content']) + + async def test_default_ascii_escaping_rejects_large_cjk_item_async(self): + """Verify the backend rejects the escaped async request body because it exceeds 2 MiB.""" + document = self._large_cjk_document('utf8-large-default-async-') + escaped_body = json.dumps(document, separators=(",", ":")).encode("utf-8") + self.assertGreater(len(escaped_body), 2 * 1024 * 1024) + + with self.assertRaises(exceptions.CosmosHttpResponseError) as context: + await self.created_container.create_item(document) + + self.assertEqual(context.exception.status_code, 413) + + async def test_compact_utf8_large_cjk_item_stays_under_wire_limit_async(self): + """Write an item that exceeds 2 MiB when escaped but not as compact UTF-8.""" + document = self._large_cjk_document('utf8-large-compact-async-') + escaped_body = json.dumps(document, separators=(",", ":")).encode("utf-8") + compact_body = json.dumps( + document, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + self.assertGreater(len(escaped_body), 2 * 1024 * 1024) + self.assertLess(len(compact_body), 2 * 1024 * 1024) + + async with test_config.TestConfig.create_data_client_async( + enable_compact_utf8_item_writes=True + ) as client: + container = client.get_database_client( + test_config.TestConfig.TEST_DATABASE_ID + ).get_container_client( + test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID + ) + created = await container.create_item(document) + try: + self.assertEqual(created['content'], document['content']) + finally: + # This item is roughly 1.2 MiB and lives in a shared container that is reused by + # the emulator, key-live, and AAD-live lanes. Delete it in a finally block so + # repeated CI runs cannot accumulate large indexed documents, and so an assertion + # failure above does not leak one either. + try: + await container.delete_item(document['id'], partition_key=document['pk']) + except exceptions.CosmosResourceNotFoundError: + pass + if __name__ == "__main__": unittest.main() diff --git a/sdk/cosmos/azure-cosmos/tests/test_hybrid_search_aggregator.py b/sdk/cosmos/azure-cosmos/tests/test_hybrid_search_aggregator.py new file mode 100644 index 000000000000..ec7a11fd2713 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_hybrid_search_aggregator.py @@ -0,0 +1,44 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +"""Deterministic unit tests for client-side hybrid-search ranking.""" + +import pytest + +from azure.cosmos._execution_context.hybrid_search_aggregator import ( + _compute_ranks, + _compute_rrf_scores, +) + + +@pytest.mark.cosmosEmulator +def test_mixed_component_rrf_scores_have_exact_expected_order(): + """Pin dense ranking, tied component scores, weighted RRF math, and the + final descending order without depending on service-generated scores.""" + # Each tuple contains (component score, original result index). The first + # component is already sorted from highest to lowest, while the second is + # sorted from lowest to highest, as the query pipeline does before ranking. + component_scores = [ + [(0.9, 0), (0.8, 1), (0.8, 2), (0.5, 3)], + [(0.1, 2), (0.2, 0), (0.3, 3), (0.4, 1)], + ] + results = [{"id": item_id} for item_id in ("a", "b", "c", "d")] + + ranks = _compute_ranks(component_scores) + _compute_rrf_scores(ranks, [1, 2], results) + + assert ranks == [ + [1, 2, 2, 3], + [2, 4, 1, 3], + ] + expected_scores = { + "a": 1 / 61 + 2 / 62, + "b": 1 / 62 + 2 / 64, + "c": 1 / 62 + 2 / 61, + "d": 1 / 63 + 2 / 63, + } + for result in results: + assert result["Score"] == pytest.approx(expected_scores[result["id"]]) + + results.sort(key=lambda result: result["Score"], reverse=True) + assert [result["id"] for result in results] == ["c", "a", "d", "b"] diff --git a/sdk/cosmos/azure-cosmos/tests/test_item_body_serialization.py b/sdk/cosmos/azure-cosmos/tests/test_item_body_serialization.py new file mode 100644 index 000000000000..f4b4f1be6f63 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_item_body_serialization.py @@ -0,0 +1,974 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +"""Tests for opt-in compact UTF-8 serialization of item write bodies.""" +# cspell:ignore d83d de00 udfff + +import json +import unittest +from unittest import mock + +import pytest + +from azure.cosmos import ( + _base, + _global_endpoint_manager, + _synchronized_request, + cosmos_client, + documents, + http_constants, +) + +from azure.cosmos.aio import _asynchronous_request, _cosmos_client +from azure.cosmos.documents import _OperationType +from azure.cosmos.http_constants import HttpHeaders + + +class _DummyRequestParams: + """Stand-in for RequestObject carrying only the fields the body + serialization decision reads: resource type and operation type.""" + + def __init__( + self, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Create, + ): + self.availability_strategy = None + self.is_hedging_request = False + self.resource_type = resource_type + self.operation_type = operation_type + self.retry_write = 0 + + +class _DummyGlobalEndpointManager: + """Endpoint manager stub that reports hedging as disabled, so requests + take the plain (non-hedged) path through to the mocked executor.""" + + @staticmethod + def is_per_partition_automatic_failover_enabled(): + return False + + +class _DummyRequest: + """Minimal HttpRequest stand-in capturing just the body and headers.""" + + def __init__(self): + self.headers = {} + self.data = None + + +class _DummyClient: + """Client stub exposing only the opt-in flag the serializer checks.""" + + def __init__(self, enable_compact_utf8_item_writes): + self._enable_compact_utf8_item_writes = enable_compact_utf8_item_writes + + +class _DummyHeaderClient: + """Client stub with the attributes GetHeaders needs to build headers + without a real connection.""" + + UseMultipleWriteLocations = False + master_key = None + resource_tokens = None + client_id = None + + class connection_policy: + ResponsePayloadOnWriteDisabled = False + + +class _EncodingCountingStr(str): + """A str that counts how many times .encode() is called on it. + + Used to prove the body is encoded to UTF-8 exactly once per request: + the serializer keeps the bytes it already produced and reuses their + length for Content-Length, rather than encoding the body a second time. + """ + + def __new__(cls, value): + instance = super().__new__(cls, value) + instance.encode_calls = 0 + return instance + + def encode(self, encoding="utf-8", errors="strict"): + self.encode_calls += 1 + return super().encode(encoding, errors) + + +def _compact_text(body): + """Decode a compact request body for text assertions. + + Compact UTF-8 bodies are handed to the transport as ``bytes`` (never + ``str``) so no transport layer can re-encode them: urllib3 1.x routes + ``str`` bodies through ``http.client``, which Latin-1 encodes them. Asserting + the type here means a regression back to ``str`` fails loudly rather than + silently reintroducing that bug. + """ + assert isinstance(body, bytes), f"compact body must be bytes, got {type(body).__name__}" + return body.decode("utf-8") + + +def _capture_sync_body( + request_data, + *, + enable_compact_utf8_item_writes, + resource_type, + operation_type, +): + """Run one request through the sync path with the retry executor mocked + out, and return the body and Content-Length that would have been sent.""" + request = _DummyRequest() + captured = {} + + def _fake_execute(*args, **kwargs): + request_arg = args[6] + captured["body"] = request_arg.data + captured["content_length"] = request_arg.headers.get(HttpHeaders.ContentLength) + return {}, {} + + with mock.patch.object( + _synchronized_request._retry_utility, + "Execute", + side_effect=_fake_execute, + ): + _synchronized_request.SynchronizedRequest( + client=_DummyClient(enable_compact_utf8_item_writes), + request_params=_DummyRequestParams(resource_type, operation_type), + global_endpoint_manager=_DummyGlobalEndpointManager(), + connection_policy=object(), + pipeline_client=object(), + request=request, + request_data=request_data, + ) + return captured + + +async def _capture_async_body( + request_data, + *, + enable_compact_utf8_item_writes, + resource_type, + operation_type, +): + """Async twin of _capture_sync_body. Same capture, async pipeline.""" + request = _DummyRequest() + captured = {} + + async def _fake_execute(*args, **kwargs): + request_arg = args[6] + captured["body"] = request_arg.data + captured["content_length"] = request_arg.headers.get(HttpHeaders.ContentLength) + return {}, {} + + with mock.patch.object( + _asynchronous_request._retry_utility_async, + "ExecuteAsync", + side_effect=_fake_execute, + ): + await _asynchronous_request.AsynchronousRequest( + client=_DummyClient(enable_compact_utf8_item_writes), + request_params=_DummyRequestParams(resource_type, operation_type), + global_endpoint_manager=_DummyGlobalEndpointManager(), + connection_policy=object(), + pipeline_client=object(), + request=request, + request_data=request_data, + ) + return captured + + +def _capture_sync_query_builder(query): + """Build a real sync QueryFeed request and capture its serialized body.""" + # Bypass __init__ because it builds the full connection infrastructure and + # performs account discovery. These assignments are the minimal QueryFeed + # fixture; add an attribute here if QueryFeed gains another dependency. + connection = object.__new__(cosmos_client.CosmosClientConnection) + connection.default_headers = {} + connection.last_response_headers = {} + connection._query_compatibility_mode = ( + cosmos_client.CosmosClientConnection._QueryCompatibilityMode.Default + ) + connection.availability_strategy = None + connection.availability_strategy_executor = None + connection._global_endpoint_manager = _DummyGlobalEndpointManager() + connection.connection_policy = object() + connection.pipeline_client = mock.Mock() + connection.pipeline_client.post.return_value = _DummyRequest() + connection._enable_compact_utf8_item_writes = True + connection._UpdateSessionIfRequired = mock.Mock() + captured = {} + + def _fake_execute( # pylint: disable=too-many-arguments,too-many-positional-arguments + _client, + _global_endpoint_manager, + _request_function, + request_params, + _connection_policy, + _pipeline_client, + request, + **_kwargs, + ): + captured["resource_type"] = request_params.resource_type + captured["operation_type"] = request_params.operation_type + captured["body"] = request.data + captured["content_length"] = request.headers.get(HttpHeaders.ContentLength) + return {"Documents": []}, {} + + with ( + mock.patch.object(_base, "GetHeaders", return_value={}), + mock.patch.object(_base, "set_session_token_header"), + mock.patch.object( + _synchronized_request._retry_utility, + "Execute", + side_effect=_fake_execute, + ), + ): + connection.QueryFeed( + "dbs/db/colls/container/docs", + "container-rid", + query, + {}, + ) + return captured + + +async def _capture_async_query_builder(query): + """Build a real async QueryFeed request and capture its serialized body.""" + # Bypass __init__ because it builds the full async connection infrastructure. + # These assignments are the minimal QueryFeed fixture; add an attribute here + # if QueryFeed gains another dependency. + connection = object.__new__(_cosmos_client.CosmosClientConnection) + connection.default_headers = {} + connection.last_response_headers = {} + connection._query_compatibility_mode = ( + _cosmos_client.CosmosClientConnection._QueryCompatibilityMode.Default + ) + connection.availability_strategy = None + connection.availability_strategy_max_concurrency = None + connection._global_endpoint_manager = _DummyGlobalEndpointManager() + connection.connection_policy = object() + connection.pipeline_client = mock.Mock() + connection.pipeline_client.post.return_value = _DummyRequest() + connection._enable_compact_utf8_item_writes = True + connection._UpdateSessionIfRequired = mock.Mock() + captured = {} + + async def _fake_execute( # pylint: disable=too-many-arguments,too-many-positional-arguments + _client, + _global_endpoint_manager, + _request_function, + request_params, + _connection_policy, + _pipeline_client, + request, + **_kwargs, + ): + captured["resource_type"] = request_params.resource_type + captured["operation_type"] = request_params.operation_type + captured["body"] = request.data + captured["content_length"] = request.headers.get(HttpHeaders.ContentLength) + return {"Documents": []}, {} + + with ( + mock.patch.object(_base, "GetHeaders", return_value={}), + mock.patch.object(_base, "set_session_token_header_async"), + mock.patch.object( + _asynchronous_request._retry_utility_async, + "ExecuteAsync", + side_effect=_fake_execute, + ), + ): + await connection.QueryFeed( + "dbs/db/colls/container/docs", + "container-rid", + query, + {}, + ) + return captured + + +# These tests need no emulator, but the Cosmos CI lane selects tests with +# "-m cosmosEmulator" (see eng/pipelines/templates/stages/cosmos-sdk-client.yml), +# so an unmarked test is silently deselected and would never run. +@pytest.mark.cosmosEmulator +class TestItemBodySerialization(unittest.TestCase): + """Sync path: how item write bodies are serialized on and off the option.""" + + def test_default_serialization_remains_ascii_escaped(self): + """With the option off (the default), bodies keep the historical + \\uXXXX-escaped form, so existing customers see no change on the wire.""" + data = {"text": "café 日本 🎉"} + + captured = _capture_sync_body( + data, + enable_compact_utf8_item_writes=False, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Create, + ) + + self.assertEqual(captured["body"], json.dumps(data, separators=(",", ":"))) + self.assertNotIn("日本", captured["body"]) + + def test_enabled_item_write_operations_use_compact_utf8(self): + """With the option on, every operation the feature is scoped to + (create, upsert, replace, patch) sends unescaped UTF-8, and + Content-Length matches the body's real byte count.""" + data = {"text": "café 日本 🎉"} + expected = json.dumps(data, separators=(",", ":"), ensure_ascii=False) + + for operation_type in ( + _OperationType.Create, + _OperationType.Upsert, + _OperationType.Replace, + _OperationType.Patch, + ): + with self.subTest(operation_type=operation_type): + captured = _capture_sync_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=operation_type, + ) + self.assertEqual(captured["body"], expected.encode("utf-8")) + self.assertEqual(captured["content_length"], len(expected.encode("utf-8"))) + + def test_batch_wire_shape_is_what_the_detector_keys_off(self): + """Pin the service's batch wire contract independently of the SDK. + + _batch_contains_item_body decides on the presence of a resourceBody + key. The formatter-driven tests below would still pass if the + formatter and the detector were renamed together, so hard-code the + shape the service actually expects here: write operations carry a + resourceBody, read and delete carry only an id.""" + item = {"id": "item-日本", "text": "ไทย 🎉"} + with_body = _base._format_batch_operations([ + ("create", (item,)), + ("upsert", (item,)), + ("replace", ("item-日本", item)), + ("patch", ("item-日本", [{"op": "add", "path": "/text", "value": "ไทย"}])), + ]) + without_body = _base._format_batch_operations([ + ("read", ("item-日本",)), + ("delete", ("item-日本",)), + ]) + + for operation in with_body: + with self.subTest(operation=operation["operationType"]): + self.assertIn("resourceBody", operation) + for operation in without_body: + with self.subTest(operation=operation["operationType"]): + self.assertNotIn("resourceBody", operation) + self.assertIn("id", operation) + + def test_body_free_batches_remain_ascii_escaped(self): + """Read and delete operations both carry only an id, so any batch made + up solely of them has no item body and must keep the escaped form.""" + body_free_batches = ( + [("read", ("item-日本",)), ("read", ("item-ไทย",))], + [("delete", ("item-日本",)), ("delete", ("item-ไทย",))], + [("read", ("item-日本",)), ("delete", ("item-ไทย",))], + [], + ) + + for batch in body_free_batches: + with self.subTest(batch=batch): + data = _base._format_batch_operations(batch) + captured = _capture_sync_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Batch, + ) + self.assertEqual(captured["body"], json.dumps(data, separators=(",", ":"))) + self.assertNotIn("日本", captured["body"]) + + def test_batches_built_by_the_sdk_are_detected_as_item_writes(self): + """Drive the real batch formatter rather than a hand-written payload, + so renaming the resourceBody key can never silently disable compact + UTF-8 while these tests still pass. Patch is included because its body + is nested one level deeper than the others.""" + item = {"id": "item-日本", "text": "ไทย 🎉"} + write_batches = ( + [("create", (item,))], + [("upsert", (item,))], + [("replace", ("item-日本", item))], + [("patch", ("item-日本", [{"op": "add", "path": "/text", "value": "ไทย"}]))], + [("read", ("item-日本",)), ("create", (item,))], + ) + + for batch in write_batches: + with self.subTest(batch=batch[0][0]): + data = _base._format_batch_operations(batch) + self.assertTrue(_synchronized_request._batch_contains_item_body(data)) + + captured = _capture_sync_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Batch, + ) + expected = json.dumps(data, separators=(",", ":"), ensure_ascii=False) + self.assertEqual(captured["body"], expected.encode("utf-8")) + self.assertEqual(captured["content_length"], len(expected.encode("utf-8"))) + + def test_compact_body_reuses_encoded_byte_length(self): + """The body is encoded to UTF-8 exactly once. Guards the memory fix: + the Content-Length header reuses the bytes already produced during + serialization instead of re-encoding the whole body a second time.""" + serialized = _EncodingCountingStr('{"text":"日本"}') + + with mock.patch.object(_synchronized_request.json, "dumps", return_value=serialized): + captured = _capture_sync_body( + {"text": "日本"}, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Create, + ) + + self.assertEqual(serialized.encode_calls, 1) + self.assertEqual(captured["content_length"], len('{"text":"日本"}'.encode("utf-8"))) + + def test_enabled_option_does_not_change_other_request_bodies(self): + """The option is scoped to item writes only. Control-plane bodies + (databases, containers, users, permissions, sprocs, triggers, UDFs) + and document queries all stay ASCII-escaped even when it is on.""" + data = {"text": "日本"} + expected = json.dumps(data, separators=(",", ":")) + unaffected_requests = ( + (http_constants.ResourceType.Database, _OperationType.Create), + (http_constants.ResourceType.Collection, _OperationType.Replace), + (http_constants.ResourceType.User, _OperationType.Upsert), + (http_constants.ResourceType.Permission, _OperationType.Create), + (http_constants.ResourceType.StoredProcedure, _OperationType.Create), + (http_constants.ResourceType.StoredProcedure, _OperationType.ExecuteJavaScript), + (http_constants.ResourceType.Trigger, _OperationType.Create), + (http_constants.ResourceType.UserDefinedFunction, _OperationType.Create), + (http_constants.ResourceType.Document, _OperationType.SqlQuery), + ) + + for resource_type, operation_type in unaffected_requests: + with self.subTest(resource_type=resource_type, operation_type=operation_type): + captured = _capture_sync_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=resource_type, + operation_type=operation_type, + ) + self.assertEqual(captured["body"], expected) + + def test_query_builder_keeps_non_ascii_parameters_escaped(self): + """The real QueryFeed builder uses SqlQuery metadata and keeps query + parameters on the historical escaped-string serialization path.""" + data = { + "query": "SELECT * FROM c WHERE c.name = @name", + "parameters": [{"name": "@name", "value": "日本"}], + } + expected = json.dumps(data, separators=(",", ":")) + + captured = _capture_sync_query_builder(data) + + self.assertEqual(captured["resource_type"], http_constants.ResourceType.Document) + self.assertEqual(captured["operation_type"], _OperationType.SqlQuery) + self.assertEqual(captured["body"], expected) + self.assertNotIn("日本", captured["body"]) + self.assertEqual(captured["content_length"], len(expected.encode("utf-8"))) + + def test_pre_serialized_string_is_unchanged(self): + """A body the caller already serialized to a str is passed through + untouched. The SDK never re-parses or rewrites it.""" + data = '{"text":"\\u65e5\\u672c"}' + + captured = _capture_sync_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Create, + ) + + self.assertEqual(captured["body"], data) + + def test_surrogate_pairs_are_compact_and_lone_surrogates_are_escaped(self): + """Adjacent surrogate pairs become compact scalars while unpaired + surrogates remain valid JSON escapes.""" + data = { + "paired": "\ud83d\ude00", + "lone_high": "\ud800", + "lone_low": "\udfff", + "mixed": "before\ud800\ud83d\ude00\udfff" "after", + "reverse_lone": "\udfff\ud800", + "literal_escape": "\\ud83d\\ude00", + "key\ud800": "value", + "key\ud83d\ude00": "paired key", + "emoji": "🎉", + } + expected = { + "paired": "😀", + "lone_high": "\ud800", + "lone_low": "\udfff", + "mixed": "before\ud800😀\udfff" "after", + "reverse_lone": "\udfff\ud800", + "literal_escape": "\\ud83d\\ude00", + "key\ud800": "value", + "key😀": "paired key", + "emoji": "🎉", + } + expected_body = json.dumps( + expected, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8", "backslashreplace") + + captured = _capture_sync_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Create, + ) + body = _compact_text(captured["body"]) + + self.assertIn("\\ud800", body) + self.assertIn("\\udfff", body) + self.assertIn("😀", body) + self.assertIn("🎉", body) + self.assertEqual(captured["body"], expected_body) + self.assertEqual(json.loads(body), expected) + self.assertEqual(captured["content_length"], len(captured["body"])) + + def test_json_required_escapes_are_preserved(self): + """Turning off ASCII escaping must not turn off JSON escaping. Quotes, + backslashes, newlines, tabs, and NUL stay escaped; U+2028/U+2029 are + emitted literally, which is valid JSON.""" + data = {"text": "\"\\\n\t\x00", "line_separators": "\u2028\u2029"} + + captured = _capture_sync_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Create, + ) + body = _compact_text(captured["body"]) + + self.assertIn('\\"', body) + self.assertIn("\\\\", body) + self.assertIn("\\n", body) + self.assertIn("\\t", body) + self.assertIn("\\u0000", body) + self.assertIn("\u2028\u2029", body) + self.assertEqual(json.loads(body), data) + + def test_large_cjk_item_stays_below_two_mib(self): + """The point of the feature: a CJK document that exceeds the 2 MB item + limit when \\uXXXX-escaped fits under it as compact UTF-8.""" + data = {"text": "日" * 400000} + + escaped = json.dumps(data, separators=(",", ":")).encode("utf-8") + captured = _capture_sync_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Create, + ) + compact = captured["body"] + self.assertIsInstance(compact, bytes) + + self.assertGreater(len(escaped), 2 * 1024 * 1024) + self.assertLess(len(compact), 2 * 1024 * 1024) + self.assertEqual(captured["content_length"], len(compact)) + + def test_large_surrogate_pair_item_stays_below_two_mib(self): + """UTF-16-derived supplementary characters are compacted rather than + remaining as 12-byte surrogate-pair escapes.""" + data = {"text": "\ud83d\ude00" * 175000} + escaped = json.dumps(data, separators=(",", ":")).encode("utf-8") + expected = json.dumps( + {"text": "😀" * 175000}, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + + captured = _capture_sync_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Create, + ) + + self.assertGreater(len(escaped), 2 * 1024 * 1024) + self.assertLess(len(captured["body"]), 2 * 1024 * 1024) + self.assertEqual(captured["body"], expected) + self.assertEqual(captured["content_length"], len(expected)) + + def test_partition_key_header_remains_ascii_escaped(self): + """The partition key header is unaffected by the option and stays + ASCII-escaped, since headers are not UTF-8 safe in transit.""" + headers = _base.GetHeaders( + _DummyHeaderClient(), + {}, + "post", + "dbs/db/colls/container/docs", + "item", + http_constants.ResourceType.Document, + _OperationType.Create, + {"partitionKey": "日本"}, + ) + + self.assertEqual(headers[HttpHeaders.PartitionKey], '["\\u65e5\\u672c"]') + + +@pytest.mark.cosmosEmulator +class TestItemBodySerializationAsync(unittest.IsolatedAsyncioTestCase): + """Async path: mirrors the sync coverage so both stacks stay in step.""" + + async def test_default_item_write_remains_ascii_escaped(self): + """Async default is escaped, same as sync, with a byte-accurate + Content-Length.""" + data = {"text": "café 日本 🎉"} + expected = json.dumps(data, separators=(",", ":")) + + captured = await _capture_async_body( + data, + enable_compact_utf8_item_writes=False, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Create, + ) + + self.assertEqual(captured["body"], expected) + self.assertEqual(captured["content_length"], len(expected.encode("utf-8"))) + + async def test_enabled_item_write_uses_compact_utf8(self): + """Async opt-in sends compact UTF-8 with a byte-accurate + Content-Length.""" + data = {"text": "café 日本 🎉"} + expected = json.dumps(data, separators=(",", ":"), ensure_ascii=False) + + captured = await _capture_async_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Upsert, + ) + + self.assertEqual(captured["body"], expected.encode("utf-8")) + self.assertEqual(captured["content_length"], len(expected.encode("utf-8"))) + + async def test_write_containing_batch_uses_compact_utf8(self): + """Async mixed batches use compact UTF-8 when one operation contains + an item body. The read operation's id is compact too, since the whole + batch body is serialized in one pass.""" + data = _base._format_batch_operations([ + ("read", ("existing-日本",)), + ("create", ({"id": "new-日本", "text": "ไทย 🎉"},)), + ]) + expected = json.dumps(data, separators=(",", ":"), ensure_ascii=False) + + captured = await _capture_async_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Batch, + ) + + self.assertEqual(captured["body"], expected.encode("utf-8")) + self.assertIn("existing-日本", _compact_text(captured["body"])) + self.assertEqual(captured["content_length"], len(expected.encode("utf-8"))) + + async def test_body_free_batches_remain_ascii_escaped(self): + """Async twin: read-only, delete-only, and mixed read/delete batches + all lack an item body and stay escaped.""" + body_free_batches = ( + [("read", ("item-日本",)), ("read", ("item-ไทย",))], + [("delete", ("item-日本",)), ("delete", ("item-ไทย",))], + [("read", ("item-日本",)), ("delete", ("item-ไทย",))], + [], + ) + + for batch in body_free_batches: + with self.subTest(batch=batch): + data = _base._format_batch_operations(batch) + captured = await _capture_async_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Batch, + ) + self.assertEqual(captured["body"], json.dumps(data, separators=(",", ":"))) + self.assertNotIn("日本", captured["body"]) + + async def test_compact_body_reuses_encoded_byte_length(self): + """Async twin of the single-encode check, so the memory fix is + pinned on both stacks.""" + serialized = _EncodingCountingStr('{"text":"日本"}') + + with mock.patch.object(_synchronized_request.json, "dumps", return_value=serialized): + captured = await _capture_async_body( + {"text": "日本"}, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Create, + ) + + self.assertEqual(serialized.encode_calls, 1) + self.assertEqual(captured["content_length"], len('{"text":"日本"}'.encode("utf-8"))) + + async def test_surrogate_pairs_are_compact_and_lone_surrogates_are_escaped(self): + """Async twin of the mixed paired and unpaired surrogate check.""" + data = { + "paired": "\ud83d\ude00", + "lone_high": "\ud800", + "lone_low": "\udfff", + "mixed": "before\ud800\ud83d\ude00\udfff" "after", + "reverse_lone": "\udfff\ud800", + "literal_escape": "\\ud83d\\ude00", + "key\ud800": "value", + "key\ud83d\ude00": "paired key", + "emoji": "🎉", + } + expected = { + "paired": "😀", + "lone_high": "\ud800", + "lone_low": "\udfff", + "mixed": "before\ud800😀\udfff" "after", + "reverse_lone": "\udfff\ud800", + "literal_escape": "\\ud83d\\ude00", + "key\ud800": "value", + "key😀": "paired key", + "emoji": "🎉", + } + expected_body = json.dumps( + expected, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8", "backslashreplace") + + captured = await _capture_async_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.Create, + ) + body = _compact_text(captured["body"]) + + self.assertIn("\\ud800", body) + self.assertIn("\\udfff", body) + self.assertIn("😀", body) + self.assertIn("🎉", body) + self.assertEqual(captured["body"], expected_body) + self.assertEqual(json.loads(body), expected) + self.assertEqual(captured["content_length"], len(captured["body"])) + + async def test_query_body_remains_ascii_escaped(self): + """Query bodies stay escaped on the async path even with the option + on, since queries are not item writes.""" + data = {"query": "SELECT * FROM c WHERE c.name = @name", "parameters": [{"value": "日本"}]} + + captured = await _capture_async_body( + data, + enable_compact_utf8_item_writes=True, + resource_type=http_constants.ResourceType.Document, + operation_type=_OperationType.SqlQuery, + ) + + self.assertEqual(captured["body"], json.dumps(data, separators=(",", ":"))) + + async def test_query_builder_keeps_non_ascii_parameters_escaped(self): + """The real async QueryFeed builder uses SqlQuery metadata and keeps + query parameters on the historical escaped-string serialization path.""" + data = { + "query": "SELECT * FROM c WHERE c.name = @name", + "parameters": [{"name": "@name", "value": "日本"}], + } + expected = json.dumps(data, separators=(",", ":")) + + captured = await _capture_async_query_builder(data) + + self.assertEqual(captured["resource_type"], http_constants.ResourceType.Document) + self.assertEqual(captured["operation_type"], _OperationType.SqlQuery) + self.assertEqual(captured["body"], expected) + self.assertNotIn("日本", captured["body"]) + self.assertEqual(captured["content_length"], len(expected.encode("utf-8"))) + + +@pytest.mark.cosmosEmulator +class TestClientOptionWiring(unittest.IsolatedAsyncioTestCase): + """How the keyword travels from the client constructor to the serializer, + and how bad values are rejected.""" + + def test_sync_client_consumes_and_forwards_option(self): + """The keyword is forwarded to the client connection and is not leaked + into the connection policy kwargs, where it would be an unknown option.""" + with ( + mock.patch.object( + cosmos_client, + "_build_connection_policy", + return_value=object(), + ) as build_policy, + mock.patch.object(cosmos_client, "CosmosClientConnection") as connection, + ): + cosmos_client.CosmosClient( + "https://example.test", + "credential", + enable_compact_utf8_item_writes=True, + ) + + self.assertNotIn( + "enable_compact_utf8_item_writes", + build_policy.call_args.args[0], + ) + self.assertTrue(connection.call_args.kwargs["enable_compact_utf8_item_writes"]) + + def test_sync_client_rejects_non_boolean_option(self): + """A string such as "true" is rejected with a TypeError at client + construction rather than being silently coerced. Guards against a + config-sourced "false" evaluating truthy and enabling the feature.""" + with self.assertRaisesRegex( + TypeError, + "enable_compact_utf8_item_writes must be a bool", + ): + cosmos_client.CosmosClient( + "https://example.test", + "credential", + enable_compact_utf8_item_writes="true", + ) + + def test_sync_connection_rejects_non_boolean_option(self): + """The same strict bool check applies at the connection layer, for + ints, strings, and None.""" + for invalid_value in (0, "false", None): + with self.subTest(invalid_value=invalid_value): + with self.assertRaisesRegex( + TypeError, + "enable_compact_utf8_item_writes must be a bool", + ): + cosmos_client.CosmosClientConnection( + "https://example.test", + {"masterKey": "credential"}, + enable_compact_utf8_item_writes=invalid_value, + ) + + def test_sync_connection_string_forwards_option(self): + """The option survives the from_connection_string entry point.""" + with mock.patch.object(cosmos_client.CosmosClient, "__init__", return_value=None) as init: + cosmos_client.CosmosClient.from_connection_string( + "AccountEndpoint=https://example.test;AccountKey=credential;", + enable_compact_utf8_item_writes=True, + ) + + self.assertTrue(init.call_args.kwargs["enable_compact_utf8_item_writes"]) + + def test_sync_real_client_stores_and_honors_option(self): + """End to end on a real client object: the flag is stored and actually + flips the serializer's decision for an item write.""" + with ( + mock.patch.object( + _global_endpoint_manager._GlobalEndpointManager, + "_GetDatabaseAccount", + return_value=documents.DatabaseAccount(), + ), + mock.patch.object( + _global_endpoint_manager._GlobalEndpointManager, + "force_refresh_on_startup", + ), + ): + client = cosmos_client.CosmosClient( + "https://example.test", + "credential", + enable_compact_utf8_item_writes=True, + ) + + try: + self.assertTrue( + client.client_connection._enable_compact_utf8_item_writes + ) + self.assertFalse( + _synchronized_request._should_escape_non_ascii_in_request_body( + client.client_connection, + _DummyRequestParams(), + {"text": "日本"}, + ) + ) + finally: + client.close() + + def test_async_client_consumes_and_forwards_option(self): + """Async twin of the forwarding check.""" + with ( + mock.patch.object( + _cosmos_client, + "_build_connection_policy", + return_value=object(), + ) as build_policy, + mock.patch.object(_cosmos_client, "CosmosClientConnection") as connection, + ): + _cosmos_client.CosmosClient( + "https://example.test", + "credential", + enable_compact_utf8_item_writes=True, + ) + + self.assertNotIn( + "enable_compact_utf8_item_writes", + build_policy.call_args.args[0], + ) + self.assertTrue(connection.call_args.kwargs["enable_compact_utf8_item_writes"]) + + def test_async_client_rejects_non_boolean_option(self): + """Async twin of the string-rejection check.""" + with self.assertRaisesRegex( + TypeError, + "enable_compact_utf8_item_writes must be a bool", + ): + _cosmos_client.CosmosClient( + "https://example.test", + "credential", + enable_compact_utf8_item_writes="true", + ) + + def test_async_connection_rejects_non_boolean_option(self): + """Async twin of the connection-layer validation check.""" + for invalid_value in (0, "false", None): + with self.subTest(invalid_value=invalid_value): + with self.assertRaisesRegex( + TypeError, + "enable_compact_utf8_item_writes must be a bool", + ): + _cosmos_client.CosmosClientConnection( + "https://example.test", + {"masterKey": "credential"}, + enable_compact_utf8_item_writes=invalid_value, + ) + + def test_async_connection_string_forwards_option(self): + """Async twin of the from_connection_string check.""" + with mock.patch.object(_cosmos_client.CosmosClient, "__init__", return_value=None) as init: + _cosmos_client.CosmosClient.from_connection_string( + "AccountEndpoint=https://example.test;AccountKey=credential;", + enable_compact_utf8_item_writes=True, + ) + + self.assertTrue(init.call_args.kwargs["enable_compact_utf8_item_writes"]) + + async def test_async_real_client_stores_and_honors_option(self): + """Async twin of the end-to-end stored-and-honored check.""" + client = _cosmos_client.CosmosClient( + "https://example.test", + "credential", + enable_compact_utf8_item_writes=True, + ) + + try: + self.assertTrue( + client.client_connection._enable_compact_utf8_item_writes + ) + self.assertFalse( + _synchronized_request._should_escape_non_ascii_in_request_body( + client.client_connection, + _DummyRequestParams(), + {"text": "日本"}, + ) + ) + finally: + await client.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search.py b/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search.py index ba45c6930dc3..accea3cc698a 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search.py @@ -150,6 +150,11 @@ def test_hybrid_search_queries(self): results = self.test_container.query_items(query, enable_cross_partition_query=True) result_list = list(results) assert len(result_list) == 13 + # Captured so the OFFSET query below can be validated against the ranking this same + # container produced, instead of against a hard coded order. Equal RRF scores are broken by + # _rid, which is assigned by the service and differs between containers, so the absolute + # order is not stable across runs - but the offset window must always be a slice of it. + full_rrf_ranking = [res['index'] for res in result_list] for res in result_list: assert res['index'] in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85] @@ -169,8 +174,10 @@ def test_hybrid_search_queries(self): results = self.test_container.query_items(query, enable_cross_partition_query=True) result_list = list(results) assert len(result_list) == 8 - for res in result_list: - assert res['index'] in [24, 77, 76, 80, 25, 22, 2, 66, 57, 85] + # Same query as the ranking captured above, only with OFFSET/LIMIT applied. The SDK must + # return exactly that window of the ranking - this still fails if offset handling or the + # RRF ordering regresses, while tolerating service side score/_rid differences. + assert [res['index'] for res in result_list] == full_rrf_ranking[5:15] query = "SELECT TOP 10 c.index, c.title FROM c " \ "ORDER BY RANK RRF(FullTextScore(c.title, 'John'), FullTextScore(c.text, 'United States'))" @@ -248,6 +255,7 @@ def test_hybrid_search_weighted_reciprocal_rank_fusion(self): # If some scores rank the same the order of the results may change for result in result_list: assert result in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85] + equal_weights_ranking = result_list # Test case 2 query = """ @@ -261,6 +269,11 @@ def test_hybrid_search_weighted_reciprocal_rank_fusion(self): # If some scores rank the same the order of the results may change for result in result_list: assert result in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85] + # [10, 10] is a positive uniform scaling of [1, 1], so every RRF score is multiplied by the + # same constant and the ranking must be identical. This is an SDK invariant that holds + # regardless of the scores the service returns, so it catches weight handling regressions + # without depending on a specific backend score distribution. + assert result_list == equal_weights_ranking # Test case 3 query = """ @@ -289,14 +302,19 @@ def test_hybrid_search_weighted_reciprocal_rank_fusion(self): # Test case 5 item_vector = self.test_container.read_item('50', '1')['vector'] + query = "SELECT c.index, c.title FROM c " \ + "ORDER BY RANK RRF(FullTextScore(c.text, 'United States'), VectorDistance(c.vector, {})) " \ + "OFFSET 0 LIMIT 10".format(item_vector) + results = self.test_container.query_items(query, enable_cross_partition_query=True) + result_list_without_weights = [res['index'] for res in results] + query = "SELECT c.index, c.title FROM c " \ "ORDER BY RANK RRF(FullTextScore(c.text, 'United States'), VectorDistance(c.vector, {}), [1,1]) " \ "OFFSET 0 LIMIT 10".format(item_vector) results = self.test_container.query_items(query, enable_cross_partition_query=True) - result_list = list(results) - assert len(result_list) == 10 - result_list = [res['index'] for res in result_list] - assert result_list == [51, 54, 28, 70, 56, 24, 26, 61, 58, 68] + result_list_with_equal_weights = [res['index'] for res in results] + assert len(result_list_with_equal_weights) == 10 + assert result_list_with_equal_weights == result_list_without_weights def test_invalid_hybrid_search_queries_weighted_reciprocal_rank_fusion(self): try: @@ -598,4 +616,3 @@ def test_hybrid_search_parameterized_with_full_text_score_scope(self): if __name__ == "__main__": unittest.main() - diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search_async.py index b5f7307d8357..1c7815f097d5 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_hybrid_search_async.py @@ -145,6 +145,11 @@ async def test_hybrid_search_queries_async(self): results = self.test_container.query_items(query) result_list = [item async for item in results] assert len(result_list) == 13 + # Captured so the OFFSET query below can be validated against the ranking this same + # container produced, instead of against a hard coded order. Equal RRF scores are broken by + # _rid, which is assigned by the service and differs between containers, so the absolute + # order is not stable across runs - but the offset window must always be a slice of it. + full_rrf_ranking = [res['index'] for res in result_list] for res in result_list: assert res['index'] in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85] @@ -164,8 +169,10 @@ async def test_hybrid_search_queries_async(self): results = self.test_container.query_items(query) result_list = [item async for item in results] assert len(result_list) == 8 - for res in result_list: - assert res['index'] in [24, 77, 76, 80, 25, 22, 2, 66, 57, 85] + # Same query as the ranking captured above, only with OFFSET/LIMIT applied. The SDK must + # return exactly that window of the ranking - this still fails if offset handling or the + # RRF ordering regresses, while tolerating service side score/_rid differences. + assert [res['index'] for res in result_list] == full_rrf_ranking[5:15] query = "SELECT TOP 10 c.index, c.title FROM c " \ "ORDER BY RANK RRF(FullTextScore(c.title, 'John'), FullTextScore(c.text, 'United States'))" @@ -246,6 +253,7 @@ async def test_hybrid_search_weighted_reciprocal_rank_fusion_async(self): # If some scores rank the same the order of the results may change for result in result_list: assert result in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85] + equal_weights_ranking = result_list # Test case 2 query = """ @@ -259,6 +267,11 @@ async def test_hybrid_search_weighted_reciprocal_rank_fusion_async(self): # If some scores rank the same the order of the results may change for result in result_list: assert result in [61, 51, 49, 54, 75, 24, 77, 76, 80, 25, 22, 2, 66, 57, 85] + # [10, 10] is a positive uniform scaling of [1, 1], so every RRF score is multiplied by the + # same constant and the ranking must be identical. This is an SDK invariant that holds + # regardless of the scores the service returns, so it catches weight handling regressions + # without depending on a specific backend score distribution. + assert result_list == equal_weights_ranking # Test case 3 query = """ @@ -288,14 +301,19 @@ async def test_hybrid_search_weighted_reciprocal_rank_fusion_async(self): # Test case 5 read_item = await self.test_container.read_item('50', '1') item_vector = read_item['vector'] + query = "SELECT c.index, c.title FROM c " \ + "ORDER BY RANK RRF(FullTextScore(c.text, 'United States'), VectorDistance(c.vector, {})) " \ + "OFFSET 0 LIMIT 10".format(item_vector) + results = self.test_container.query_items(query) + result_list_without_weights = [res['index'] async for res in results] + query = "SELECT c.index, c.title FROM c " \ "ORDER BY RANK RRF(FullTextScore(c.text, 'United States'), VectorDistance(c.vector, {}), [1,1]) " \ "OFFSET 0 LIMIT 10".format(item_vector) results = self.test_container.query_items(query) - result_list = [res async for res in results] - assert len(result_list) == 10 - result_list = [res['index'] for res in result_list] - assert result_list == [51, 54, 28, 70, 56, 24, 26, 61, 58, 68] + result_list_with_equal_weights = [res['index'] async for res in results] + assert len(result_list_with_equal_weights) == 10 + assert result_list_with_equal_weights == result_list_without_weights async def test_invalid_hybrid_search_queries_weighted_reciprocal_rank_fusion_async(self): try: @@ -592,4 +610,3 @@ async def test_hybrid_search_parameterized_with_full_text_score_scope_async(self if __name__ == "__main__": unittest.main() - diff --git a/sdk/cosmos/azure-cosmos/tests/test_transport_body_encoding.py b/sdk/cosmos/azure-cosmos/tests/test_transport_body_encoding.py new file mode 100644 index 000000000000..db638041b6c1 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_transport_body_encoding.py @@ -0,0 +1,371 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +"""Transport-level checks for compact UTF-8 item write bodies. + +These tests send requests through the real ``azure-core`` sync and async +transports to a local HTTP server and assert on the bytes that server receives. +Assertions made on ``request.data`` before it reaches a transport cannot +observe re-encoding performed by the transport itself, so the checks here are +made against the wire bytes. + +Three properties are covered: + +* A compact body is handed to the transport as ``bytes``, not ``str``. The + supported dependency range still permits urllib3 1.x (``requests`` declares + ``urllib3<3,>=1.26``), where a ``str`` body is encoded as Latin-1 by + ``http.client`` (RFC 2616 3.7.1). Passing ``bytes`` leaves no re-encoding + step for any transport to get wrong. +* The server receives the exact UTF-8 encoding of the body, with a + ``Content-Length`` matching those bytes. Two payloads are used: CJK text, + which has no Latin-1 representation, and ``é``, which has one and so would be + sent as the single byte 0xE9 instead of the UTF-8 pair 0xC3 0xA9 without + raising anything. +* A PATCH request carries the explicit Cosmos media type, so the body type + cannot change the ``Content-Type`` a transport infers for it. +""" +# pylint: disable=invalid-name,missing-class-docstring,too-few-public-methods + +import json +import threading +import unittest +from http.server import BaseHTTPRequestHandler, HTTPServer + +import pytest + +# azure-core exposes the transports lazily through a module __getattr__, so +# static analysis cannot see them even though they are public API. +from azure.core.pipeline.transport import ( # pylint: disable=no-name-in-module + AioHttpTransport, + HttpRequest, + RequestsTransport, +) + +from azure.cosmos import _base, http_constants +from azure.cosmos._synchronized_request import _request_body_from_data +from azure.cosmos.documents import _OperationType + +# One payload that cannot be represented in Latin-1 at all, and one that can be +# but encodes to different bytes than UTF-8. Written as escapes so the intent +# survives any tooling that rewrites this file with the wrong encoding. +_CJK = {"x": "\u65e5\u672c"} +_LATIN1_REPRESENTABLE = {"x": "\u00e9"} + +# The exact bytes each payload must produce on the wire. +_CJK_UTF8 = b'{"x":"\xe6\x97\xa5\xe6\x9c\xac"}' +_LATIN1_REPRESENTABLE_UTF8 = b'{"x":"\xc3\xa9"}' +# What a str body would produce instead on urllib3 1.x: the same character +# written as one Latin-1 byte rather than the required UTF-8 pair. +_LATIN1_REPRESENTABLE_LATIN1 = b'{"x":"\xe9"}' + + +def _make_echo_handler(captured): + """Build a request handler that records one request into ``captured``. + + :param dict captured: Mapping the handler writes the body and headers into. + :returns: A handler class bound to the supplied mapping. + :rtype: type + """ + + class _EchoHandler(BaseHTTPRequestHandler): + """Records the exact request body bytes and Content-Length header.""" + + def do_POST(self): # pylint: disable=invalid-name + """Record the request body, then answer with an empty 200.""" + length = int(self.headers.get("Content-Length") or 0) + captured["body"] = self.rfile.read(length) + captured["content_length"] = self.headers.get("Content-Length") + captured["content_type"] = self.headers.get("Content-Type") + self.send_response(200) + self.send_header("Content-Length", "0") + self.end_headers() + + do_PATCH = do_POST # pylint: disable=invalid-name + + def log_message(self, *args): + """Silence the default stderr request logging.""" + + return _EchoHandler + + +def _latin1_encode_like_http_client(text): + """Reproduce what ``http.client._send_request`` does to a ``str`` body. + + The stdlib comment there cites RFC 2616 3.7.1 ("text default has a default + charset of iso-8859-1") and calls ``data.encode('latin-1')``. Spelled out + here rather than calling the private helper so this test does not depend on + a private stdlib API. + + :param str text: The body text a transport would encode. + :returns: The Latin-1 encoding of the text. + :rtype: bytes + """ + return text.encode("latin-1") + + +def _compact_body(document): + """Serialize exactly the way the SDK does for a compact item write. + + :param dict document: The item body to serialize. + :returns: The serialized body and its UTF-8 byte length. + :rtype: tuple + """ + body = _request_body_from_data(document, ensure_ascii=False) + return body, len(body) + + +class _DummyHeaderClient: + """Minimal connection object for generating Cosmos request headers.""" + + UseMultipleWriteLocations = False + master_key = None + resource_tokens = None + client_id = None + + class connection_policy: + ResponsePayloadOnWriteDisabled = False + + +def _patch_headers(content_type=None): + """Generate the headers used by a Cosmos item PATCH request.""" + options = {"partitionKey": "pk"} + if content_type: + options["contentType"] = content_type + return _base.GetHeaders( + _DummyHeaderClient(), + {}, + "patch", + "dbs/db/colls/container/docs", + "item", + http_constants.ResourceType.Document, + _OperationType.Patch, + options, + ) + + +class _LocalServer: + """A localhost HTTP server that captures one request.""" + + def __init__(self): + self.captured = {} + self._server = None + self._thread = None + + def __enter__(self): + """Start the server on an ephemeral port. + + :returns: This server, ready to receive one request. + :rtype: _LocalServer + """ + self._server = HTTPServer(("127.0.0.1", 0), _make_echo_handler(self.captured)) + self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + self._thread.start() + return self + + def __exit__(self, *exc_info): + """Shut the server down and join its thread.""" + self._server.shutdown() + self._server.server_close() + self._thread.join(timeout=5) + + @property + def url(self): + """Return the base URL the server is listening on. + + :returns: The server's URL. + :rtype: str + """ + host, port = self._server.server_address + return f"http://{host}:{port}/" + + +# These tests use a local HTTP server and do not require the Cosmos emulator. +# The standard Cosmos PR CI jobs run pytest with "-m cosmosEmulator", so this +# marker is required for these transport regressions to run in PR validation. +@pytest.mark.cosmosEmulator +class TestCompactBodyReachesTransportIntact(unittest.TestCase): + """The bytes the server receives must be the UTF-8 encoding of the body.""" + + def _round_trip_sync(self, document): + """Send one compact body through the real sync transport. + + :param dict document: The item body to send. + :returns: The serialized body, its byte length, and what the server saw. + :rtype: tuple + """ + body, byte_length = _compact_body(document) + with _LocalServer() as server: + request = HttpRequest("POST", server.url) + request.data = body + request.headers["Content-Length"] = str(byte_length) + with RequestsTransport() as transport: + transport.send(request) + return body, byte_length, server.captured + + def _round_trip_patch_sync(self, document, compact): + """Send a PATCH body through the real sync transport.""" + if compact: + body, byte_length = _compact_body(document) + else: + body = json.dumps(document, separators=(",", ":")) + byte_length = len(body.encode("utf-8")) + with _LocalServer() as server: + request = HttpRequest("PATCH", server.url, headers=_patch_headers()) + request.data = body + request.headers["Content-Length"] = str(byte_length) + with RequestsTransport() as transport: + transport.send(request) + return server.captured + + def test_payload_constants_are_the_intended_characters(self): + """Guard against this file being rewritten with the wrong encoding, + which would leave the tests passing against garbled text instead of the + characters they are meant to cover.""" + self.assertEqual(_CJK["x"].encode("utf-8"), b"\xe6\x97\xa5\xe6\x9c\xac") + self.assertEqual(_LATIN1_REPRESENTABLE["x"].encode("utf-8"), b"\xc3\xa9") + self.assertEqual(len(_LATIN1_REPRESENTABLE_UTF8), 10) + self.assertEqual(len(_LATIN1_REPRESENTABLE_LATIN1), 9) + + def test_compact_body_is_bytes_not_str(self): + """The serializer returns ``bytes`` for a compact body, for both a + payload with no Latin-1 representation and one that has one.""" + for label, document in (("cjk", _CJK), ("latin1", _LATIN1_REPRESENTABLE)): + with self.subTest(payload=label): + body, _ = _compact_body(document) + self.assertIsInstance(body, bytes) + self.assertNotIsInstance(body, str) + + def test_str_body_would_be_latin1_encoded_by_http_client(self): + """Latin-1 encoding of a ``str`` body either rejects the text outright + or produces bytes that are not its UTF-8 encoding. Asserts on stdlib + behavior rather than SDK behavior, to document what passing ``bytes`` + avoids.""" + cjk_text = json.dumps(_CJK, separators=(",", ":"), ensure_ascii=False) + with self.assertRaises(UnicodeEncodeError): + # CJK has no Latin-1 representation, so nothing would be sent. + _latin1_encode_like_http_client(cjk_text) + + latin_text = json.dumps(_LATIN1_REPRESENTABLE, separators=(",", ":"), ensure_ascii=False) + latin1_bytes = _latin1_encode_like_http_client(latin_text) + + # The character is representable in Latin-1, so this case would be sent + # silently with the wrong encoding: 0xE9 where the service requires the + # UTF-8 pair 0xC3 0xA9. + self.assertEqual(latin1_bytes, _LATIN1_REPRESENTABLE_LATIN1) + self.assertNotEqual(latin1_bytes, _LATIN1_REPRESENTABLE_UTF8) + self.assertEqual(latin_text.encode("utf-8"), _LATIN1_REPRESENTABLE_UTF8) + + def test_cjk_body_arrives_as_utf8_over_sync_transport(self): + """A CJK body reaches the server as UTF-8 with a matching length.""" + body, byte_length, captured = self._round_trip_sync(_CJK) + + self.assertEqual(captured["body"], body) + self.assertEqual(captured["body"], _CJK_UTF8) + self.assertEqual(json.loads(captured["body"].decode("utf-8")), _CJK) + self.assertEqual(int(captured["content_length"]), len(captured["body"])) + self.assertEqual(int(captured["content_length"]), byte_length) + + def test_latin1_representable_body_arrives_as_utf8_over_sync_transport(self): + """This character is representable in Latin-1, so a regression to a str + body would be sent without error but with the wrong bytes. Assert the + exact UTF-8 encoding rather than only the length, because a Latin-1 body + is self-consistent with its own Content-Length.""" + body, byte_length, captured = self._round_trip_sync(_LATIN1_REPRESENTABLE) + + self.assertEqual(captured["body"], body) + self.assertEqual(captured["body"], _LATIN1_REPRESENTABLE_UTF8) + self.assertNotEqual(captured["body"], _LATIN1_REPRESENTABLE_LATIN1) + self.assertEqual(int(captured["content_length"]), len(captured["body"])) + self.assertEqual(int(captured["content_length"]), byte_length) + + def test_patch_content_type_is_json_patch_for_str_and_bytes(self): + """The explicit Cosmos media type prevents transport auto-detection + from changing PATCH semantics when compact mode changes the body type.""" + for compact in (False, True): + with self.subTest(compact=compact): + captured = self._round_trip_patch_sync(_CJK, compact) + self.assertEqual( + captured["content_type"], + "application/json-patch+json", + ) + + def test_patch_content_type_override_is_preserved(self): + """Caller-supplied content types remain authoritative.""" + custom_content_type = "application/vnd.contoso.patch+json" + self.assertEqual( + _patch_headers(custom_content_type)[http_constants.HttpHeaders.ContentType], + custom_content_type, + ) + + +@pytest.mark.cosmosEmulator +class TestCompactBodyReachesTransportIntactAsync(unittest.IsolatedAsyncioTestCase): + """Async twin of the sync checks. aiohttp encodes a ``str`` body as UTF-8 + rather than Latin-1, so these assert the same wire contract to keep the two + stacks from drifting apart.""" + + async def _round_trip_async(self, document): + """Send one compact body through the real async transport. + + :param dict document: The item body to send. + :returns: The serialized body, its byte length, and what the server saw. + :rtype: tuple + """ + body, byte_length = _compact_body(document) + with _LocalServer() as server: + request = HttpRequest("POST", server.url) + request.data = body + request.headers["Content-Length"] = str(byte_length) + async with AioHttpTransport() as transport: + await transport.send(request) + # The handler records the body before it answers, so the capture is + # complete once the response has been received. + return body, byte_length, server.captured + + async def _round_trip_patch_async(self, document, compact): + """Send a PATCH body through the real async transport.""" + if compact: + body, byte_length = _compact_body(document) + else: + body = json.dumps(document, separators=(",", ":")) + byte_length = len(body.encode("utf-8")) + with _LocalServer() as server: + request = HttpRequest("PATCH", server.url, headers=_patch_headers()) + request.data = body + request.headers["Content-Length"] = str(byte_length) + async with AioHttpTransport() as transport: + await transport.send(request) + return server.captured + + async def test_cjk_body_arrives_as_utf8_over_async_transport(self): + """Async twin of the CJK check.""" + body, byte_length, captured = await self._round_trip_async(_CJK) + + self.assertEqual(captured["body"], body) + self.assertEqual(captured["body"], _CJK_UTF8) + self.assertEqual(int(captured["content_length"]), len(captured["body"])) + self.assertEqual(int(captured["content_length"]), byte_length) + + async def test_latin1_representable_body_arrives_as_utf8_over_async_transport(self): + """Async twin of the Latin-1-representable check, asserting exact UTF-8.""" + body, byte_length, captured = await self._round_trip_async(_LATIN1_REPRESENTABLE) + + self.assertEqual(captured["body"], body) + self.assertEqual(captured["body"], _LATIN1_REPRESENTABLE_UTF8) + self.assertNotEqual(captured["body"], _LATIN1_REPRESENTABLE_LATIN1) + self.assertEqual(int(captured["content_length"]), len(captured["body"])) + self.assertEqual(int(captured["content_length"]), byte_length) + + async def test_patch_content_type_is_json_patch_for_str_and_bytes(self): + """Async transport sees the same explicit media type in both modes.""" + for compact in (False, True): + with self.subTest(compact=compact): + captured = await self._round_trip_patch_async(_CJK, compact) + self.assertEqual( + captured["content_type"], + "application/json-patch+json", + ) + + +if __name__ == "__main__": + unittest.main()