Fix KeyError: 'version' in SessionContainer.get_session_token#47798
Conversation
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>
There was a problem hiding this comment.
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 constructingPartitionKeyin the sync and asyncget_session_tokenpaths. - Add a
CHANGELOG.mdbug-fix entry. - Add unit tests covering the missing/present
partitionKey.versioncases 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. |
NaluTripician
left a comment
NaluTripician
left a comment
There was a problem hiding this comment.
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.
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.
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>
Co-authored-by: Simon Moreno <30335873+simorenoh@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
/check-enforcer override |
|
@microsoft-github-policy-service rerun |
Description
The service-returned "partitionKey" definition does not always include the optional
versionfield. 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:
General Guidelines and Best Practices
Testing Guidelines