Consolidate Container Properties cache in Client Instance (#35293)

* Add client cache to hold container properties

This adds a client level cache for container properties.

* Added tests for container properties cache

* Update _cosmos_client_connection_async.py

* Update _cosmos_client_connection_async.py

* Fix typos

* Update container.py

* Update _container.py

* Update cache test

* Update test_container_properties_cache.py

* update tests

* Update test_container_properties_cache.py

* Update tests

updated tests that were compensating for extra container reads, the extra reads no longer happen with adding the container cache properties in the client.

* tests updates

* Mypy fixes

* mypy fix

* Update

* Cache update

Fixed some reccomnded changes, and added optimizations for the cache. Only including needed properties info that will be unlikely to change.

* test update

* Update CHANGELOG.md

* added set_properties_cache method

Added a method in base.py to set the properties cache so we only have to change one location in case in the future we want to cache additional properties

* Update _base.py

* Update _cosmos_client_connection.py

* MyPy fixes

* Fix mypy issue

* Update _cosmos_client_connection.py

* Update _cosmos_client_connection.py

* Update _cosmos_client_connection_async.py

* Update _cosmos_client_connection.py

* Update container cache

Changed container cache to be a private attribute.
Also added tests to verify the contents of the cache to be correct ones.

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

Co-authored-by: Kushagra Thapar <[email protected]>

* Additional pylint updates

* more pylint fixes

* pylint fix

* Update test_container_properties_cache.py

---------

Co-authored-by: Kushagra Thapar <[email protected]>

Author: bambriz

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -10,6 +10,7 @@
 #### Bugs Fixed
 * Fixed a bug where change feed query in Async client was not returning all pages due to case-sensitive response headers. See [PR 35090](https://github.com/Azure/azure-sdk-for-python/pull/35090)
 * Fixed a bug when a retryable exception occurs in the first page of a query execution causing query to return 0 results. See [PR 35090](https://github.com/Azure/azure-sdk-for-python/pull/35090).
+* Consolidated Container Properties Cache to be in the Client to cache partition key definition and container rid to avoid unnecessary container reads. See [PR 35293](https://github.com/Azure/azure-sdk-for-python/pull/35293)
 
 
 #### Other Changes

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -853,3 +853,10 @@ def _format_batch_operations(
         final_operations.append(operation)
 
     return final_operations
+
+
+def _set_properties_cache(properties: Dict[str, Any]) -> Dict[str, Any]:
+    return {
+        "_self": properties.get("_self", None), "_rid": properties.get("_rid", None),
+        "partitionKey": properties.get("partitionKey", None)
+    }

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -45,6 +45,7 @@
 )
 
 from . import _base as base
+from ._base import _set_properties_cache
 from . import documents
 from .documents import ConnectionPolicy, DatabaseAccount
 from ._constants import _Constants as Constants
@@ -143,7 +144,7 @@ def __init__(
 
         self.connection_policy = connection_policy or ConnectionPolicy()
         self.partition_resolvers: Dict[str, RangePartitionResolver] = {}
-        self.partition_key_definition_cache: Dict[str, Any] = {}
+        self.__container_properties_cache: Dict[str, Dict[str, Any]] = {}
         self.default_headers: Dict[str, Any] = {
             http_constants.HttpHeaders.CacheControl: "no-cache",
             http_constants.HttpHeaders.Version: http_constants.Versions.CurrentVersion,
@@ -230,6 +231,26 @@ def __init__(
         self.session: Optional[_session.Session] = None
         self._set_client_consistency_level(database_account, consistency_level)
 
+    @property
+    def _container_properties_cache(self) -> Dict[str, Dict[str, Any]]:
+        """Gets the container properties cache from the client.
+        :returns: the container properties cache for the client.
+        :rtype: Dict[str, Dict[str, Any]]"""
+        return self.__container_properties_cache
+
+    def _set_container_properties_cache(self, container_link: str, properties: Optional[Dict[str, Any]]) -> None:
+        """Sets the container properties cache for the specified container.
+
+        This will only update the properties cache for a specified container.
+        :param container_link: The container link will be used as the key to cache the container properties.
+        :type container_link: str
+        :param properties: These are the container properties to cache.
+        :type properties:  Optional[Dict[str, Any]]"""
+        if properties:
+            self.__container_properties_cache[container_link] = properties
+        else:
+            self.__container_properties_cache[container_link] = {}
+
     def _set_client_consistency_level(
         self,
         database_account: DatabaseAccount,
@@ -1261,7 +1282,6 @@ def CreateItem(
 
         if base.IsItemContainerLink(database_or_container_link):
             options = self._AddPartitionKey(database_or_container_link, document, options)
-
         return self.Create(document, path, "docs", collection_id, None, options, **kwargs)
 
     def UpsertItem(
@@ -3270,13 +3290,14 @@ def _UpdateSessionIfRequired(
             self.session.update_session(response_result, response_headers)
 
     def _get_partition_key_definition(self, collection_link: str) -> Optional[Dict[str, Any]]:
-        partition_key_definition = None
+        partition_key_definition: Optional[Dict[str, Any]]
         # If the document collection link is present in the cache, then use the cached partitionkey definition
-        if collection_link in self.partition_key_definition_cache:
-            partition_key_definition = self.partition_key_definition_cache.get(collection_link)
+        if collection_link in self.__container_properties_cache:
+            cached_container: Dict[str, Any] = self.__container_properties_cache.get(collection_link, {})
+            partition_key_definition = cached_container.get("partitionKey")
         # Else read the collection from backend and add it to the cache
         else:
-            collection = self.ReadContainer(collection_link)
-            partition_key_definition = collection.get("partitionKey")
-            self.partition_key_definition_cache[collection_link] = partition_key_definition
+            container = self.ReadContainer(collection_link)
+            partition_key_definition = container.get("partitionKey")
+            self.__container_properties_cache[collection_link] = _set_properties_cache(container)
         return partition_key_definition

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

@@ -36,7 +36,8 @@
     validate_cache_staleness_value,
     _deserialize_throughput,
     _replace_throughput,
-    GenerateGuidId
+    GenerateGuidId,
+    _set_properties_cache
 )
 from ..exceptions import CosmosResourceNotFoundError
 from ..http_constants import StatusCodes
@@ -81,19 +82,21 @@ def __init__(
     ) -> None:
         self.client_connection = client_connection
         self.id = id
-        self._properties = properties
         self.database_link = database_link
         self.container_link = "{}/colls/{}".format(database_link, self.id)
         self._is_system_key: Optional[bool] = None
         self._scripts: Optional[ScriptsProxy] = None
+        if properties:
+            self.client_connection._set_container_properties_cache(self.container_link,
+                                                                   _set_properties_cache(properties))  # pylint: disable=protected-access, line-too-long
 
     def __repr__(self) -> str:
         return "<ContainerProxy [{}]>".format(self.container_link)[:1024]
 
     async def _get_properties(self) -> Dict[str, Any]:
-        if self._properties is None:
-            self._properties = await self.read()
-        return self._properties
+        if self.container_link not in self.client_connection._container_properties_cache:  # pylint: disable=protected-access, line-too-long
+            await self.read()
+        return self.client_connection._container_properties_cache[self.container_link]  # pylint: disable=protected-access, line-too-long
 
     @property
     async def is_system_key(self) -> bool:
@@ -167,12 +170,10 @@ async def read(
             request_options["populatePartitionKeyRangeStatistics"] = populate_partition_key_range_statistics
         if populate_quota_info is not None:
             request_options["populateQuotaInfo"] = populate_quota_info
-
-        collection_link = self.container_link
-        self._properties = await self.client_connection.ReadContainer(
-            collection_link, options=request_options, **kwargs
-        )
-        return self._properties
+        container = await self.client_connection.ReadContainer(self.container_link, options=request_options, **kwargs)
+        # Only cache Container Properties that will not change in the lifetime of the container
+        self.client_connection._set_container_properties_cache(self.container_link, _set_properties_cache(container))  # pylint: disable=protected-access, line-too-long
+        return container
 
     @distributed_trace_async
     async def create_item(

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

@@ -47,6 +47,7 @@
     ProxyPolicy)
 
 from .. import _base as base
+from .._base import _set_properties_cache
 from .. import documents
 from .._routing import routing_range
 from ..documents import ConnectionPolicy, DatabaseAccount
@@ -144,7 +145,7 @@ def __init__(
 
         self.connection_policy = connection_policy or ConnectionPolicy()
         self.partition_resolvers: Dict[str, RangePartitionResolver] = {}
-        self.partition_key_definition_cache: Dict[str, Any] = {}
+        self.__container_properties_cache: Dict[str, Dict[str, Any]] = {}
         self.default_headers: Dict[str, Any] = {
             http_constants.HttpHeaders.CacheControl: "no-cache",
             http_constants.HttpHeaders.Version: http_constants.Versions.CurrentVersion,
@@ -226,6 +227,26 @@ def __init__(
         # Routing map provider
         self._routing_map_provider: SmartRoutingMapProvider = SmartRoutingMapProvider(self)
 
+    @property
+    def _container_properties_cache(self) -> Dict[str, Dict[str, Any]]:
+        """Gets the container properties cache from the client.
+        :returns: the container properties cache for the client.
+        :rtype: Dict[str, Dict[str, Any]]"""
+        return self.__container_properties_cache
+
+    def _set_container_properties_cache(self, container_link: str, properties: Optional[Dict[str, Any]]) -> None:
+        """Sets the container properties cache for the specified container.
+
+        This will only update the properties cache for a specified container.
+        :param container_link: The container link will be used as the key to cache the container properties.
+        :type container_link: str
+        :param properties: These are the container properties to cache.
+        :type properties:  Optional[Dict[str, Any]]"""
+        if properties:
+            self.__container_properties_cache[container_link] = properties
+        else:
+            self.__container_properties_cache[container_link] = {}
+
     @property
     def _Session(self) -> Optional[_session.Session]:
         """Gets the session object from the client.
@@ -3034,13 +3055,14 @@ async def _AddPartitionKey(self, collection_link, document, options):
         # TODO: Refresh the cache if partition is extracted automatically and we get a 400.1001
 
         # If the document collection link is present in the cache, then use the cached partitionkey definition
-        if collection_link in self.partition_key_definition_cache:
-            partitionKeyDefinition = self.partition_key_definition_cache.get(collection_link)
+        if collection_link in self.__container_properties_cache:
+            cached_container = self.__container_properties_cache.get(collection_link)
+            partitionKeyDefinition = cached_container.get("partitionKey")
         # Else read the collection from backend and add it to the cache
         else:
-            collection = await self.ReadContainer(collection_link)
-            partitionKeyDefinition = collection.get("partitionKey")
-            self.partition_key_definition_cache[collection_link] = partitionKeyDefinition
+            container = await self.ReadContainer(collection_link)
+            partitionKeyDefinition = container.get("partitionKey")
+            self.__container_properties_cache[collection_link] = _set_properties_cache(container)
 
         # If the collection doesn't have a partition key definition, skip it as it's a legacy collection
         if partitionKeyDefinition:

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

@@ -36,7 +36,8 @@
     validate_cache_staleness_value,
     _deserialize_throughput,
     _replace_throughput,
-    GenerateGuidId
+    GenerateGuidId,
+    _set_properties_cache
 )
 from .exceptions import CosmosResourceNotFoundError
 from .http_constants import StatusCodes
@@ -83,17 +84,19 @@ def __init__(
         self.id = id
         self.container_link = "{}/colls/{}".format(database_link, self.id)
         self.client_connection = client_connection
-        self._properties = properties
         self._is_system_key: Optional[bool] = None
         self._scripts: Optional[ScriptsProxy] = None
+        if properties:
+            self.client_connection._set_container_properties_cache(self.container_link,
+                                                                   _set_properties_cache(properties))  # pylint: disable=protected-access, line-too-long
 
     def __repr__(self) -> str:
         return "<ContainerProxy [{}]>".format(self.container_link)[:1024]
 
     def _get_properties(self) -> Dict[str, Any]:
-        if self._properties is None:
-            self._properties = self.read()
-        return self._properties
+        if self.container_link not in self.client_connection._container_properties_cache:  # pylint: disable=protected-access, line-too-long
+            self.read()
+        return self.client_connection._container_properties_cache[self.container_link]  # pylint: disable=protected-access, line-too-long
 
     @property
     def is_system_key(self) -> bool:
@@ -173,11 +176,10 @@ def read(  # pylint:disable=docstring-missing-param
             request_options["populatePartitionKeyRangeStatistics"] = populate_partition_key_range_statistics
         if populate_quota_info is not None:
             request_options["populateQuotaInfo"] = populate_quota_info
-        collection_link = self.container_link
-        self._properties = self.client_connection.ReadContainer(
-            collection_link, options=request_options, **kwargs
-        )
-        return self._properties
+        container = self.client_connection.ReadContainer(self.container_link, options=request_options, **kwargs)
+        # Only cache Container Properties that will not change in the lifetime of the container
+        self.client_connection._set_container_properties_cache(self.container_link, _set_properties_cache(container))  # pylint: disable=protected-access, line-too-long
+        return container
 
     @distributed_trace
     def read_item(  # pylint:disable=docstring-missing-param

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

@@ -0,0 +1,82 @@
+# -*- coding: utf-8 -*-
+# The MIT License (MIT)
+# Copyright (c) Microsoft Corporation. All rights reserved.
+
+"""End-to-end test.
+"""
+
+import unittest
+import uuid
+
+import pytest
+
+import azure.cosmos.cosmos_client as cosmos_client
+import azure.cosmos.exceptions as exceptions
+import test_config
+from azure.cosmos.partition_key import PartitionKey
+
+
+@pytest.mark.cosmosEmulator
+class TestContainerPropertiesCache(unittest.TestCase):
+    """Python CRUD Tests.
+        """
+
+    configs = test_config.TestConfig
+    host = configs.host
+    masterKey = configs.masterKey
+    connectionPolicy = configs.connectionPolicy
+    last_headers = []
+    client: cosmos_client.CosmosClient = None
+
+    @classmethod
+    def setUpClass(cls):
+        if (cls.masterKey == '[YOUR_KEY_HERE]' or
+                cls.host == '[YOUR_ENDPOINT_HERE]'):
+            raise Exception(
+                "You must specify your Azure Cosmos account values for "
+                "'masterKey' and 'host' at the top of this class to run the "
+                "tests.")
+        cls.client = cosmos_client.CosmosClient(cls.host, cls.masterKey)
+        cls.databaseForTest = cls.client.get_database_client(cls.configs.TEST_DATABASE_ID)
+
+    def test_container_properties_cache(self):
+        client = self.client
+        database_name = "Container Properties Cache Test DB " + str(uuid.uuid4())
+        created_db = client.create_database(database_name)
+        container_name = str(uuid.uuid4())
+        container_pk = "PK"
+        # Create The Container
+        try:
+            client.get_database_client(database_name).create_container(id=container_name, partition_key=PartitionKey(
+                path="/" + container_pk))
+        except exceptions.CosmosResourceExistsError:
+            pass
+        # Delete The cache as this is meant to test calling operations on a preexisting container
+        # and not a freshly made one. It's a private attribute so use mangled name.
+        client.client_connection._CosmosClientConnection__container_properties_cache = {}
+        # We will hot path operations to verify cache persists
+        # This will extract partition key from the item body, which will need partition key definition from
+        # container properties. We test to check the cache is empty since we just created the container
+        self.assertTrue(client.client_connection._container_properties_cache == {})
+        client.get_database_client(database_name).get_container_client(container_name).create_item(
+            body={'id': 'item1', container_pk: 'value'})
+        # Since the cache was empty, it should have called a container read to get properties. So now Cache should
+        # be populated and available even when we don't have a container instance
+        self.assertTrue(client.client_connection._container_properties_cache != {})
+        # We can test if the cache properties are correct by comparing them to a fresh read.
+        # First lets save the old cache values
+        cached_properties = created_db.get_container_client(container_name)._get_properties()
+        # Get the container dictionary out of a fresh container read
+        fresh_container_read = created_db.get_container_client(container_name).read()
+        # Now we can compare the RID and Partition Key Definition
+        self.assertEqual(cached_properties.get("_rid"), fresh_container_read.get("_rid"))
+        self.assertEqual(cached_properties.get("partitionKey"), fresh_container_read.get("partitionKey"))
+        client.delete_database(created_db)
+
+
+if __name__ == '__main__':
+    try:
+        unittest.main()
+    except SystemExit as inst:
+        if inst.args[0] is True:  # raised by sys.exit(True) when tests failed
+            raise