Session Token Management APIs

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -10,7 +10,8 @@ This version and all future versions will support Python 3.13.
 * Added option to disable write payload on writes. See [PR 37365](https://github.com/Azure/azure-sdk-for-python/pull/37365)
 * Added get feed ranges API. See [PR 37687](https://github.com/Azure/azure-sdk-for-python/pull/37687)
 * Added feed range support in `query_items_change_feed`. See [PR 37687](https://github.com/Azure/azure-sdk-for-python/pull/37687)
-* Added **preview** helper APIs for managing session tokens. See [PR 36971](https://github.com/Azure/azure-sdk-for-python/pull/36971)
+* Added **provisional** helper APIs for managing session tokens. See [PR 36971](https://github.com/Azure/azure-sdk-for-python/pull/36971)
+* Added ability to get feed range for a partition key. See [PR 36971](https://github.com/Azure/azure-sdk-for-python/pull/36971)
 
 #### Breaking Changes
 * Item-level point operations will now return `CosmosDict` and `CosmosList` response types. 

sdk/cosmos/azure-cosmos/azure/cosmos/_change_feed/change_feed_state.py

@@ -383,9 +383,7 @@ def from_initial_state(
 
         feed_range: Optional[FeedRangeInternal] = None
         if change_feed_state_context.get("feedRange"):
-            feed_range_str = base64.b64decode(change_feed_state_context["feedRange"]).decode('utf-8')
-            feed_range_json = json.loads(feed_range_str)
-            feed_range = FeedRangeInternalEpk.from_json(feed_range_json)
+            feed_range = FeedRangeInternalEpk.from_json(change_feed_state_context["feedRange"])
         elif change_feed_state_context.get("partitionKey"):
             if change_feed_state_context.get("partitionKeyFeedRange"):
                 feed_range =\

sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py

@@ -21,9 +21,7 @@
 
 """Internal Helper functions for manipulating session tokens.
 """
-import base64
-import json
-from typing import Tuple, List
+from typing import Tuple, List, Dict, Any
 
 from azure.cosmos._routing.routing_range import Range
 from azure.cosmos._vector_session_token import VectorSessionToken
@@ -165,18 +163,15 @@ def merge_ranges_with_subsets(overlapping_ranges: List[Tuple[Range, str]]) -> Li
         overlapping_ranges.remove(overlapping_ranges[0])
     return processed_ranges
 
-def get_latest_session_token(feed_ranges_to_session_tokens: List[Tuple[str, str]], target_feed_range: str):
+def get_latest_session_token(feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]],
+                             target_feed_range: Dict[str, Any]):
 
-    target_feed_range_str = base64.b64decode(target_feed_range).decode('utf-8')
-    feed_range_json = json.loads(target_feed_range_str)
-    target_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json)
+    target_feed_range_epk = FeedRangeInternalEpk.from_json(target_feed_range)
     target_feed_range_normalized = target_feed_range_epk.get_normalized_range()
     # filter out tuples that overlap with target_feed_range and normalizes all the ranges
     overlapping_ranges = []
     for feed_range_to_session_token in feed_ranges_to_session_tokens:
-        feed_range_str = base64.b64decode(feed_range_to_session_token[0]).decode('utf-8')
-        feed_range_json = json.loads(feed_range_str)
-        feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json)
+        feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_to_session_token[0])
         if Range.overlaps(target_feed_range_normalized,
                           feed_range_epk.get_normalized_range()):
             overlapping_ranges.append((feed_range_epk.get_normalized_range(),

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py

@@ -21,11 +21,9 @@
 
 """Create, read, update and delete items in the Azure Cosmos DB SQL API service.
 """
-import base64
-import json
 import warnings
 from datetime import datetime
-from typing import Any, Dict, Mapping, Optional, Sequence, Type, Union, List, Tuple, cast, overload
+from typing import Any, Dict, Mapping, Optional, Sequence, Type, Union, List, Tuple, cast, overload, Iterable
 from typing_extensions import Literal
 
 from azure.core import MatchConditions
@@ -537,15 +535,15 @@ def query_items_change_feed(
     def query_items_change_feed(
             self,
             *,
-            feed_range: str,
+            feed_range: Dict[str, Any],
             max_item_count: Optional[int] = None,
             start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None,
             priority: Optional[Literal["High", "Low"]] = None,
             **kwargs: Any
     ) -> AsyncItemPaged[Dict[str, Any]]:
         """Get a sorted list of items that were changed, in the order in which they were modified.
 
-        :keyword str feed_range: The feed range that is used to define the scope.
+        :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope.
         :keyword int max_item_count: Max number of items to be returned in the enumeration operation.
         :keyword start_time: The start time to start processing chang feed items.
             Beginning: Processing the change feed items from the beginning of the change feed.
@@ -623,7 +621,7 @@ def query_items_change_feed( # pylint: disable=unused-argument
         """Get a sorted list of items that were changed, in the order in which they were modified.
 
         :keyword str continuation: The continuation token retrieved from previous response.
-        :keyword str feed_range: The feed range that is used to define the scope.
+        :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope.
         :keyword partition_key: The partition key that is used to define the scope
             (logical partition or a subset of a container)
         :type partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]]
@@ -1299,13 +1297,17 @@ async def read_feed_ranges(
             *,
             force_refresh: Optional[bool] = False,
             **kwargs: Any
-    ) -> List[str]:
+    ) -> Iterable[Dict[str, Any]]:
         """ Obtains a list of feed ranges that can be used to parallelize feed operations.
 
         :keyword bool force_refresh:
             Flag to indicate whether obtain the list of feed ranges directly from cache or refresh the cache.
         :returns: A list representing the feed ranges in base64 encoded string
-        :rtype: List[str]
+        :rtype: Iterable[Dict[str, Any]]
+
+        .. note::
+          For each feed range, even through a Dict has been returned,
+          but in the future, the structure may change. Please just treat it as opaque and do not take any dependent on it.
 
         """
         if force_refresh is True:
@@ -1318,50 +1320,58 @@ async def read_feed_ranges(
                 [Range("", "FF", True, False)],
                 **kwargs)
 
-        return [FeedRangeInternalEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)).__str__()
+        feed_ranges = [FeedRangeInternalEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)).to_dict()
                 for partitionKeyRange in partition_key_ranges]
