Prevent Unnecessary Read Container Calls (#42143)

* fix container cache

* update changelog

* fix tests

Author: tvaron3

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -8,6 +8,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed bug where container cache was not being properly updated resulting in unnecessary extra requests. See [PR 42143](https://github.com/Azure/azure-sdk-for-python/pull/42143).
 
 #### Other Changes
 

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

@@ -21,6 +21,7 @@
 
 """Create, read, update and delete items in the Azure Cosmos DB SQL API service.
 """
+import asyncio # pylint: disable=do-not-import-asyncio
 from datetime import datetime
 from typing import (Any, Dict, Mapping, Optional, Sequence, Type, Union, List, Tuple, cast, overload, AsyncIterable,
                     Callable)
@@ -90,6 +91,7 @@ def __init__(
         properties: Optional[Dict[str, Any]] = None
     ) -> None:
         self.client_connection = client_connection
+        self.container_cache_lock = asyncio.Lock()
         self.id = id
         self.database_link = database_link
         self.container_link = "{}/colls/{}".format(database_link, self.id)
@@ -111,7 +113,9 @@ async def _get_properties_with_options(self, options: Optional[Dict[str, Any]] =
 
     async def _get_properties(self, **kwargs: Any) -> Dict[str, Any]:
         if self.container_link not in self.client_connection._container_properties_cache:
-            await self.read(**kwargs)
+            async with self.container_cache_lock:
+                if self.container_link not in self.client_connection._container_properties_cache:
+                    await self.read(**kwargs)
         return self.client_connection._container_properties_cache[self.container_link]
 
     @property

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

@@ -21,6 +21,7 @@
 
 """Create, read, update and delete items in the Azure Cosmos DB SQL API service.
 """
+import threading
 import warnings
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Sequence, Union, Tuple, Mapping, cast, overload, Iterable, Callable
@@ -94,6 +95,7 @@ def __init__(
         self.id = id
         self.container_link = "{}/colls/{}".format(database_link, self.id)
         self.client_connection = client_connection
+        self.container_cache_lock = threading.Lock()
         self._is_system_key: Optional[bool] = None
         self._scripts: Optional[ScriptsProxy] = None
         if properties:
@@ -112,7 +114,9 @@ def _get_properties_with_options(self, options: Optional[Dict[str, Any]] = None)
 
     def _get_properties(self, **kwargs: Any) -> Dict[str, Any]:
         if self.container_link not in self.__get_client_container_caches():
-            self.read(**kwargs)
+            with self.container_cache_lock:
+                if self.container_link not in self.__get_client_container_caches():
+                    self.read(**kwargs)
         return self.__get_client_container_caches()[self.container_link]
 
     @property

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

@@ -2,6 +2,7 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 
 import collections
+import logging
 import os
 import time
 import unittest
@@ -553,3 +554,16 @@ def create_range(range_min: str, range_max: str, is_min_inclusive: bool = True,
 
 def create_feed_range_in_dict(feed_range):
     return FeedRangeInternalEpk(feed_range).to_dict()
+
+
+class MockHandler(logging.Handler):
+
+    def __init__(self):
+        super(MockHandler, self).__init__()
+        self.messages = []
+
+    def reset(self):
+        self.messages = []
+
+    def emit(self, record):
+        self.messages.append(record)

sdk/cosmos/azure-cosmos/tests/test_cosmos_http_logging_policy.py

@@ -23,18 +23,6 @@
     from mock import Mock  # type: ignore
 
 
-class MockHandler(logging.Handler):
-
-    def __init__(self):
-        super(MockHandler, self).__init__()
-        self.messages = []
-
-    def reset(self):
-        self.messages = []
-
-    def emit(self, record):
-        self.messages.append(record)
-
 class FilterStatusCode(logging.Filter):
     def filter(self, record):
         if hasattr(record, 'status_code') and record.status_code >= 400:
@@ -51,7 +39,7 @@ def filter(self, record):
     L2_URL: L2}
 
 
-def create_logger(name: str, mock_handler: MockHandler, level: int = logging.INFO) -> logging.Logger:
+def create_logger(name: str, mock_handler: test_config.MockHandler, level: int = logging.INFO) -> logging.Logger:
     logger = logging.getLogger(name)
     logger.addHandler(mock_handler)
     logger.setLevel(level)
@@ -87,9 +75,9 @@ def setUpClass(cls):
                 "You must specify your Azure Cosmos account values for "
                 "'masterKey' and 'host' at the top of this class to run the "
                 "tests.")
-        cls.mock_handler_default = MockHandler()
-        cls.mock_handler_diagnostic = MockHandler()
-        cls.mock_handler_filtered_diagnostic = MockHandler()
+        cls.mock_handler_default = test_config.MockHandler()
+        cls.mock_handler_diagnostic = test_config.MockHandler()
+        cls.mock_handler_filtered_diagnostic = test_config.MockHandler()
 
         # Add filter to the filtered diagnostics handler
 
@@ -233,7 +221,7 @@ def test_client_settings(self):
         multiple_write_locations = True
 
         # Client setup
-        mock_handler = MockHandler()
+        mock_handler = test_config.MockHandler()
         logger = create_logger("test_logger_client_settings", mock_handler)
 
         custom_transport = FaultInjectionTransport()
@@ -288,7 +276,7 @@ def test_client_settings(self):
 
     def test_activity_id_logging_policy(self):
         # Create a mock handler and logger for the new client
-        self.mock_handler_activity_id = MockHandler()
+        self.mock_handler_activity_id = test_config.MockHandler()
         self.logger_activity_id = create_logger("test_logger_activity_id", self.mock_handler_activity_id)
 
         # Create a new client with the logger and enable diagnostics logging

sdk/cosmos/azure-cosmos/tests/test_cosmos_http_logging_policy_async.py

@@ -16,7 +16,7 @@
 from test_fault_injection_transport_async import TestFaultInjectionTransportAsync
 
 from test_cosmos_http_logging_policy import create_logger, L1, L2, CONFIG, \
-    get_locations_list, FilterStatusCode, MockHandler
+    get_locations_list, FilterStatusCode
 
 try:
     from unittest.mock import Mock
@@ -42,10 +42,10 @@ def setUpClass(cls):
                 "tests.")
 
     async def asyncSetUp(self):
-        self.mock_handler_default = MockHandler()
-        self.mock_handler_diagnostic = MockHandler()
-        self.mock_handler_filtered_diagnostic = MockHandler()
-        self.mock_handler_activity_id = MockHandler()
+        self.mock_handler_default = test_config.MockHandler()
+        self.mock_handler_diagnostic = test_config.MockHandler()
+        self.mock_handler_filtered_diagnostic = test_config.MockHandler()
+        self.mock_handler_activity_id = test_config.MockHandler()
         # Add filter to the filtered diagnostics handler
 
         self.mock_handler_filtered_diagnostic.addFilter(FilterStatusCode())
@@ -208,7 +208,7 @@ async def test_client_settings_async(self):
         multiple_write_locations = True
 
         # Client setup
-        mock_handler = MockHandler()
+        mock_handler = test_config.MockHandler()
         logger = create_logger("test_logger_client_settings", mock_handler)
 
         custom_transport = FaultInjectionTransportAsync()

sdk/cosmos/azure-cosmos/tests/test_crud.py

@@ -4,7 +4,8 @@
 
 """End-to-end test.
 """
-
+import logging
+import threading
 import time
 import unittest
 import urllib.parse as urllib
@@ -617,6 +618,33 @@ def test_document_crud(self):
                                            replaced_document['id'],
                                            replaced_document['id'])
 
+    def test_read_collection_only_once(self):
+        # Add filter to the filtered diagnostics handler
+        mock_handler = test_config.MockHandler()
+        logger = logging.getLogger("azure.cosmos")
+        logger.setLevel(logging.INFO)
+        logger.addHandler(mock_handler)
+
+        client = cosmos_client.CosmosClient(self.host, self.masterKey)
+        database = client.get_database_client(self.configs.TEST_DATABASE_ID)
+        container = database.get_container_client(self.configs.TEST_SINGLE_PARTITION_CONTAINER_ID)
+
+        def upsert_worker():
+            # Synchronously perform the upsert item operation
+            container.upsert_item(body={'id': str(uuid.uuid4()), 'name': f'sample-'})
+
+        # Perform 10 concurrent upserts
+        threads = []
+        for i in range(10):
+            t = threading.Thread(target=upsert_worker)
+            threads.append(t)
+            t.start()
+
+        for t in threads:
+            t.join()
+        assert sum("'x-ms-thinclient-proxy-resource-type': 'colls'" in response.message for response in mock_handler.messages) == 1
+
+
     def test_document_upsert(self):
         # create database
         created_db = self.databaseForTest

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

@@ -5,6 +5,7 @@
 """End-to-end test.
 """
 import asyncio
+import logging
 import os
 import time
 import unittest
@@ -441,7 +442,7 @@ async def test_document_upsert_async(self):
         await created_collection.delete_item(item=upserted_document, partition_key=upserted_document['pk'])
         await created_collection.delete_item(item=new_document, partition_key=new_document['pk'])
 
-        # read documents after delete and verify count is same as original
+        # read documents after delete and verify count remains the same
         document_list = [document async for document in created_collection.read_all_items()]
         assert len(document_list) == before_create_documents_count
 
@@ -920,6 +921,25 @@ async def test_client_request_timeout_when_connection_retry_configuration_specif
                 await container.read_item(item=item['id'], partition_key=item['id'])
                 print('Async Initialization')
 
+    async def test_read_collection_only_once_async(self):
+        # Add filter to the filtered diagnostics handler
+        mock_handler = test_config.MockHandler()
+        logger = logging.getLogger("azure.cosmos")
+        logger.setLevel(logging.INFO)
+        logger.addHandler(mock_handler)
+
+        async with CosmosClient(self.host, self.masterKey) as client:
+            database = client.get_database_client(self.configs.TEST_DATABASE_ID)
+            container = database.get_container_client(self.configs.TEST_SINGLE_PARTITION_CONTAINER_ID)
+
+            # Perform 10 concurrent upserts
+            tasks = [
+                container.upsert_item(body={'id': str(uuid.uuid4()), 'name': f'sample-{i}'})
+                for i in range(10)
+            ]
+            await asyncio.gather(*tasks)
+        assert sum("'x-ms-thinclient-proxy-resource-type': 'colls'" in response.message for response in mock_handler.messages) == 1
+
     # TODO: Skipping this test to debug later
     @unittest.skip
     async def test_client_connection_retry_configuration_async(self):
@@ -1500,3 +1520,4 @@ async def _mock_execute_function(self, function, *args, **kwargs):
 
 if __name__ == '__main__':
     unittest.main()
+