diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index b8b1c451ad10..371c5358d893 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -7,6 +7,7 @@ #### Breaking Changes #### Bugs Fixed +* Fixed `KeyError: 'version'` in `SessionContainer.get_session_token` (sync and async) when the container's `partitionKey` definition returned by the service does not include the optional `version` field. The error was silently swallowed by a broad `except`, causing the client to send no `x-ms-session-token` header on subsequent reads. Against the Dedicated Gateway, this turned every Session-consistency read into an Integrated Cache miss. `partitionKey.version` is now treated as optional and defaults to `1`, matching how `PartitionKey` handles a missing version. See [PR 47143](https://github.com/Azure/azure-sdk-for-python/pull/47143) #### Other Changes diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session.py index bb1229b57662..5bf73e655b30 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_session.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_session.py @@ -119,7 +119,7 @@ def get_session_token( collection_pk_definition = container_properties_cache[collection_name]["partitionKey"] partition_key = PartitionKey(path=collection_pk_definition['paths'], kind=collection_pk_definition['kind'], - version=collection_pk_definition['version']) + version=collection_pk_definition.get('version', 1)) epk_range = partition_key._get_epk_range_for_partition_key(pk_value=pk_value) pk_range = routing_map_provider.get_overlapping_ranges(collection_name, [epk_range], @@ -213,7 +213,7 @@ async def get_session_token_async( collection_pk_definition = container_properties_cache[collection_name]["partitionKey"] partition_key = PartitionKey(path=collection_pk_definition['paths'], kind=collection_pk_definition['kind'], - version=collection_pk_definition['version']) + version=collection_pk_definition.get('version', 1)) epk_range = partition_key._get_epk_range_for_partition_key(pk_value=pk_value) pk_range = await routing_map_provider.get_overlapping_ranges(collection_name, [epk_range], diff --git a/sdk/cosmos/azure-cosmos/tests/test_session_token_unit.py b/sdk/cosmos/azure-cosmos/tests/test_session_token_unit.py index be368ef03321..0c5de8578251 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_session_token_unit.py +++ b/sdk/cosmos/azure-cosmos/tests/test_session_token_unit.py @@ -321,6 +321,152 @@ def __init__(self, t): self.assertEqual(result, token.session_token) +class TestGetSessionTokenWithoutPartitionKeyVersion(unittest.TestCase): + + COLLECTION_LINK = "dbs/db1/colls/c1" + COLLECTION_RID = "abc123==" + PK_RANGE_ID = "0" + PK_VALUE = "some-pk-value" + RAW_TOKEN = "1#100" + + class _CapturingRoutingMapProvider: + """Records the ranges argument passed to get_overlapping_ranges so tests + can assert that get_session_token computed the correct effective partition + key. Returns a single fixed partition range.""" + + def __init__(self, pk_range_id): + self._pk_range_id = pk_range_id + self.captured_ranges = None + + def get_overlapping_ranges(self, collection_link, ranges, options): + del collection_link, options + self.captured_ranges = list(ranges) + return [{ + "id": self._pk_range_id, + "minInclusive": "", + "maxExclusive": "FF", + "parents": [], + }] + + class _TokenWrap: + """Mimics what ``SessionContainer.rid_to_session_token`` stores per + partition: an object exposing ``.session_token`` as the string form.""" + + def __init__(self, vector_session_token): + self.session_token = vector_session_token.session_token + + def _build_container_with_seeded_token(self): + container = _session.SessionContainer() + container.collection_name_to_rid[self.COLLECTION_LINK] = self.COLLECTION_RID + container.rid_to_session_token[self.COLLECTION_RID] = { + self.PK_RANGE_ID: self._TokenWrap(VectorSessionToken.create(self.RAW_TOKEN)), + } + return container + + def _build_cache(self, partition_key_definition): + return { + self.COLLECTION_LINK: { + "_rid": self.COLLECTION_RID, + "partitionKey": partition_key_definition, + }, + } + + def test_partitionkey_definition_without_version_returns_token(self): + """Service responses without a 'version' key must not silently nuke the token.""" + container = self._build_container_with_seeded_token() + cache = self._build_cache({"paths": ["/pk"], "kind": "Hash"}) # no 'version' + + token = container.get_session_token( + resource_path=self.COLLECTION_LINK, + pk_value=self.PK_VALUE, + container_properties_cache=cache, + routing_map_provider=self._CapturingRoutingMapProvider(self.PK_RANGE_ID), + partition_key_range_id=None, + options={}, + ) + + self.assertEqual( + token, "{0}:{1}".format(self.PK_RANGE_ID, self.RAW_TOKEN), + "Expected a non-empty session token when partitionKey lacks 'version'." + ) + + def test_partitionkey_definition_with_version_returns_same_token(self): + """When 'version' IS present, behavior must be unchanged.""" + container = self._build_container_with_seeded_token() + cache = self._build_cache({"paths": ["/pk"], "kind": "Hash", "version": 2}) + + token = container.get_session_token( + resource_path=self.COLLECTION_LINK, + pk_value=self.PK_VALUE, + container_properties_cache=cache, + routing_map_provider=self._CapturingRoutingMapProvider(self.PK_RANGE_ID), + partition_key_range_id=None, + options={}, + ) + + self.assertEqual( + token, "{0}:{1}".format(self.PK_RANGE_ID, self.RAW_TOKEN), + "Expected an unchanged session token when partitionKey includes 'version'." + ) + + def test_missing_version_produces_v1_hash_epk_not_raw_binary(self): + """Missing 'version' must default to V1 hashing (matching the SDK's + _get_partition_key_from_partition_key_definition convention). If + version resolves to None, PartitionKey.__init__ stores None + (kwargs.get('version', V2) returns None when the key is present) and + _get_hashed_partition_key_string matches neither V1 nor V2 — the EPK + becomes the raw-binary encoding and targets the wrong physical + partition on multi-partition Hash containers.""" + from azure.cosmos.partition_key import PartitionKey + expected_epk = PartitionKey( + path=["/pk"], kind="Hash", version=1 + )._get_epk_range_for_partition_key(pk_value=self.PK_VALUE) + + container = self._build_container_with_seeded_token() + cache = self._build_cache({"paths": ["/pk"], "kind": "Hash"}) # no 'version' + capturing = self._CapturingRoutingMapProvider(self.PK_RANGE_ID) + + container.get_session_token( + resource_path=self.COLLECTION_LINK, + pk_value=self.PK_VALUE, + container_properties_cache=cache, + routing_map_provider=capturing, + partition_key_range_id=None, + options={}, + ) + + self.assertIsNotNone( + capturing.captured_ranges, + "get_overlapping_ranges was never called by get_session_token.", + ) + self.assertEqual(len(capturing.captured_ranges), 1) + actual_epk = capturing.captured_ranges[0] + self.assertEqual( + actual_epk.min, expected_epk.min, + "Missing 'version' must default to V1 hashing. If the actual EPK " + "differs from the expected V1 hash, the fix regressed to " + "version=None and produced the raw-binary encoding instead.", + ) + self.assertEqual(actual_epk.max, expected_epk.max) + + def test_v1_and_v2_produce_different_effective_partition_keys(self): + """Sanity check: V1 and V2 hashing produce distinct EPKs for the same + pk_value. If this ever regresses to equal, the V1-EPK assertion in + test_missing_version_produces_v1_hash_epk_not_raw_binary becomes vacuous.""" + from azure.cosmos.partition_key import PartitionKey + v1_epk = PartitionKey( + path=["/pk"], kind="Hash", version=1 + )._get_epk_range_for_partition_key(pk_value=self.PK_VALUE) + v2_epk = PartitionKey( + path=["/pk"], kind="Hash", version=2 + )._get_epk_range_for_partition_key(pk_value=self.PK_VALUE) + self.assertNotEqual( + v1_epk.min, v2_epk.min, + "V1 and V2 must produce distinct EPKs — if identical, the " + "'version' parameter no longer influences hashing.", + ) + + # Unit tests for set_session_token_header. When a request targets a single # partition, the helper must send only that partition's token, not a # comma-joined token covering every cached partition.