From 1e72d4ef7722080a2cd99b56fbde59e4c59ae1e8 Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Thu, 18 Dec 2025 21:59:06 -0600 Subject: [PATCH 1/2] fix: fixing sending redundant session token information --- .../azure/cosmos/_cosmos_client_connection.py | 4 ++ .../aio/_cosmos_client_connection_async.py | 4 ++ sdk/cosmos/azure-cosmos/tests/test_session.py | 45 +++++++++++++++++++ .../azure-cosmos/tests/test_session_async.py | 45 +++++++++++++++++++ 4 files changed, 98 insertions(+) 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 0c7b619b0ae8..7e8df25a3dca 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -3336,6 +3336,10 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]: EPK_sub_range = routing_range.Range(range_min=max(single_range.min, feed_range_epk.min), range_max=min(single_range.max, feed_range_epk.max), isMinInclusive=True, isMaxInclusive=False) + + # set the session token for this specific partition to avoid sending compound token for all partitions + base.set_session_token_header(self, req_headers, path, request_params, options, + over_lapping_range["id"]) if single_range.min == EPK_sub_range.min and EPK_sub_range.max == single_range.max: # The Epk Sub Range spans exactly one physical partition # In this case we can route to the physical pk range id 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 efa5188c70fe..69a6a5430b88 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 @@ -3128,6 +3128,10 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]: EPK_sub_range = routing_range.Range(range_min=max(single_range.min, feed_range_epk.min), range_max=min(single_range.max, feed_range_epk.max), isMinInclusive=True, isMaxInclusive=False) + + # set the session token for this specific partition to avoid sending compound token for all partitions + base.set_session_token_header(self, req_headers, path, request_params, options, + over_lapping_range["id"]) if single_range.min == EPK_sub_range.min and EPK_sub_range.max == single_range.max: # The Epk Sub Range spans exactly one physical partition # In this case we can route to the physical pk range id diff --git a/sdk/cosmos/azure-cosmos/tests/test_session.py b/sdk/cosmos/azure-cosmos/tests/test_session.py index e4a675236125..55e4e86def26 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_session.py +++ b/sdk/cosmos/azure-cosmos/tests/test_session.py @@ -153,6 +153,51 @@ def test_session_token_sm_for_ops(self): assert self.created_db.client_connection.last_response_headers.get(HttpHeaders.SessionToken) is not None assert self.created_db.client_connection.last_response_headers.get(HttpHeaders.SessionToken) != batch_response_token + def test_session_token_compound_not_sent_for_single_partition_query(self): + """ + Verify that when querying with a feed range (single physical partition), + only that partition's session token is sent, not the entire compound token. + """ + test_container = self.created_db.create_container( + "Container query test" + str(uuid.uuid4()), + PartitionKey(path="/pk"), + offer_throughput=11000 + ) + + try: + # Create items across multiple partition keys + for i in range(100): + test_container.create_item({ + 'id': str(uuid.uuid4()), + 'pk': f"pk_{i:04d}" + }) + + # Get feed ranges and verify multiple exist + feed_ranges = list(test_container.read_feed_ranges()) + self.assertGreater(len(feed_ranges), 1, "Expected multiple feed ranges") + + # Capture session token sent with feed range query + captured_session_token = {} + + def capture_session_token(request): + captured_session_token['token'] = request.http_request.headers.get(HttpHeaders.SessionToken) + + # Query with single feed range + list(test_container.query_items( + query="SELECT * FROM c", + feed_range=feed_ranges[0], + raw_request_hook=capture_session_token + )) + + # Verify only single partition token was sent + token = captured_session_token.get('token') + self.assertIsNotNone(token, "Session token should be present") + self.assertNotIn(',', token, + f"Expected single partition token, got compound token: {token}") + + finally: + self.created_db.delete_container(test_container) + def test_session_token_with_space_in_container_name(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 6a991be540ab..e6f4aeb75cfb 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_session_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_session_async.py @@ -86,6 +86,51 @@ def manual_token_hook(request): raw_request_hook=manual_token_hook ) + async def test_session_token_compound_not_sent_for_single_partition_query_async(self): + """ + Verify that when querying with a feed range (single physical partition), + only that partition's session token is sent, not the entire compound token. + """ + test_container = await self.created_db.create_container( + "Container query test" + str(uuid.uuid4()), + PartitionKey(path="/pk"), + offer_throughput=11000 + ) + + try: + # Create items across multiple partition keys + for i in range(100): + await test_container.create_item({ + 'id': str(uuid.uuid4()), + 'pk': f"pk_{i:04d}" + }) + + # Get feed ranges and verify multiple exist + feed_ranges = [feed_range async for feed_range in test_container.read_feed_ranges()] + self.assertGreater(len(feed_ranges), 1, "Expected multiple feed ranges") + + # Capture session token sent with feed range query + captured_session_token = {} + + def capture_session_token(request): + captured_session_token['token'] = request.http_request.headers.get(HttpHeaders.SessionToken) + + # Query with single feed range + _ = [item async for item in test_container.query_items( + query="SELECT * FROM c", + feed_range=feed_ranges[0], + raw_request_hook=capture_session_token + )] + + # Verify only single partition token was sent + token = captured_session_token.get('token') + self.assertIsNotNone(token, "Session token should be present") + self.assertNotIn(',', token, + f"Expected single partition token, got compound token: {token}") + + finally: + await self.created_db.delete_container(test_container) + 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 7007f4873191e2367054dce88fad94c8dac7f6a2 Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Thu, 18 Dec 2025 22:12:46 -0600 Subject: [PATCH 2/2] fix: adding async /await --- .../azure/cosmos/aio/_cosmos_client_connection_async.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 69a6a5430b88..6ba5e1cf9b5c 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 @@ -3130,7 +3130,7 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]: isMinInclusive=True, isMaxInclusive=False) # set the session token for this specific partition to avoid sending compound token for all partitions - base.set_session_token_header(self, req_headers, path, request_params, options, + await base.set_session_token_header_async(self, req_headers, path, request_params, options, over_lapping_range["id"]) if single_range.min == EPK_sub_range.min and EPK_sub_range.max == single_range.max: # The Epk Sub Range spans exactly one physical partition