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

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

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

* Apply suggestions from code review

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

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

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

* Update sdk/cosmos/azure-cosmos/CHANGELOG.md

Co-authored-by: Simon Moreno <30335873+simorenoh@users.noreply.github.com>

* Rename test to spell out effective partition keys (fix cspell)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Re-trigger CI (emulator collection-count flake)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: REDMOND\gaausfel <gaausfel@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Fabian Meiswinkel <fabianm@microsoft.com>
Co-authored-by: Nalu Tripician <27316859+NaluTripician@users.noreply.github.com>
Co-authored-by: Simon Moreno <30335873+simorenoh@users.noreply.github.com>

Author: 6 people

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -8,6 +8,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
 

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],

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.