[Cosmos] Session container fixes

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -20,8 +20,11 @@
 #### Bugs Fixed
 * Fixed how resource tokens are parsed for metadata calls in the lifecycle of a document operation. See [PR 40302](https://github.com/Azure/azure-sdk-for-python/pull/40302).
 * Fixed issue where Query Change Feed did not return items if the container uses legacy Hash V1 Partition Keys. This also fixes issues with not being able to change feed query for Specific Partition Key Values for HPK. See [PR 41270](https://github.com/Azure/azure-sdk-for-python/pull/41270/)
+* Fixed session container compound session token logic. The SDK will now only send the relevant partition-local session tokens for each read request, as opposed to the entire compound session token for the container. See [PR 40366](https://github.com/Azure/azure-sdk-for-python/pull/40366).
+* Write requests for single-write region accounts will no longer send session tokens when using session consistency. See [PR 40366](https://github.com/Azure/azure-sdk-for-python/pull/40366).
 
 #### Other Changes
+* Cross-partition queries will now always send a query plan before attempting to execute. See [PR 40366](https://github.com/Azure/azure-sdk-for-python/pull/40366).
 * Added Client Generated Activity IDs to all Requests. Cosmos Diagnostics Logs will more clearly show the Activity ID for each request and response. [PR 41013](https://github.com/Azure/azure-sdk-for-python/pull/41013)
 
 ### 4.12.0b1 (2025-05-19)

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

@@ -44,7 +44,9 @@
 if TYPE_CHECKING:
     from ._cosmos_client_connection import CosmosClientConnection
     from .aio._cosmos_client_connection_async import CosmosClientConnection as AsyncClientConnection
+    from ._request_object import RequestObject
 
+# pylint: disable=protected-access
 
 _COMMON_OPTIONS = {
     'initial_headers': 'initialHeaders',
@@ -172,37 +174,9 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
     if options.get("indexingDirective"):
         headers[http_constants.HttpHeaders.IndexingDirective] = options["indexingDirective"]
 
-    consistency_level = None
-
-    # get default client consistency level
-    default_client_consistency_level = headers.get(http_constants.HttpHeaders.ConsistencyLevel)
-
-    # set consistency level. check if set via options, this will override the default
+    # set request consistency level - if session consistency, the client should be setting this on its own
     if options.get("consistencyLevel"):
-        consistency_level = options["consistencyLevel"]
-        # TODO: move this line outside of if-else cause to remove the code duplication
-        headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level
-    elif default_client_consistency_level is not None:
-        consistency_level = default_client_consistency_level
-        headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level
-
-    # figure out if consistency level for this request is session
-    is_session_consistency = consistency_level == documents.ConsistencyLevel.Session
-
-    # set session token if required
-    if is_session_consistency is True and not IsMasterResource(resource_type):
-        # if there is a token set via option, then use it to override default
-        if options.get("sessionToken"):
-            headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"]
-        else:
-            # check if the client's default consistency is session (and request consistency level is same),
-            # then update from session container
-            if default_client_consistency_level == documents.ConsistencyLevel.Session and \
-                    cosmos_client_connection.session:
-                # populate session token from the client's session container
-                headers[http_constants.HttpHeaders.SessionToken] = cosmos_client_connection.session.get_session_token(
-                    path
-                )
+        headers[http_constants.HttpHeaders.ConsistencyLevel] = options["consistencyLevel"]
 
     if options.get("enableScanInQuery"):
         headers[http_constants.HttpHeaders.EnableScanInQuery] = options["enableScanInQuery"]
@@ -345,6 +319,76 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
 
     return headers
 
+def _is_session_token_request(
+        cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"],
+        headers: dict,
+        request_object: "RequestObject") -> bool:
+    consistency_level = headers.get(http_constants.HttpHeaders.ConsistencyLevel)
+    # Figure out if consistency level for this request is session
+    is_session_consistency = consistency_level == documents.ConsistencyLevel.Session
+
+    # Verify that it is not a metadata request, and that it is either a read request, batch request, or an account
+    # configured to use multiple write regions
+    return (is_session_consistency is True and cosmos_client_connection.session is not None
+            and not IsMasterResource(request_object.resource_type)
+            and (documents._OperationType.IsReadOnlyOperation(request_object.operation_type)
+                 or request_object.operation_type == "Batch"
+                 or cosmos_client_connection._global_endpoint_manager.can_use_multiple_write_locations(request_object)))
+
+
+def set_session_token_header(
+        cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"],
+        headers: dict,
+        path: str,
+        request_object: "RequestObject",
+        options: Mapping[str, Any],
+        partition_key_range_id: Optional[str] = None) -> None:
+    # set session token if required
+    if _is_session_token_request(cosmos_client_connection, headers, request_object):
+        # if there is a token set via option, then use it to override default
+        if options.get("sessionToken"):
+            headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"]
+        else:
+            # check if the client's default consistency is session (and request consistency level is same),
+            # then update from session container
+            if headers[http_constants.HttpHeaders.ConsistencyLevel] == documents.ConsistencyLevel.Session and \
+                    cosmos_client_connection.session:
+                # populate session token from the client's session container
+                session_token = (
+                    cosmos_client_connection.session.get_session_token(path,
+                                                                options.get('partitionKey'),
+                                                                cosmos_client_connection._container_properties_cache,
+                                                                cosmos_client_connection._routing_map_provider,
+                                                                partition_key_range_id))
+                if session_token != "":
+                    headers[http_constants.HttpHeaders.SessionToken] = session_token
+
+async def set_session_token_header_async(
+        cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"],
+        headers: dict,
+        path: str,
+        request_object: "RequestObject",
+        options: Mapping[str, Any],
+        partition_key_range_id: Optional[str] = None) -> None:
+    # set session token if required
+    if _is_session_token_request(cosmos_client_connection, headers, request_object):
+        # if there is a token set via option, then use it to override default
+        if options.get("sessionToken"):
+            headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"]
+        else:
+            # check if the client's default consistency is session (and request consistency level is same),
+            # then update from session container
+            if headers[http_constants.HttpHeaders.ConsistencyLevel] == documents.ConsistencyLevel.Session and \
+                    cosmos_client_connection.session:
+                # populate session token from the client's session container
+                session_token = \
+                    await cosmos_client_connection.session.get_session_token_async(path,
+                                                                options.get('partitionKey'),
+                                                                cosmos_client_connection._container_properties_cache,
+                                                                cosmos_client_connection._routing_map_provider,
+                                                                partition_key_range_id)
+                if session_token != "":
+                    headers[http_constants.HttpHeaders.SessionToken] = session_token
 
 def GetResourceIdOrFullNameFromLink(resource_link: str) -> str:
     """Gets resource id or full name from resource link.