+        return (feed_range for feed_range in feed_ranges)
 
     async def get_latest_session_token(self,
-                                       feed_ranges_to_session_tokens: List[Tuple[str, str]],
-                                       target_feed_range: str
+                                       feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]],
+                                       target_feed_range: Dict[str, Any]
                                        ) -> str:
-        """Gets the the most up to date session token from the list of session token and feed range tuples
+        """ **provisional** Gets the the most up to date session token from the list of session token and feed range tuples
         for a specific target feed range. The feed range can be obtained from a logical partition or by reading the
         container feed ranges. This should only be used if maintaining own session token or else the sdk will
         keep track of session token. Session tokens and feed ranges are scoped to a container. Only input session
         tokens and feed ranges obtained from the same container.
         :param feed_ranges_to_session_tokens: List of partition key and session token tuples.
-        :type feed_ranges_to_session_tokens: List[Tuple[str, FeedRange]]
+        :type feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]]
         :param target_feed_range: feed range to get most up to date session token.
-        :type target_feed_range: FeedRange
+        :type target_feed_range: Dict[str, Any]
         :returns: a session token
         :rtype: str
         """
         return get_latest_session_token(feed_ranges_to_session_tokens, target_feed_range)
 
-    async def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> str:
+    async def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> Dict[str, Any]:
         """Gets the feed range for a given partition key.
         :param partition_key: partition key to get feed range.
         :type partition_key: PartitionKey
         :returns: a feed range
-        :rtype: str
+        :rtype: Dict[str, Any]
+
+        .. note::
+          For the feed range, even through a Dict has been returned, but in the future,
+          the structure may change. Please just treat it as opaque and do not take any dependence on it.
+
         """
-        return FeedRangeInternalEpk(await self._get_epk_range_for_partition_key(partition_key)).__str__()
+        return FeedRangeInternalEpk(await self._get_epk_range_for_partition_key(partition_key)).to_dict()
 
-    async def is_feed_range_subset(self, parent_feed_range: str, child_feed_range: str) -> bool:
+    async def is_feed_range_subset(self, parent_feed_range: Dict[str, Any],
+                                   child_feed_range: Dict[str, Any]) -> bool:
         """Checks if child feed range is a subset of parent feed range.
         :param parent_feed_range: left feed range
-        :type parent_feed_range: str
+        :type parent_feed_range: Dict[str, Any]
         :param child_feed_range: right feed range
-        :type child_feed_range: str
+        :type child_feed_range: Dict[str, Any]
         :returns: a boolean indicating if child feed range is a subset of parent feed range
         :rtype: bool
