From 45d270a9284676e7b72e9ec6477e9a18e871c1c9 Mon Sep 17 00:00:00 2001 From: "REDMOND\\gaausfel" Date: Mon, 29 Jun 2026 15:54:05 -0700 Subject: [PATCH 1/8] [Cosmos] Fix KeyError: 'version' in SessionContainer.get_session_token The sync and async `SessionContainer.get_session_token` accessed `collection_pk_definition['version']` with bracket notation. When the service-returned `partitionKey` payload omits the optional `version` field (common for containers created without an explicit V2 partition key), this raised `KeyError`, which was silently swallowed by the surrounding `except (KeyError, AttributeError)`. The net effect was that `get_session_token` returned an empty string, so the client sent no `x-ms-session-token` header on the next read. Against the Dedicated Gateway / Integrated Cache, every Session-consistency read was rejected as a cache miss (verified at the server with `ComputeRequest5M.CacheHitLevel`: 695 reads / 0 CacheFullHit pre-fix vs 678 reads / 677 CacheFullHit post-fix on the same item). Treating `version` as optional and passing `None` to `PartitionKey` when it is missing matches the existing `PartitionKey` constructor contract and restores Session-consistency cache hits. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos/azure/cosmos/_session.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index b8b1c451ad10..96249aca4357 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 `None`, matching how `PartitionKey` handles a missing version. #### Other Changes diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session.py index bb1229b57662..4d25e5e15ce6 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')) 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')) 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], From f26dc0258306a3ab44a20007401a3c0243bcec8e Mon Sep 17 00:00:00 2001 From: "REDMOND\\gaausfel" Date: Tue, 30 Jun 2026 08:29:05 -0700 Subject: [PATCH 2/8] [Cosmos] Add regression tests for partitionKey without 'version' Covers SessionContainer.get_session_token in tests/test_session_token_unit.py: test_partitionkey_definition_without_version_returns_token Mirrors the live failure: container_properties_cache holds a partitionKey definition without 'version' (as the service returns it for containers created without explicit V2 partition-key versioning). Asserts a real per-partition token is returned (was '' before the fix, because the KeyError was silently swallowed). test_partitionkey_definition_with_version_returns_same_token Pins the no-regression case: when 'version' IS present, behavior is unchanged. Verified by reverting the fix locally and observing that the without_version test fails with AssertionError: '' != '0:1#100'. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../tests/test_session_token_unit.py | 85 +++++++++++++++++++ 1 file changed, 85 insertions(+) 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..a9cfb3284dc8 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,91 @@ 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 _StubRoutingMapProvider: + """Returns a single partition range covering any EPK.""" + + def __init__(self, pk_range_id): + self._pk_range_id = pk_range_id + + def get_overlapping_ranges(self, collection_link, ranges, options): + del collection_link, ranges, options + 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._StubRoutingMapProvider(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._StubRoutingMapProvider(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'." + ) + + # 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. From 6b6f3196f0688ac67dfebaaeeb9ebff8f2aae753 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 2 Jul 2026 00:52:24 +0200 Subject: [PATCH 3/8] Apply suggestions from code review Co-authored-by: Nalu Tripician <27316859+NaluTripician@users.noreply.github.com> --- sdk/cosmos/azure-cosmos/azure/cosmos/_session.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_session.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_session.py index 4d25e5e15ce6..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.get('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.get('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], From 01141cecc9735c3168ccbab7f78f6660fa2ba570 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Thu, 2 Jul 2026 00:53:34 +0200 Subject: [PATCH 4/8] Fix KeyError in SessionContainer for optional version Updated the default value of `partitionKey.version` to 1 in `SessionContainer.get_session_token` to match `PartitionKey` behavior. --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 96249aca4357..1b97d394ec25 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -7,7 +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 `None`, matching how `PartitionKey` handles a missing version. +* 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. #### Other Changes From e427dbe1c055f6a7470bb6faa10bc85e3ba7d4d0 Mon Sep 17 00:00:00 2001 From: "REDMOND\\gaausfel" Date: Wed, 1 Jul 2026 17:17:37 -0700 Subject: [PATCH 5/8] Strengthen regression tests to assert V1 EPK The prior tests used a routing-map stub that discarded the ranges argument and returned a fixed partition range regardless of the computed effective partition key, so a wrong-EPK regression (e.g. version=None falling through to the raw-binary encoding) would still have passed the assertions on the returned session token. Replace the stub with a capturing variant that records what get_overlapping_ranges received, and add two tests: test_missing_version_produces_v1_hash_epk_not_raw_binary Computes the V1-hash EPK independently and asserts the captured EPK matches. Fails if the fix ever regresses to version=None (the actual EPK would be the raw-binary encoding). test_v1_and_v2_produce_different_epks Sanity check: V1 and V2 must produce distinct EPKs, otherwise the V1-EPK assertion above becomes vacuous. Verified by temporarily reverting to .get('version') (no default) and observing the new EPK test fails with AssertionError comparing the raw-binary encoding against the expected V1 hash. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../tests/test_session_token_unit.py | 71 +++++++++++++++++-- 1 file changed, 66 insertions(+), 5 deletions(-) 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 a9cfb3284dc8..390fdc4ce6a7 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_session_token_unit.py +++ b/sdk/cosmos/azure-cosmos/tests/test_session_token_unit.py @@ -329,14 +329,18 @@ class TestGetSessionTokenWithoutPartitionKeyVersion(unittest.TestCase): PK_VALUE = "some-pk-value" RAW_TOKEN = "1#100" - class _StubRoutingMapProvider: - """Returns a single partition range covering any EPK.""" + 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, ranges, options + del collection_link, options + self.captured_ranges = list(ranges) return [{ "id": self._pk_range_id, "minInclusive": "", @@ -376,7 +380,7 @@ def test_partitionkey_definition_without_version_returns_token(self): resource_path=self.COLLECTION_LINK, pk_value=self.PK_VALUE, container_properties_cache=cache, - routing_map_provider=self._StubRoutingMapProvider(self.PK_RANGE_ID), + routing_map_provider=self._CapturingRoutingMapProvider(self.PK_RANGE_ID), partition_key_range_id=None, options={}, ) @@ -395,7 +399,7 @@ def test_partitionkey_definition_with_version_returns_same_token(self): resource_path=self.COLLECTION_LINK, pk_value=self.PK_VALUE, container_properties_cache=cache, - routing_map_provider=self._StubRoutingMapProvider(self.PK_RANGE_ID), + routing_map_provider=self._CapturingRoutingMapProvider(self.PK_RANGE_ID), partition_key_range_id=None, options={}, ) @@ -405,6 +409,63 @@ def test_partitionkey_definition_with_version_returns_same_token(self): "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_epks(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 From 4f1a86bd232d5760731e7b5ce9b201f0f81cf227 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 7 Jul 2026 18:56:07 +0200 Subject: [PATCH 6/8] Update sdk/cosmos/azure-cosmos/CHANGELOG.md Co-authored-by: Simon Moreno <30335873+simorenoh@users.noreply.github.com> --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 1b97d394ec25..371c5358d893 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -7,7 +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. +* 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 From b6c7f21f7a1b6701caead8e590c37f19f6f3ad0a Mon Sep 17 00:00:00 2001 From: "REDMOND\\gaausfel" Date: Tue, 7 Jul 2026 10:37:15 -0700 Subject: [PATCH 7/8] Rename test to spell out effective partition keys (fix cspell) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sdk/cosmos/azure-cosmos/tests/test_session_token_unit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 390fdc4ce6a7..0c5de8578251 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_session_token_unit.py +++ b/sdk/cosmos/azure-cosmos/tests/test_session_token_unit.py @@ -449,7 +449,7 @@ def test_missing_version_produces_v1_hash_epk_not_raw_binary(self): ) self.assertEqual(actual_epk.max, expected_epk.max) - def test_v1_and_v2_produce_different_epks(self): + 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 a398ccf8a22fd772e31c91a43b2acd4b53c69cc4 Mon Sep 17 00:00:00 2001 From: "REDMOND\\gaausfel" Date: Thu, 9 Jul 2026 07:51:30 -0700 Subject: [PATCH 8/8] Re-trigger CI (emulator collection-count flake) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>