Skip to content

Commit f6ac36d

Browse files
dibahlfidibahlfi
authored
Session Token Manual Override Fix (#42965)
* fix: fixing session token bug * fix: adding more tests * fix: refactoring * fix: adding space
1 parent e4fde9e commit f6ac36d

File tree

4 files changed

+138
-3
lines changed

4 files changed

+138
-3
lines changed

‎sdk/cosmos/azure-cosmos/CHANGELOG.md‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77
#### Breaking Changes
88

99
#### Bugs Fixed
10-
10+
* 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)
11+
1112
#### Other Changes
1213

1314
### 4.14.0b3 (2025-09-09)

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -335,8 +335,7 @@ def _is_session_token_request(
335335
# Verify that it is not a metadata request, and that it is either a read request, batch request, or an account
336336
# configured to use multiple write regions. Batch requests are special-cased because they can contain both read and
337337
# write operations, and we want to use session consistency for the read operations.
338-
return (is_session_consistency is True and cosmos_client_connection.session is not None
339-
and not IsMasterResource(request_object.resource_type)
338+
return (is_session_consistency is True and not IsMasterResource(request_object.resource_type)
340339
and (documents._OperationType.IsReadOnlyOperation(request_object.operation_type)
341340
or request_object.operation_type == "Batch"
342341
or cosmos_client_connection._global_endpoint_manager.can_use_multiple_write_locations(request_object)))

‎sdk/cosmos/azure-cosmos/tests/test_session.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,73 @@ def setUpClass(cls):
4646
cls.created_db = cls.client.get_database_client(cls.TEST_DATABASE_ID)
4747
cls.created_collection = cls.created_db.get_container_client(cls.TEST_COLLECTION_ID)
4848

49+
def test_manual_session_token_takes_precedence(self):
50+
# Establish an initial session state for the primary client. After this call, self.client has an internal session token.
51+
self.created_collection.create_item(
52+
body={'id': 'precedence_doc_1' + str(uuid.uuid4()), 'pk': 'mypk'}
53+
)
54+
# Capture the session token from the primary client (Token A)
55+
token_A = self.client.client_connection.last_response_headers.get(HttpHeaders.SessionToken)
56+
self.assertIsNotNone(token_A)
57+
58+
# Use a separate client to create a second item. This gives us a new, distinct session token from the response.
59+
with cosmos_client.CosmosClient(self.host, self.masterKey) as other_client:
60+
other_collection = other_client.get_database_client(self.TEST_DATABASE_ID) \
61+
.get_container_client(self.TEST_COLLECTION_ID)
62+
item2 = other_collection.create_item(
63+
body={'id': 'precedence_doc_2' + str(uuid.uuid4()), 'pk': 'mypk'}
64+
)
65+
# Capture the session token from the second client (Token B)
66+
manual_session_token = other_client.client_connection.last_response_headers.get(HttpHeaders.SessionToken)
67+
self.assertIsNotNone(manual_session_token)
68+
69+
# Assert that the two tokens are different to ensure we are testing a real override scenario.
70+
self.assertNotEqual(token_A, manual_session_token)
71+
72+
# 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.
73+
def manual_token_hook(request):
74+
# Assert that the header contains the manually provided second token not the client's automatic first token.
75+
self.assertIn(HttpHeaders.SessionToken, request.http_request.headers)
76+
self.assertEqual(request.http_request.headers[HttpHeaders.SessionToken], manual_session_token)
77+
78+
#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.
79+
self.created_collection.read_item(
80+
item=item2['id'], # Reading the item associated with second token
81+
partition_key='mypk',
82+
session_token=manual_session_token, # Manually provide second token
83+
raw_request_hook=manual_token_hook
84+
)
85+
86+
def test_manual_session_token_override(self):
87+
# Create an item to get a valid session token from the response
88+
created_document = self.created_collection.create_item(
89+
body={'id': 'doc_for_manual_session' + str(uuid.uuid4()), 'pk': 'mypk'}
90+
)
91+
session_token = self.client.client_connection.last_response_headers.get(HttpHeaders.SessionToken)
92+
self.assertIsNotNone(session_token)
93+
94+
# temporarily disable client-side session management to test manual override
95+
original_session = self.client.client_connection.session
96+
self.client.client_connection.session = None
97+
98+
try:
99+
# Define a hook to inspect the request headers
100+
def manual_token_hook(request):
101+
self.assertIn(HttpHeaders.SessionToken, request.http_request.headers)
102+
self.assertEqual(request.http_request.headers[HttpHeaders.SessionToken], session_token)
103+
104+
# Read the item, passing the session token manually.
105+
# The hook will verify it's correctly added to the request headers.
106+
self.created_collection.read_item(
107+
item=created_document['id'],
108+
partition_key='mypk',
109+
session_token=session_token, # Manually provide the session token
110+
raw_request_hook=manual_token_hook
111+
)
112+
finally:
113+
# Restore the original session object to avoid affecting other tests
114+
self.client.client_connection.session = original_session
115+
49116
def test_session_token_sm_for_ops(self):
50117

51118
# Session token should not be sent for control plane operations

‎sdk/cosmos/azure-cosmos/tests/test_session_async.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,74 @@ async def asyncSetUp(self):
4848
async def asyncTearDown(self):
4949
await self.client.close()
5050

51+
async def test_manual_session_token_takes_precedence_async(self):
52+
# Establish an initial session state for the primary async client.
53+
await self.created_container.create_item(
54+
body={'id': 'precedence_doc_1_async' + str(uuid.uuid4()), 'pk': 'mypk'}
55+
)
56+
# Capture the session token from the primary client (Token A)
57+
token_A = self.client.client_connection.last_response_headers.get(HttpHeaders.SessionToken)
58+
self.assertIsNotNone(token_A)
59+
60+
# Use a separate async client to create a second item. This gives us a new, distinct session token.
61+
async with CosmosClient(self.host, self.masterKey) as other_client:
62+
other_collection = other_client.get_database_client(self.TEST_DATABASE_ID) \
63+
.get_container_client(self.TEST_COLLECTION_ID)
64+
item2 = await other_collection.create_item(
65+
body={'id': 'precedence_doc_2_async' + str(uuid.uuid4()), 'pk': 'mypk'}
66+
)
67+
# Capture the session token from the second client (Token B)
68+
manual_session_token = other_client.client_connection.last_response_headers.get(HttpHeaders.SessionToken)
69+
self.assertIsNotNone(manual_session_token)
70+
71+
# Assert that the two tokens are different to ensure we are testing a real override scenario.
72+
self.assertNotEqual(token_A, manual_session_token)
73+
74+
# Define a hook to verify the correct token is sent.
75+
def manual_token_hook(request):
76+
# Assert that the header contains the manually provided Token B, not the client's automatic Token A.
77+
self.assertIn(HttpHeaders.SessionToken, request.http_request.headers)
78+
self.assertEqual(request.http_request.headers[HttpHeaders.SessionToken], manual_session_token)
79+
80+
# Read an item using the primary client, but manually providing Token B.
81+
# The hook will verify that Token B overrides the client's internal Token A.
82+
await self.created_container.read_item(
83+
item=item2['id'],
84+
partition_key='mypk',
85+
session_token=manual_session_token, # Manually provide Token B
86+
raw_request_hook=manual_token_hook
87+
)
88+
89+
async def test_manual_session_token_override_async(self):
90+
# Create an item to get a valid session token from the response
91+
created_document = await self.created_container.create_item(
92+
body={'id': 'doc_for_manual_session' + str(uuid.uuid4()), 'pk': 'mypk'}
93+
)
94+
session_token = self.client.client_connection.last_response_headers.get(HttpHeaders.SessionToken)
95+
self.assertIsNotNone(session_token)
96+
97+
# temporarily disable client-side session management to test manual override
98+
original_session = self.client.client_connection.session
99+
self.client.client_connection.session = None
100+
101+
try:
102+
# Define a hook to inspect the request headers
103+
def manual_token_hook(request):
104+
self.assertIn(HttpHeaders.SessionToken, request.http_request.headers)
105+
self.assertEqual(request.http_request.headers[HttpHeaders.SessionToken], session_token)
106+
107+
# Read the item, passing the session token manually.
108+
# The hook will verify it's correctly added to the request headers.
109+
await self.created_container.read_item(
110+
item=created_document['id'],
111+
partition_key='mypk',
112+
session_token=session_token, # Manually provide the session token
113+
raw_request_hook=manual_token_hook
114+
)
115+
finally:
116+
# Restore the original session object to avoid affecting other tests
117+
self.client.client_connection.session = original_session
118+
51119
async def test_session_token_swr_for_ops_async(self):
52120
# Session token should not be sent for control plane operations
53121
test_container = await self.created_db.create_container(str(uuid.uuid4()), PartitionKey(path="/id"), raw_response_hook=test_config.no_token_response_hook)

0 commit comments

Comments
 (0)