Session Token Management APIs

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -8,6 +8,7 @@
 * 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 helper APIs for managing session tokens. See [PR 36971](https://github.com/Azure/azure-sdk-for-python/pull/36971)
 
 #### Bugs Fixed
 * Consolidated Container Properties Cache to be in the Client to cache partition key definition and container rid to avoid unnecessary container reads. See [PR 35731](https://github.com/Azure/azure-sdk-for-python/pull/35731).

sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py

@@ -186,7 +186,6 @@ def _compare_helper(a, b):
 
     @staticmethod
     def overlaps(range1, range2):
-
         if range1 is None or range2 is None:
             return False
         if range1.isEmpty() or range2.isEmpty():
@@ -195,10 +194,36 @@ def overlaps(range1, range2):
         cmp1 = Range._compare_helper(range1.min, range2.max)
         cmp2 = Range._compare_helper(range2.min, range1.max)
 
-        if cmp1 <= 0 or cmp2 <= 0:
+        if cmp1 <= 0 and cmp2 <= 0:
             if (cmp1 == 0 and not (range1.isMinInclusive and range2.isMaxInclusive)) or (
                 cmp2 == 0 and not (range2.isMinInclusive and range1.isMaxInclusive)
             ):
                 return False
             return True
         return False
+
+    def can_merge(self, other):
+        if self.isSingleValue() and other.isSingleValue():
+            return self.min == other.min
+        # if share the same boundary, they can merge
+        overlap_boundary1 = self.max == other.min and self.isMaxInclusive or other.isMinInclusive
+        overlap_boundary2 = other.max == self.min and other.isMaxInclusive or self.isMinInclusive
+        if overlap_boundary1 or overlap_boundary2:
+            return True
+        return self.overlaps(self, other)
+
+    def merge(self, other):
+        if not self.can_merge(other):
+            raise ValueError("Ranges do not overlap")
+        min_val = self.min if self.min < other.min else other.min
+        max_val = self.max if self.max > other.max else other.max
+        is_min_inclusive = self.isMinInclusive if self.min < other.min else other.isMinInclusive
+        is_max_inclusive = self.isMaxInclusive if self.max > other.max else other.isMaxInclusive
+        return Range(min_val, max_val, is_min_inclusive, is_max_inclusive)
+
+    def is_subset(self, parent_range) -> bool:
+        normalized_parent_range = parent_range.to_normalized_range()
+        normalized_child_range = self.to_normalized_range()
+        return normalized_parent_range.contains(normalized_child_range.min) and \
+            (normalized_parent_range.contains(normalized_child_range.max)
+             or normalized_parent_range.max == normalized_child_range.max)

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

@@ -0,0 +1,191 @@
+# The MIT License (MIT)
+# Copyright (c) 2014 Microsoft Corporation
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Internal Helper functions for manipulating session tokens.
+"""
+from azure.cosmos._routing.routing_range import Range
+from azure.cosmos._vector_session_token import VectorSessionToken
+
+# pylint: disable=protected-access
+
+def merge_session_tokens_with_same_range(session_token1, session_token2):
+    pk_range_id1, vector_session_token1 = parse_session_token(session_token1)
+    pk_range_id2, vector_session_token2 = parse_session_token(session_token2)
+    pk_range_id = pk_range_id1
+    if pk_range_id1 != pk_range_id2:
+        pk_range_id = pk_range_id1 \
+            if vector_session_token1.is_greater(vector_session_token2) else pk_range_id2
+    vector_session_token = vector_session_token1.merge(vector_session_token2)
+    return pk_range_id + ":" +  vector_session_token.session_token
+
+def is_compound_session_token(session_token):
+    return "," in session_token
+
+def parse_session_token(session_token):
+    tokens = session_token.split(":")
+    return tokens[0], VectorSessionToken.create(tokens[1])
+
+def split_compound_session_tokens(compound_session_tokens):
+    session_tokens = []
+    for _, session_token in compound_session_tokens:
+        if is_compound_session_token(session_token):
+            tokens = session_token.split(",")
+            for token in tokens:
+                session_tokens.append(token)
+        else:
+            session_tokens.append(session_token)
+    return session_tokens
+
+def merge_session_tokens_with_same_pkrangeid(session_tokens):
+    i = 0
+    while i < len(session_tokens):
+        j = i + 1
+        while j < len(session_tokens):
+            pk_range_id1, vector_session_token1 = parse_session_token(session_tokens[i])
+            pk_range_id2, vector_session_token2 = parse_session_token(session_tokens[j])
+            if pk_range_id1 == pk_range_id2:
+                vector_session_token = vector_session_token1.merge(vector_session_token2)
+                session_tokens.append(pk_range_id1 + ":" + vector_session_token.session_token)
+                remove_session_tokens = [session_tokens[i], session_tokens[j]]
+                for token in remove_session_tokens:
+                    session_tokens.remove(token)
+                i = -1
+                break
+            j += 1
+        i += 1
+
+    return session_tokens
+
+def merge_ranges_with_subsets(overlapping_ranges):
+    processed_ranges = []
+    while len(overlapping_ranges) != 0:
+        feed_range_cmp, session_token_cmp = overlapping_ranges[0]
+        # compound session tokens are not considered for merging
+        if is_compound_session_token(session_token_cmp):
+            processed_ranges.append(overlapping_ranges[0])
+            overlapping_ranges.remove(overlapping_ranges[0])
+            continue
+        _, vector_session_token_cmp = parse_session_token(session_token_cmp)
+        subsets = []
+        # finding the subset feed ranges of the current feed range
+        for j in range(1, len(overlapping_ranges)):
+            feed_range = overlapping_ranges[j][0]
+            if not is_compound_session_token(overlapping_ranges[j][1]) and \
+                    feed_range.is_subset(feed_range_cmp):
+                subsets.append(overlapping_ranges[j] + (j,))
+
+        # go through subsets to see if can create current feed range from the subsets
+        not_found = True
+        j = 0
+        while not_found and j < len(subsets):
+            merged_range = subsets[j][0]
+            session_tokens = [subsets[j][1]]
+            merged_indices = [subsets[j][2]]
+            if len(subsets) == 1:
+                _, vector_session_token = parse_session_token(session_tokens[0])
+                if vector_session_token_cmp.is_greater(vector_session_token):
+                    overlapping_ranges.remove(overlapping_ranges[merged_indices[0]])
+            else:
+                for k, subset in enumerate(subsets):
+                    if j == k:
+                        continue
+                    if merged_range.can_merge(subset[0]):
+                        merged_range = merged_range.merge(subset[0])
+                        session_tokens.append(subset[1])
+                        merged_indices.append(subset[2])
+                    if feed_range_cmp == merged_range:
+                        # if feed range can be created from the subsets
+                        # take the subsets if their global lsn is larger
+                        # else take the current feed range
+                        children_more_updated = True
+                        for session_token in session_tokens:
+                            _, vector_session_token = parse_session_token(session_token)
+                            if vector_session_token_cmp.is_greater(vector_session_token):
+                                children_more_updated = False
+                        feed_ranges_to_remove = [overlapping_ranges[i] for i in merged_indices]
+                        for feed_range_to_remove in feed_ranges_to_remove:
+                            overlapping_ranges.remove(feed_range_to_remove)
+                        if children_more_updated:
+                            overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens))))
+                            overlapping_ranges.remove(overlapping_ranges[0])
+                        not_found = False
+                        break
+
+            j += 1
+
+        processed_ranges.append(overlapping_ranges[0])
+        overlapping_ranges.remove(overlapping_ranges[0])
+    return processed_ranges
+
+def get_updated_session_token(feed_ranges_to_session_tokens, target_feed_range):
+    target_feed_range_normalized = target_feed_range._feed_range_internal.get_normalized_range()
+    # filter out tuples that overlap with target_feed_range and normalizes all the ranges
+    overlapping_ranges = [(feed_range_to_session_token[0]._feed_range_internal.get_normalized_range(),
+                           feed_range_to_session_token[1])
+                          for feed_range_to_session_token in feed_ranges_to_session_tokens if Range.overlaps(
+            target_feed_range_normalized, feed_range_to_session_token[0]._feed_range_internal.get_normalized_range())]
+    # Is there a feed_range that is a superset of some of the other feed_ranges excluding tuples
+    # with compound session tokens?
+    if len(overlapping_ranges) == 0:
+        raise ValueError('There were no overlapping feed ranges with the target.')
+
+    # merge any session tokens that are the same exact feed range
+    i = 0
+    j = 1
+    while i < len(overlapping_ranges) and j < len(overlapping_ranges):
+        cur_feed_range = overlapping_ranges[i][0]
+        session_token = overlapping_ranges[i][1]
+        session_token_1 = overlapping_ranges[j][1]
+        if (not is_compound_session_token(session_token) and
+                not is_compound_session_token(overlapping_ranges[j][1]) and
+                cur_feed_range == overlapping_ranges[j][0]):
+            session_token = merge_session_tokens_with_same_range(session_token, session_token_1)
+            feed_ranges_to_remove = [overlapping_ranges[i], overlapping_ranges[j]]
+            for feed_range_to_remove in feed_ranges_to_remove:
+                overlapping_ranges.remove(feed_range_to_remove)
+            overlapping_ranges.append((cur_feed_range, session_token))
+            i, j = 0, 1
+        else:
+            j += 1
+            if j == len(overlapping_ranges):
+                i += 1
+                j = i + 1
+
+    # checking for merging of feed ranges that can be created from other feed ranges
+    processed_ranges = merge_ranges_with_subsets(overlapping_ranges)
+
+    # break up session tokens that are compound
+    remaining_session_tokens = split_compound_session_tokens(processed_ranges)
+
+    if len(remaining_session_tokens) == 1:
+        return remaining_session_tokens[0]
+    # merging any session tokens with same physical partition key range id
+    remaining_session_tokens = merge_session_tokens_with_same_pkrangeid(remaining_session_tokens)
+
+    updated_session_token = ""
+    # compound the remaining session tokens
+    for i in range(len(remaining_session_tokens)):
+        if i == len(remaining_session_tokens) - 1:
+            updated_session_token += remaining_session_tokens[i]
+        else:
+            updated_session_token += remaining_session_tokens[i] + ","
+
+    return updated_session_token

sdk/cosmos/azure-cosmos/azure/cosmos/_vector_session_token.py

@@ -113,6 +113,11 @@ def equals(self, other):
             and self.are_region_progress_equal(other.local_lsn_by_region)
         )
 
+    def is_greater(self, other):
+        if self.global_lsn > other.global_lsn:
+            return True
+        return False
+
     def merge(self, other):
         if other is None:
             raise ValueError("Invalid Session Token (should not be None)")