Merging latest hotfix released back to main branch (Azure#39584)

* add new policy, add logic to use policy

* added small test file I was using

* initial regional endpoint work

* groundwork

* re-add AzureError logic, refactor, fix tests

* Update _retry_utility.py

* Updated location_cache with new design

* Fixed key error with most_preferred_location

* Update test_cosmos_http_logging_policy.py

* Update _retry_utility.py

* Added logic to refresh cache on previous endpoint usage

* Added business logic update the regional endpoint based on success or failures

* implementation

* Update _retry_utility_async.py

* fix some tests

* changelog, versions, fixes

* fixes

* fix some tests

* remove fake logic, count fix

* fix some tests

* Update _service_request_retry_policy.py

* Update _retry_utility_async.py

* retry utilities fixing

* Update _retry_utility.py

* additional enhancements

* Update setup.py

* Update _retry_utility_async.py

* add tests, remove previous retry logic for ServiceRequestExceptions

* clean up with finally

* tests

* retry utilities

* disable tests

* add logging to policies

* GetDatabaseAccount Fix

* Update _base.py

* retry utilities fixes

* Update _retry_utility.py

* retry utulities part 34

* Update _service_request_retry_policy.py

* remove extra logs

* policy updates

* Update _service_response_retry_policy.py

* Update _service_response_retry_policy.py

* policies updates and update operation types

* trying out fixes

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

Co-authored-by: Abhijeet Mohanty <[email protected]>

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

Co-authored-by: Abhijeet Mohanty <[email protected]>

* Skipped proxy test for debugging

* annotation fix

* Fixed some tests cases

* test fixes

* Update test_service_retry_policies_async.py

* Fixed some mocking behavior

* fixed pylint issues

* Added aiohttp minimum dependency

* Updated changelog and setup.py

* Updated changelog

* Add changelog and fix tests.

* Fix tests

* bootstrapping with global endpoint as previous for writes

* Add headers and cleanup

* cleanup and retry all service request headers

* Don't retry on a none previous

* Updated the business logic with current and previous, fixed database account refresh and some retry policies

* fix client id

* Reacting to comments

* Added print statements and fixed some retry logic

* Revert getDatabase in mark endpoint

* Fixed some pylint and changelog issues

* Fixed version

* fix bug with type check, update tests

* Update test_service_retry_policies_async.py

* sync tests updates

* Reacting to comments and fixing service request retry policy

* Code review comments and pylint issues

* Fixed tests and pylint

* more sync mock tests - missing async copies

* Fixed min aiohttp requirements

* Update _retry_utility_async.py

* Change to check operation type in operations

* push initial GEM mock test

* Update test_service_retry_policies.py

* Fixed extra retries

* sync tests

* Update test_service_retry_policies_async.py

* Fixed extra retries and relevant tests

* Only delay retry by one second

* async tests - need to split up inheritance ones since endpoint unavailable stops extra retries

* Change retry strategy

* add sub-class errors tests

* change old tests, refactoring, fix mocking bleed

* Fix a test

* clear last routed location pythonic

* Removed aiohttp dependency

* catch import errors

* Skipped global endpoint manager test for debugging

* Fixed tests

* Removed skips

* fix live tests and print statements for debugging

* cleanup of few tests

* updated globaldb mock

* Moved some of the high  offer throughput tests to live tests

* Fixed global endpoint retry async test

* Tried fixing global endpoint retry async test

* no swaps on success test

* fix import

* Tried fixing global endpoint retry async test

* Added separate split live tests

* Added live platform matrix

* some test fixes

* Fixed live test pipeline

* Moved test resource id to cosmosLong

* Updated live tests

* Running live tests with proper flag

* testing logging experiments

* fix tests

* honor testmark argument through a safe environment variable, versus accessing the value directly

* more test fixes

* remove accidental log files

* Fixed issues with swapping and retry policies

* Fixed issues with swapping and retry policies

* Marking endpoint as down fix

* more test fixes

* Remove print statements

* Fixed some minor issues with emulator tests

* split change feed tests

* Fixed emulator tests

* updated changelog

* Fixed emulator tests again

* Fixed emulator tests and event loop

* vector/fts query tests

* Fix session token live tests

* hybrid search query fixes

* Fixed live test name

* fallback to regional

* fix ci tests

* Update conftest.py

* Database accounts call will timeout in 5 seconds

* Change timeouts and update docs

* call updates to endpoint policy and location cache

* Health check for endpoitns

* database account retry policy

* Fix parameter error

* Retry on cosmos error fix

* Retry on service request error fix

* None checks for request in retry utilities

* lowercase constructed regional endpoint

* fix global endpoint as unhealthy

* fix parsing test

* Added logic for swapping on health check failed

* Fixed log statement

* fix pylint, docs, and remove print statements

* fix pylint

* fix some tests

* Prepared for release

* Increment package version after release of azure-cosmos (Azure#39570)

* Moved container create tests to emulator

* Fix timeout test

* Optimized tests to run live tests with correct configs

* Fixed some errors

* combine session token split tests

* Moved query split test to query

* Update sdk/cosmos/azure-cosmos/test/test_latest_session_token_async.py

Co-authored-by: Simon Moreno <[email protected]>

* rename to async

* remove unnecessary imports

---------

Co-authored-by: Simon Moreno <[email protected]>
Co-authored-by: tvaron3 <[email protected]>
Co-authored-by: Abhijeet Mohanty <[email protected]>
Co-authored-by: Tomas Varon Saldarriaga <[email protected]>
Co-authored-by: Scott Beddall <[email protected]>
Co-authored-by: Azure SDK Bot <[email protected]>
Co-authored-by: Tomas Varon <[email protected]>

Author: 6 people

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -1,5 +1,15 @@
 ## Release History
 
+### 4.9.1b4 (Unreleased)
+
+#### Features Added
+
+#### Breaking Changes
+
+#### Bugs Fixed
+
+#### Other Changes
+
 ### 4.9.1b3 (2025-02-04)
 
 #### Features Added

sdk/cosmos/azure-cosmos/azure/cosmos/_version.py

@@ -19,4 +19,4 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-VERSION = "4.9.1b3"
+VERSION = "4.9.1b4"

sdk/cosmos/azure-cosmos/test/test_auto_scale.py

@@ -11,7 +11,7 @@
 from azure.cosmos import CosmosClient, cosmos_client
 from azure.cosmos import ThroughputProperties, PartitionKey
 
-@pytest.mark.cosmosSplit
+@pytest.mark.cosmosLong
 class TestAutoScale(unittest.TestCase):
     client: CosmosClient = None
     host = test_config.TestConfig.host

sdk/cosmos/azure-cosmos/test/test_auto_scale_async.py

@@ -11,7 +11,7 @@
 from azure.cosmos import ThroughputProperties, PartitionKey
 from azure.cosmos.aio import CosmosClient, DatabaseProxy
 
-@pytest.mark.cosmosSplit
+@pytest.mark.cosmosLong
 class TestAutoScaleAsync(unittest.IsolatedAsyncioTestCase):
     host = test_config.TestConfig.host
     masterKey = test_config.TestConfig.masterKey
@@ -35,7 +35,7 @@ async def asyncSetUp(self):
         self.client = CosmosClient(self.host, self.masterKey)
         self.created_database = self.client.get_database_client(self.TEST_DATABASE_ID)
 
-    async def tearDown(self):
+    async def asyncTearDown(self):
         await self.client.close()
 
     async def test_autoscale_create_container_async(self):

sdk/cosmos/azure-cosmos/test/test_change_feed_split.py

@@ -46,25 +46,8 @@ def test_query_change_feed_with_split(self):
         assert len(iter_list) == 1
         continuation = created_collection.client_connection.last_response_headers['etag']
 
-        print("Triggering a split in test_query_change_feed_with_split")
-        created_collection.replace_throughput(11000)
-        print("changed offer to 11k")
-        print("--------------------------------")
-        print("Waiting for split to complete")
-        start_time = time.time()
+        test_config.TestConfig.trigger_split(created_collection, 11000)
 
-        while True:
-            offer = created_collection.get_throughput()
-            if offer.properties['content'].get('isOfferReplacePending', False):
-                if time.time() - start_time > 60 * 25:  # timeout test at 25 minutes
-                    unittest.skip("Partition split didn't complete in time.")
-                else:
-                    print("Waiting for split to complete")
-                    time.sleep(60)
-            else:
-                break
-
-        print("Split in test_query_change_feed_with_split has completed")
         print("creating few more documents")
         new_documents = [{'pk': 'pk2', 'id': 'doc2'}, {'pk': 'pk3', 'id': 'doc3'}, {'pk': 'pk4', 'id': 'doc4'}]
         expected_ids = ['doc2', 'doc3', 'doc4']

sdk/cosmos/azure-cosmos/test/test_change_feed_split_async.py

@@ -36,7 +36,7 @@ async def asyncSetUp(self):
         self.client = CosmosClient(self.host, self.masterKey)
         self.created_database = self.client.get_database_client(self.TEST_DATABASE_ID)
 
-    async def tearDown(self):
+    async def asyncTearDown(self):
         await self.client.close()
 
     async def test_query_change_feed_with_split_async(self):
@@ -59,25 +59,8 @@ async def test_query_change_feed_with_split_async(self):
         assert len(iter_list) == 1
         continuation = created_collection.client_connection.last_response_headers['etag']
 
-        print("Triggering a split in test_query_change_feed_with_split")
-        await created_collection.replace_throughput(11000)
-        print("changed offer to 11k")
-        print("--------------------------------")
-        print("Waiting for split to complete")
-        start_time = time.time()
-
-        while True:
-            offer = await created_collection.get_throughput()
-            if offer.properties['content'].get('isOfferReplacePending', False):
-                if time.time() - start_time > 60 * 25:  # timeout test at 25 minutes
-                    unittest.skip("Partition split didn't complete in time.")
-                else:
-                    print("Waiting for split to complete")
-                    time.sleep(60)
-            else:
-                break
-
-        print("Split in test_query_change_feed_with_split has completed")
+        await test_config.TestConfig.trigger_split_async(created_collection, 11000)
+
         print("creating few more documents")
         new_documents = [{'pk': 'pk2', 'id': 'doc2'}, {'pk': 'pk3', 'id': 'doc3'}, {'pk': 'pk4', 'id': 'doc4'}]
         expected_ids = ['doc2', 'doc3', 'doc4']

sdk/cosmos/azure-cosmos/test/test_config.py

@@ -4,6 +4,7 @@
 import collections
 import os
 import time
+import unittest
 import uuid
 
 from azure.cosmos.cosmos_client import CosmosClient
@@ -173,6 +174,48 @@ async def _validate_offset_limit(cls, created_collection, query, results):
         query_iterable = created_collection.query_items(query=query)
         assert list(map(lambda doc: doc['pk'], [item async for item in query_iterable])) == results
 
+    @staticmethod
+    def trigger_split(container, throughput):
+        print("Triggering a split in session token helpers")
+        container.replace_throughput(throughput)
+        print("changed offer to 11k")
+        print("--------------------------------")
+        print("Waiting for split to complete")
+        start_time = time.time()
+
+        while True:
+            offer = container.get_throughput()
+            if offer.properties['content'].get('isOfferReplacePending', False):
+                if time.time() - start_time > 60 * 25:  # timeout test at 25 minutes
+                    unittest.skip("Partition split didn't complete in time.")
+                else:
+                    print("Waiting for split to complete")
+                    time.sleep(60)
+            else:
+                break
+        print("Split in session token helpers has completed")
+
+    @staticmethod
+    async def trigger_split_async(container, throughput):
+        print("Triggering a split in session token helpers")
+        await container.replace_throughput(throughput)
+        print("changed offer to 11k")
+        print("--------------------------------")
+        print("Waiting for split to complete")
+        start_time = time.time()
+
+        while True:
+            offer = await container.get_throughput()
+            if offer.properties['content'].get('isOfferReplacePending', False):
+                if time.time() - start_time > 60 * 25:  # timeout test at 25 minutes
+                    unittest.skip("Partition split didn't complete in time.")
+                else:
+                    print("Waiting for split to complete")
+                    time.sleep(60)
+            else:
+                break
+        print("Split in session token helpers has completed")
+
 
 def get_vector_indexing_policy(embedding_type):
     return {

sdk/cosmos/azure-cosmos/test/test_container_properties_cache.py

@@ -16,7 +16,7 @@
 from azure.cosmos.partition_key import PartitionKey
 
 
-@pytest.mark.cosmosEmulator
+@pytest.mark.cosmosLong
 class TestContainerPropertiesCache(unittest.TestCase):
     """Python CRUD Tests.
         """

sdk/cosmos/azure-cosmos/test/test_container_properties_cache_async.py

@@ -17,7 +17,7 @@
 from azure.cosmos.partition_key import PartitionKey
 
 
-@pytest.mark.cosmosEmulator
+@pytest.mark.cosmosLong
 class TestContainerPropertiesCache(unittest.IsolatedAsyncioTestCase):
     """Python CRUD Tests.
         """
@@ -42,7 +42,7 @@ async def asyncSetUp(self):
         self.client = CosmosClient(self.host, self.masterKey)
         self.databaseForTest = await self.client.create_database_if_not_exists(self.configs.TEST_DATABASE_ID)
 
-    async def tearDown(self):
+    async def asyncTearDown(self):
         await self.client.close()
 
     async def test_container_properties_cache_async(self):

sdk/cosmos/azure-cosmos/test/test_crud_async.py

@@ -85,7 +85,7 @@ async def asyncSetUp(self):
         self.client = CosmosClient(self.host, self.masterKey)
         self.database_for_test = self.client.get_database_client(self.configs.TEST_DATABASE_ID)
 
-    async def tearDown(self):
+    async def asyncTearDown(self):
         await self.client.close()
 
     async def test_database_crud_async(self):
@@ -1708,37 +1708,24 @@ async def test_client_connection_retry_configuration_async(self):
         total_time_for_three_retries = await self.initialize_client_with_connection_urllib_retry_config(3)
         assert total_time_for_three_retries > total_time_for_two_retries
 
-        total_time_for_two_retries = await self.initialize_client_with_connection_core_retry_config(2)
-        total_time_for_three_retries = await self.initialize_client_with_connection_core_retry_config(3)
-        assert total_time_for_three_retries > total_time_for_two_retries
-
     async def initialize_client_with_connection_urllib_retry_config(self, retries):
-        start_time = time.time()
-        try:
-            async with CosmosClient("https://localhost:9999", TestCRUDOperationsAsync.masterKey,
-                                    retry_total=retries, retry_connect=retries, retry_read=retries,
-                                    retry_backoff_max=0.3, retry_on_status_codes=[500, 502, 504]) as client:
-                print('Async initialization')
-            self.fail()
-        except AzureError as e:
-            end_time = time.time()
-            return end_time - start_time
+        async with CosmosClient(TestCRUDOperationsAsync.host, TestCRUDOperationsAsync.masterKey,
+                                retry_total=retries, retry_connect=retries, retry_read=retries,
+                                retry_backoff_max=0.3,
+                                connection_timeout=.000000001) as client:
+            databaseForTest = client.get_database_client(self.configs.TEST_DATABASE_ID)
+            container = databaseForTest.get_container_client(self.configs.TEST_SINGLE_PARTITION_CONTAINER_ID)
+            start_time = time.time()
+            # client does a getDatabaseAccount on initialization, which will not time out because
+            # there is a forced timeout for those calls
 
-    async def initialize_client_with_connection_core_retry_config(self, retries):
-        start_time = time.time()
-        try:
-            async with CosmosClient(
-                    "https://localhost:9999",
-                    TestCRUDOperationsAsync.masterKey,
-                    retry_total=retries,
-                    retry_read=retries,
-                    retry_connect=retries,
-                    retry_status=retries) as client:
+            try:
+                await container.create_item(body={'id': str(uuid.uuid4()), 'name': 'sample'})
                 print('Async initialization')
-            self.fail()
-        except AzureError as e:
-            end_time = time.time()
-            return end_time - start_time
+                self.fail()
+            except AzureError as e:
+                end_time = time.time()
+                return end_time - start_time
 
     # TODO: Skipping this test to debug later
     @unittest.skip
@@ -2501,7 +2488,7 @@ async def priority_mock_execute_function(function, *args, **kwargs):
         # Negative Test: Verify that if we send a value other than High or Low that it will not set the header value
         # and result in bad request
         try:
-            item2_read = created_container.read_item("item2", "pk2", priority="Medium")
+            item2_read = await created_container.read_item("item2", "pk2", priority="Medium")
         except exceptions.CosmosHttpResponseError as e:
             assert e.status_code == StatusCodes.BAD_REQUEST
         _retry_utility_async.ExecuteFunctionAsync = self.OriginalExecuteFunction