Fix KeyError: 'version' in SessionContainer.get_session_token by ausfeldt · Pull Request #47798 · Azure/azure-sdk-for-python

ausfeldt · 2026-07-01T21:00:33Z

Description

The service-returned "partitionKey" definition does not always include the optional version field. When it does not, "collection_pk_definition['version']" raises "KeyError", which the surrounding "except (KeyError, AttributeError)" silently swallows. The function returns`"" and the client omits the "x-ms-session-token" header on the next read. Against the Dedicated Gateway / Integrated Cache, every Session-consistency read is then rejected as "CacheMiss" even against a warm cache for the same document.

Two-line change in sdk/cosmos/azure-cosmos/azure/cosmos/_session.py (sync at line 122 and async at line 216):
"version=collection_pk_definition['version']" to "version=collection_pk_definition.get('version')".

"PartitionKey" already accepts "version=None" and defaults to V2 internally, so the semantic behavior is unchanged when "version" is present.

All SDK Contribution checklist:

The pull request does not introduce [breaking changes]
CHANGELOG is updated for new features, bug fixes or other significant changes.
I have read the contribution guidelines.

General Guidelines and Best Practices

Title of the pull request is clear and informative.
There are a small number of commits, each of which have an informative message. This means that previously merged commits do not appear in the history of the PR. For more information on cleaning up the commits in your PR, see this page.

Testing Guidelines

Pull request includes test coverage for the included changes.

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>

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>

Copilot

Pull request overview

This PR intends to fix a KeyError: 'version' in SessionContainer.get_session_token (sync and async) that occurs when the service-returned partitionKey definition omits the optional version field. The KeyError was being silently swallowed by a broad except, causing the client to omit the x-ms-session-token header on subsequent reads and turning every Session-consistency read against the Dedicated Gateway / Integrated Cache into a cache miss. The change replaces collection_pk_definition['version'] with collection_pk_definition.get('version') in both code paths, plus a CHANGELOG entry and new unit tests.

Changes:

Use .get('version') instead of ['version'] when constructing PartitionKey in the sync and async get_session_token paths.
Add a CHANGELOG.md bug-fix entry.
Add unit tests covering the missing/present partitionKey.version cases for the sync path.

Key concern: The chosen fix avoids the KeyError but does not correctly restore behavior. Passing version=None (what .get('version') returns when the key is missing) is not the same as omitting the keyword: PartitionKey.__init__ only applies its V2 default when version is absent, so None is stored and the effective-partition-key hashing falls through to a non-hash encoding, producing an incorrect EPK. A concrete default (e.g. 1, matching the existing _get_partition_key_from_partition_key_definition helper) is needed. The new tests don't catch this because the routing-map stub ignores the computed EPK.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
`sdk/cosmos/azure-cosmos/azure/cosmos/_session.py`	Changes `version=collection_pk_definition['version']` to `.get('version')` in sync (L122) and async (L216) paths; avoids `KeyError` but passes `None`, yielding an incorrect EPK instead of the V2 default.
`sdk/cosmos/azure-cosmos/tests/test_session_token_unit.py`	Adds sync-only tests for missing/present `partitionKey.version`; the routing-map stub discards the computed EPK ranges, so it doesn't validate the fix's core resolution behavior.
`sdk/cosmos/azure-cosmos/CHANGELOG.md`	Adds bug-fix entry; wording ("defaults to `None`, matching how `PartitionKey` handles a missing version") is inaccurate.

FabianMeiswinkel · 2026-07-01T22:50:38Z

                        partition_key = PartitionKey(path=collection_pk_definition['paths'],
                                                     kind=collection_pk_definition['kind'],
-                                                     version=collection_pk_definition['version'])
+                                                     version=collection_pk_definition.get('version'))


This change makes sense to me

                        partition_key = PartitionKey(path=collection_pk_definition['paths'],
                                                     kind=collection_pk_definition['kind'],
-                                                     version=collection_pk_definition['version'])
+                                                     version=collection_pk_definition.get('version'))


+        def get_overlapping_ranges(self, collection_link, ranges, options):
+            del collection_link, ranges, options
+            return [{
+                "id": self._pk_range_id,
+                "minInclusive": "",
+                "maxExclusive": "FF",
+                "parents": [],
+            }]


 #### 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.


NaluTripician

Thanks for the thorough root-cause writeup — the diagnosis (swallowed KeyError -> empty session token -> Integrated Cache misses against the Dedicated Gateway) is spot-on and worth fixing.

I did dig into the actual partition_key.py code paths, though, and I think the fix as written trades the loud failure for a quieter correctness bug. The premise in the description — "PartitionKey accepts version=None and defaults to V2 internally" — doesn't hold: passing version=None explicitly keeps self.version = None, and for a Hash container that falls through _get_hashed_partition_key_string to the raw-binary EPK path rather than V1/V2 hashing. It's harmless on single-partition containers (which the new tests cover) but can target the wrong partition — or produce an empty token — on multi-partition Hash containers.

The one-line change to .get('version', 1) (matching the existing _get_partition_key_from_partition_key_definition convention), plus a test that exercises real EPK->range resolution, should make this solid. Details inline. Defaulting to 1 also matches .NET/Java (absent version => V1 for legacy containers), whereas None is a Python-only divergence.

NaluTripician · 2026-07-01T21:20:40Z

 #### 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.


This wording isn't accurate: PartitionKey does not treat a missing version as None — its constructor default is V2, and the SDK's own convention for a service definition missing version is to default to 1 (see _get_partition_key_from_partition_key_definition in partition_key.py). Passing None actually makes a Hash container's EPK fall through to the raw-binary path. Once the fix defaults to 1, please reword, e.g. "...defaults to version 1, consistent with _get_partition_key_from_partition_key_definition."

NaluTripician · 2026-07-01T21:20:40Z

+            self._pk_range_id = pk_range_id
+
+        def get_overlapping_ranges(self, collection_link, ranges, options):
+            del collection_link, ranges, options


This stub discards the ranges argument (del ... ranges ...) and always returns a single fixed range, so the computed EPK is never actually exercised. Both the with-version and without-version tests would pass even if the EPK were computed incorrectly — they only prove "no KeyError was raised," not "the correct partition was targeted." As written they'd stay green even with the wrong-EPK behavior from the .get('version') change.

Consider a test that does not stub out EPK->range resolution: build a realistic routing map with two ranges and assert that a missing version resolves to the same range/token as an explicit version=1. At minimum, assert the EPK produced when version is omitted equals the V1-hash EPK (not the hex-binary string). That would catch the issue flagged on _session.py.

Co-authored-by: Nalu Tripician <27316859+NaluTripician@users.noreply.github.com>

Updated the default value of `partitionKey.version` to 1 in `SessionContainer.get_session_token` to match `PartitionKey` behavior.

FabianMeiswinkel

LGTM

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>

tvaron3 · 2026-07-02T05:21:18Z

@@ -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.


changelog entry should include pr link like the other entries in this file

simorenoh · 2026-07-02T14:36:16Z

@@ -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.


Suggested change

* 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)

REDMOND\gaausfel and others added 2 commits June 29, 2026 15:54

Copilot AI review requested due to automatic review settings July 1, 2026 21:00

ausfeldt requested a review from a team as a code owner July 1, 2026 21:00

github-actions Bot added the Cosmos label Jul 1, 2026

github-project-automation Bot added this to CosmosDB Python Eco-System Jul 1, 2026

Copilot started reviewing on behalf of ausfeldt July 1, 2026 21:01 View session

Copilot AI reviewed Jul 1, 2026

View reviewed changes

NaluTripician reviewed Jul 1, 2026

View reviewed changes

FabianMeiswinkel reviewed Jul 1, 2026

View reviewed changes

Comment thread sdk/cosmos/azure-cosmos/CHANGELOG.md Outdated

FabianMeiswinkel and others added 2 commits July 2, 2026 00:52

Apply suggestions from code review

6b6f319

Co-authored-by: Nalu Tripician <27316859+NaluTripician@users.noreply.github.com>

Fix KeyError in SessionContainer for optional version

01141ce

Updated the default value of `partitionKey.version` to 1 in `SessionContainer.get_session_token` to match `PartitionKey` behavior.

FabianMeiswinkel approved these changes Jul 1, 2026

View reviewed changes

tvaron3 reviewed Jul 2, 2026

View reviewed changes

tvaron3 approved these changes Jul 2, 2026

View reviewed changes

simorenoh reviewed Jul 2, 2026

View reviewed changes

simorenoh approved these changes Jul 2, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fix KeyError: 'version' in SessionContainer.get_session_token#47798

Fix KeyError: 'version' in SessionContainer.get_session_token#47798
ausfeldt wants to merge 5 commits into
Azure:mainfrom
ausfeldt:users/gaausfel/fix-session-version

ausfeldt commented Jul 1, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

FabianMeiswinkel Jul 1, 2026

Uh oh!

NaluTripician left a comment

Uh oh!

Uh oh!

Uh oh!

NaluTripician Jul 1, 2026

Uh oh!

NaluTripician Jul 1, 2026

Uh oh!

Uh oh!

FabianMeiswinkel left a comment

Uh oh!

tvaron3 Jul 2, 2026

Uh oh!

simorenoh Jul 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

ausfeldt commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

All SDK Contribution checklist:

General Guidelines and Best Practices

Testing Guidelines

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

FabianMeiswinkel Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

NaluTripician left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

NaluTripician Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

NaluTripician Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

FabianMeiswinkel left a comment

Choose a reason for hiding this comment

Uh oh!

tvaron3 Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

simorenoh Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ausfeldt commented Jul 1, 2026 •

edited

Loading