+
+        .. note::
+          For the feed range, even through a Dict has been returned, but in the future,
+          the structure may change. Please just treat it as opaque and do not take any dependence on it.
+
         """
-        parent_feed_range_str = base64.b64decode(parent_feed_range).decode('utf-8')
-        feed_range_json = json.loads(parent_feed_range_str)
-        parent_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json)
-        child_feed_range_str = base64.b64decode(child_feed_range).decode('utf-8')
-        feed_range_json = json.loads(child_feed_range_str)
-        child_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json)
+        parent_feed_range_epk = FeedRangeInternalEpk.from_json(parent_feed_range)
+        child_feed_range_epk = FeedRangeInternalEpk.from_json(child_feed_range)
         return child_feed_range_epk.get_normalized_range().is_subset(
             parent_feed_range_epk.get_normalized_range())

sdk/cosmos/azure-cosmos/azure/cosmos/container.py

@@ -21,11 +21,9 @@
 
 """Create, read, update and delete items in the Azure Cosmos DB SQL API service.
 """
-import base64
-import json
 import warnings
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Sequence, Union, Tuple, Mapping, Type, cast, overload
+from typing import Any, Dict, List, Optional, Sequence, Union, Tuple, Mapping, Type, cast, overload, Iterable
 from typing_extensions import Literal
 
 from azure.core import MatchConditions
@@ -357,7 +355,7 @@ def query_items_change_feed(
     def query_items_change_feed(
             self,
             *,
-            feed_range: str,
+            feed_range: Dict[str, Any],
             max_item_count: Optional[int] = None,
             start_time: Optional[Union[datetime, Literal["Now", "Beginning"]]] = None,
             priority: Optional[Literal["High", "Low"]] = None,
@@ -366,7 +364,7 @@ def query_items_change_feed(
 
         """Get a sorted list of items that were changed, in the order in which they were modified.
 
-        :keyword str feed_range: The feed range that is used to define the scope.
+        :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope.
         :keyword int max_item_count: Max number of items to be returned in the enumeration operation.
         :keyword start_time: The start time to start processing chang feed items.
             Beginning: Processing the change feed items from the beginning of the change feed.
@@ -443,7 +441,7 @@ def query_items_change_feed(
         """Get a sorted list of items that were changed, in the order in which they were modified.
 
         :keyword str continuation: The continuation token retrieved from previous response.
-        :keyword str feed_range: The feed range that is used to define the scope.
+        :keyword Dict[str, Any] feed_range: The feed range that is used to define the scope.
         :keyword partition_key: The partition key that is used to define the scope
             (logical partition or a subset of a container)
         :type partition_key: Union[str, int, float, bool, List[Union[str, int, float, bool]]]
@@ -1368,14 +1366,18 @@ def read_feed_ranges(
             self,
             *,
             force_refresh: Optional[bool] = False,
-            **kwargs: Any) -> List[str]:
+            **kwargs: Any) -> Iterable[Dict[str, Any]]:
 
         """ Obtains a list of feed ranges that can be used to parallelize feed operations.
 
         :keyword bool force_refresh:
             Flag to indicate whether obtain the list of feed ranges directly from cache or refresh the cache.
         :returns: A list representing the feed ranges in base64 encoded string
-        :rtype: List[str]
+        :rtype: Iterable[Dict[str, Any]]
+
+        .. note::
+          For each feed range, even through a Dict has been returned,
+          but in the future, the structure may change. Please just treat it as opaque and do not take any dependent on it.
 
         """
         if force_refresh is True:
@@ -1387,50 +1389,55 @@ def read_feed_ranges(
                 [Range("", "FF", True, False)], # default to full range
                 **kwargs)
 
-        return [FeedRangeInternalEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)).__str__()
+        feed_ranges = [FeedRangeInternalEpk(Range.PartitionKeyRangeToRange(partitionKeyRange)).to_dict()
                 for partitionKeyRange in partition_key_ranges]
+        return (feed_range for feed_range in feed_ranges)
 
     def get_latest_session_token(
             self,
-            feed_ranges_to_session_tokens: List[Tuple[str, str]],
-            target_feed_range: str) -> str:
-        """Gets the the most up to date session token from the list of session token and feed range tuples
+            feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]],
+            target_feed_range: Dict[str, Any]) -> str:
+        """ **provisional** Gets the the most up to date session token from the list of session token and feed range tuples
         for a specific target feed range. The feed range can be obtained from a logical partition or by reading the
         container feed ranges. This should only be used if maintaining own session token or else the sdk will
         keep track of session token. Session tokens and feed ranges are scoped to a container. Only input session
         tokens and feed ranges obtained from the same container.
         :param feed_ranges_to_session_tokens: List of partition key and session token tuples.
-        :type feed_ranges_to_session_tokens: List[Tuple[str, str]]
+        :type feed_ranges_to_session_tokens: List[Tuple[Dict[str, Any], str]]
         :param target_feed_range: feed range to get most up to date session token.
-        :type target_feed_range: str
+        :type target_feed_range: Dict[str, Any]
         :returns: a session token
         :rtype: str
         """
         return get_latest_session_token(feed_ranges_to_session_tokens, target_feed_range)
 
