From 91aafb5cd10b274921cf31adf0a6837845706a88 Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Thu, 11 Sep 2025 11:01:52 -0500 Subject: [PATCH 1/4] fix: fixing session token bug --- sdk/cosmos/azure-cosmos/azure/cosmos/_base.py | 88 +++++++++---------- sdk/cosmos/azure-cosmos/tests/test_session.py | 30 +++++++ .../azure-cosmos/tests/test_session_async.py | 30 +++++++ 3 files changed, 104 insertions(+), 44 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py index 040b98d476c4..bfe491f026c7 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py @@ -349,28 +349,28 @@ def set_session_token_header( 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: - # urllib_unquote is used to decode the path, as it may contain encoded characters - path = urllib_unquote(path) - # 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, - options)) - if session_token != "": - headers[http_constants.HttpHeaders.SessionToken] = session_token + # If a session token is explicitly provided in options, use it. This allows manual override + # even when client-side session management is disabled. + if options.get("sessionToken"): + headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"] + # If no manual token is provided, check if we should use the client's session management. + elif _is_session_token_request(cosmos_client_connection, headers, request_object): + # 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: + # urllib_unquote is used to decode the path, as it may contain encoded characters + path = urllib_unquote(path) + # 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, + options)) + if session_token != "": + headers[http_constants.HttpHeaders.SessionToken] = session_token async def set_session_token_header_async( cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"], @@ -379,28 +379,28 @@ async def set_session_token_header_async( 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 - # urllib_unquote is used to decode the path, as it may contain encoded characters - path = urllib_unquote(path) - 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, - options) - if session_token != "": - headers[http_constants.HttpHeaders.SessionToken] = session_token + # If a session token is explicitly provided in options, use it. This allows manual override + # even when client-side session management is disabled. + if options.get("sessionToken"): + headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"] + # If no manual token is provided, check if we should use the client's session management. + elif _is_session_token_request(cosmos_client_connection, headers, request_object): + # 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 + # urllib_unquote is used to decode the path, as it may contain encoded characters + path = urllib_unquote(path) + 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, + options) + 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. diff --git a/sdk/cosmos/azure-cosmos/tests/test_session.py b/sdk/cosmos/azure-cosmos/tests/test_session.py index 3a88c5ca5fbe..8a7a3debcc9e 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_session.py +++ b/sdk/cosmos/azure-cosmos/tests/test_session.py @@ -46,6 +46,36 @@ def setUpClass(cls): cls.created_db = cls.client.get_database_client(cls.TEST_DATABASE_ID) cls.created_collection = cls.created_db.get_container_client(cls.TEST_COLLECTION_ID) + def test_manual_session_token_override(self): + # Create an item to get a valid session token from the response + created_document = self.created_collection.create_item( + body={'id': 'doc_for_manual_session' + str(uuid.uuid4()), 'pk': 'mypk'} + ) + session_token = self.client.client_connection.last_response_headers.get(HttpHeaders.SessionToken) + self.assertIsNotNone(session_token) + + # temporarily disable client-side session management to test manual override + original_session = self.client.client_connection.session + self.client.client_connection.session = None + + try: + # Define a hook to inspect the request headers + def manual_token_hook(request): + self.assertIn(HttpHeaders.SessionToken, request.http_request.headers) + self.assertEqual(request.http_request.headers[HttpHeaders.SessionToken], session_token) + + # Read the item, passing the session token manually. + # The hook will verify it's correctly added to the request headers. + self.created_collection.read_item( + item=created_document['id'], + partition_key='mypk', + session_token=session_token, # Manually provide the session token + raw_request_hook=manual_token_hook + ) + finally: + # Restore the original session object to avoid affecting other tests + self.client.client_connection.session = original_session + def test_session_token_sm_for_ops(self): # Session token should not be sent for control plane operations diff --git a/sdk/cosmos/azure-cosmos/tests/test_session_async.py b/sdk/cosmos/azure-cosmos/tests/test_session_async.py index 9947f378daa5..1e0586e36283 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_session_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_session_async.py @@ -48,6 +48,36 @@ async def asyncSetUp(self): async def asyncTearDown(self): await self.client.close() + async def test_manual_session_token_override_async(self): + # Create an item to get a valid session token from the response + created_document = await self.created_container.create_item( + body={'id': 'doc_for_manual_session' + str(uuid.uuid4()), 'pk': 'mypk'} + ) + session_token = self.client.client_connection.last_response_headers.get(HttpHeaders.SessionToken) + self.assertIsNotNone(session_token) + + # temporarily disable client-side session management to test manual override + original_session = self.client.client_connection.session + self.client.client_connection.session = None + + try: + # Define a hook to inspect the request headers + def manual_token_hook(request): + self.assertIn(HttpHeaders.SessionToken, request.http_request.headers) + self.assertEqual(request.http_request.headers[HttpHeaders.SessionToken], session_token) + + # Read the item, passing the session token manually. + # The hook will verify it's correctly added to the request headers. + await self.created_container.read_item( + item=created_document['id'], + partition_key='mypk', + session_token=session_token, # Manually provide the session token + raw_request_hook=manual_token_hook + ) + finally: + # Restore the original session object to avoid affecting other tests + self.client.client_connection.session = original_session + async def test_session_token_swr_for_ops_async(self): # Session token should not be sent for control plane operations test_container = await self.created_db.create_container(str(uuid.uuid4()), PartitionKey(path="/id"), raw_response_hook=test_config.no_token_response_hook) From ed9e457e223fe467749c32997f69b3fc7a7999d0 Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Thu, 11 Sep 2025 11:46:05 -0500 Subject: [PATCH 2/4] fix: adding more tests --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 3 +- sdk/cosmos/azure-cosmos/tests/test_session.py | 37 ++++++++++++++++++ .../azure-cosmos/tests/test_session_async.py | 38 +++++++++++++++++++ 3 files changed, 77 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 90a29db5c680..57b1a0df00bf 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -7,7 +7,8 @@ #### Breaking Changes #### Bugs Fixed - +* Fixed bug where client provided session token was not respected when client-side session management was disabled. See [PR 42965](https://github.com/Azure/azure-sdk-for-python/pull/42965) + #### Other Changes ### 4.14.0b3 (2025-09-09) diff --git a/sdk/cosmos/azure-cosmos/tests/test_session.py b/sdk/cosmos/azure-cosmos/tests/test_session.py index 8a7a3debcc9e..5bccd7a30039 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_session.py +++ b/sdk/cosmos/azure-cosmos/tests/test_session.py @@ -46,6 +46,43 @@ def setUpClass(cls): cls.created_db = cls.client.get_database_client(cls.TEST_DATABASE_ID) cls.created_collection = cls.created_db.get_container_client(cls.TEST_COLLECTION_ID) + def test_manual_session_token_takes_precedence(self): + # Establish an initial session state for the primary client.After this call, self.client has an internal session token. + self.created_collection.create_item( + body={'id': 'precedence_doc_1' + str(uuid.uuid4()), 'pk': 'mypk'} + ) + # Capture the session token from the primary client (Token A) + token_A = self.client.client_connection.last_response_headers.get(HttpHeaders.SessionToken) + self.assertIsNotNone(token_A) + + # Use a separate client to create a second item. This gives us a new, distinct session token from the response. + with cosmos_client.CosmosClient(self.host, self.masterKey) as other_client: + other_collection = other_client.get_database_client(self.TEST_DATABASE_ID) \ + .get_container_client(self.TEST_COLLECTION_ID) + item2 = other_collection.create_item( + body={'id': 'precedence_doc_2' + str(uuid.uuid4()), 'pk': 'mypk'} + ) + # Capture the session token from the second client (Token B) + manual_session_token = other_client.client_connection.last_response_headers.get(HttpHeaders.SessionToken) + self.assertIsNotNone(manual_session_token) + + # Assert that the two tokens are different to ensure we are testing a real override scenario. + self.assertNotEqual(token_A, manual_session_token) + + # At this point, self.client's session is at first token, but we are holding second token. We will now manually use second token in a request on self.client. + def manual_token_hook(request): + # Assert that the header contains the manually provided second token not the client's automatic first token. + self.assertIn(HttpHeaders.SessionToken, request.http_request.headers) + self.assertEqual(request.http_request.headers[HttpHeaders.SessionToken], manual_session_token) + + #Read an item using the primary client, but manually providing second token. The hook will verify that second token overrides the client's internal first token. + self.created_collection.read_item( + item=item2['id'], # Reading the item associated with second token + partition_key='mypk', + session_token=manual_session_token, # Manually provide second token + raw_request_hook=manual_token_hook + ) + def test_manual_session_token_override(self): # Create an item to get a valid session token from the response created_document = self.created_collection.create_item( diff --git a/sdk/cosmos/azure-cosmos/tests/test_session_async.py b/sdk/cosmos/azure-cosmos/tests/test_session_async.py index 1e0586e36283..6a991be540ab 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_session_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_session_async.py @@ -48,6 +48,44 @@ async def asyncSetUp(self): async def asyncTearDown(self): await self.client.close() + async def test_manual_session_token_takes_precedence_async(self): + # Establish an initial session state for the primary async client. + await self.created_container.create_item( + body={'id': 'precedence_doc_1_async' + str(uuid.uuid4()), 'pk': 'mypk'} + ) + # Capture the session token from the primary client (Token A) + token_A = self.client.client_connection.last_response_headers.get(HttpHeaders.SessionToken) + self.assertIsNotNone(token_A) + + # Use a separate async client to create a second item. This gives us a new, distinct session token. + async with CosmosClient(self.host, self.masterKey) as other_client: + other_collection = other_client.get_database_client(self.TEST_DATABASE_ID) \ + .get_container_client(self.TEST_COLLECTION_ID) + item2 = await other_collection.create_item( + body={'id': 'precedence_doc_2_async' + str(uuid.uuid4()), 'pk': 'mypk'} + ) + # Capture the session token from the second client (Token B) + manual_session_token = other_client.client_connection.last_response_headers.get(HttpHeaders.SessionToken) + self.assertIsNotNone(manual_session_token) + + # Assert that the two tokens are different to ensure we are testing a real override scenario. + self.assertNotEqual(token_A, manual_session_token) + + # Define a hook to verify the correct token is sent. + def manual_token_hook(request): + # Assert that the header contains the manually provided Token B, not the client's automatic Token A. + self.assertIn(HttpHeaders.SessionToken, request.http_request.headers) + self.assertEqual(request.http_request.headers[HttpHeaders.SessionToken], manual_session_token) + + # Read an item using the primary client, but manually providing Token B. + # The hook will verify that Token B overrides the client's internal Token A. + await self.created_container.read_item( + item=item2['id'], + partition_key='mypk', + session_token=manual_session_token, # Manually provide Token B + raw_request_hook=manual_token_hook + ) + async def test_manual_session_token_override_async(self): # Create an item to get a valid session token from the response created_document = await self.created_container.create_item( From fac473f122c2144def4e0132da6b2a5bd39df2fd Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Thu, 11 Sep 2025 11:57:58 -0500 Subject: [PATCH 3/4] fix: refactoring --- sdk/cosmos/azure-cosmos/azure/cosmos/_base.py | 91 +++++++++---------- 1 file changed, 45 insertions(+), 46 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py index bfe491f026c7..8aabfae1ffe2 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py @@ -335,8 +335,7 @@ def _is_session_token_request( # 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. Batch requests are special-cased because they can contain both read and # write operations, and we want to use session consistency for the read operations. - return (is_session_consistency is True and cosmos_client_connection.session is not None - and not IsMasterResource(request_object.resource_type) + return (is_session_consistency is True 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))) @@ -349,28 +348,28 @@ def set_session_token_header( request_object: "RequestObject", options: Mapping[str, Any], partition_key_range_id: Optional[str] = None) -> None: - # If a session token is explicitly provided in options, use it. This allows manual override - # even when client-side session management is disabled. - if options.get("sessionToken"): - headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"] - # If no manual token is provided, check if we should use the client's session management. - elif _is_session_token_request(cosmos_client_connection, headers, request_object): - # 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: - # urllib_unquote is used to decode the path, as it may contain encoded characters - path = urllib_unquote(path) - # 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, - options)) - if session_token != "": - headers[http_constants.HttpHeaders.SessionToken] = session_token + # 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: + # urllib_unquote is used to decode the path, as it may contain encoded characters + path = urllib_unquote(path) + # 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, + options)) + if session_token != "": + headers[http_constants.HttpHeaders.SessionToken] = session_token async def set_session_token_header_async( cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"], @@ -379,28 +378,28 @@ async def set_session_token_header_async( request_object: "RequestObject", options: Mapping[str, Any], partition_key_range_id: Optional[str] = None) -> None: - # If a session token is explicitly provided in options, use it. This allows manual override - # even when client-side session management is disabled. - if options.get("sessionToken"): - headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"] - # If no manual token is provided, check if we should use the client's session management. - elif _is_session_token_request(cosmos_client_connection, headers, request_object): - # 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 - # urllib_unquote is used to decode the path, as it may contain encoded characters - path = urllib_unquote(path) - 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, - options) - if session_token != "": - headers[http_constants.HttpHeaders.SessionToken] = session_token + # 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 + # urllib_unquote is used to decode the path, as it may contain encoded characters + path = urllib_unquote(path) + 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, + options) + 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. From f3a2a38c52b3dabf6112939123e3c40b202ebe05 Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Thu, 11 Sep 2025 14:05:47 -0500 Subject: [PATCH 4/4] fix: adding space --- sdk/cosmos/azure-cosmos/tests/test_session.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_session.py b/sdk/cosmos/azure-cosmos/tests/test_session.py index 5bccd7a30039..e4a675236125 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_session.py +++ b/sdk/cosmos/azure-cosmos/tests/test_session.py @@ -47,7 +47,7 @@ def setUpClass(cls): cls.created_collection = cls.created_db.get_container_client(cls.TEST_COLLECTION_ID) def test_manual_session_token_takes_precedence(self): - # Establish an initial session state for the primary client.After this call, self.client has an internal session token. + # Establish an initial session state for the primary client. After this call, self.client has an internal session token. self.created_collection.create_item( body={'id': 'precedence_doc_1' + str(uuid.uuid4()), 'pk': 'mypk'} )