-    def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> str:
-        """Gets the feed range for a given partition key.
+    def feed_range_from_partition_key(self, partition_key: PartitionKeyType) -> Dict[str, Any]:
+        """ Gets the feed range for a given partition key.
         :param partition_key: partition key to get feed range.
         :type partition_key: PartitionKey
         :returns: a feed range
-        :rtype: str
+        :rtype: Dict[str, Any]
+        .. note::
+          For the feed range, even through a Dict has been returned, but in the future,
+          the structure may change. Please just treat it as opaque and do not take any dependence on it.
+
         """
-        return FeedRangeInternalEpk(self._get_epk_range_for_partition_key(partition_key)).__str__()
+        return FeedRangeInternalEpk(self._get_epk_range_for_partition_key(partition_key)).to_dict()
 
-    def is_feed_range_subset(self, parent_feed_range: str, child_feed_range: str) -> bool:
-        """Checks if child feed range is a subset of parent feed range.
+    def is_feed_range_subset(self, parent_feed_range: Dict[str, Any], child_feed_range: Dict[str, Any]) -> bool:
+        """ Checks if child feed range is a subset of parent feed range.
         :param parent_feed_range: left feed range
-        :type parent_feed_range: str
+        :type parent_feed_range: Dict[str, Any]
         :param child_feed_range: right feed range
-        :type child_feed_range: str
+        :type child_feed_range: Dict[str, Any]
         :returns: a boolean indicating if child feed range is a subset of parent feed range
         :rtype: bool
+
+        .. note::
+          For the feed range, even through a Dict has been returned, but in the future,
+          the structure may change. Please just treat it as opaque and do not take any dependence on it.
         """
-        parent_feed_range_str = base64.b64decode(parent_feed_range).decode('utf-8')
-        feed_range_json = json.loads(parent_feed_range_str)
-        parent_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json)
-        child_feed_range_str = base64.b64decode(child_feed_range).decode('utf-8')
-        feed_range_json = json.loads(child_feed_range_str)
-        child_feed_range_epk = FeedRangeInternalEpk.from_json(feed_range_json)
+        parent_feed_range_epk = FeedRangeInternalEpk.from_json(parent_feed_range)
+        child_feed_range_epk = FeedRangeInternalEpk.from_json(child_feed_range)
         return child_feed_range_epk.get_normalized_range().is_subset(
             parent_feed_range_epk.get_normalized_range())

sdk/cosmos/azure-cosmos/samples/examples.py

@@ -11,6 +11,7 @@
 # All interaction with Cosmos DB starts with an instance of the CosmosClient
 # [START create_client]
 from azure.cosmos import exceptions, CosmosClient, PartitionKey
+from typing import Dict, Any
 
 import os
 
@@ -259,12 +260,25 @@
 
 # Get the feed ranges list from container.
 # [START read_feed_ranges]
-container.read_feed_ranges()
+feed_ranges = list(container.read_feed_ranges())
 # [END read_feed_ranges]
 
+# Get a feed range from a partition key.
+# [START feed_range_from_partition_key ]
+feed_range_from_pk = container.feed_range_from_partition_key(["GA", "Atlanta", 30363])
+# [END feed_range_from_partition_key]
+
+# Figure out if a feed range is a subset of another feed range.
+# This example sees in which feed range from the container a feed range from a partition key is part of.
+# [START is_feed_range_subset]
+parent_feed_range: Dict[str, Any] = next(
+    (feed_range for feed_range in feed_ranges if container.is_feed_range_subset(feed_range, feed_range_from_pk)),
+    {}
+)
+# [END is_feed_range_subset]
+
 # Query a sorted list of items that were changed for one feed range
 # [START query_items_change_feed]
-feed_ranges = container.read_feed_ranges()
 for item in container.query_items_change_feed(feed_range=feed_ranges[0]):
     print(json.dumps(item, indent=True))
 # [END query_items_change_feed]