Revert "Consolidate Container Properties cache in Client Instance (#35293)" (#35642)

This reverts commit 35f9733.

Author: bambriz

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -10,7 +10,6 @@
 #### 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,10 +853,3 @@ 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,7 +45,6 @@
 )
 
 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
@@ -144,7 +143,7 @@ def __init__(
 
         self.connection_policy = connection_policy or ConnectionPolicy()
         self.partition_resolvers: Dict[str, RangePartitionResolver] = {}
-        self.__container_properties_cache: Dict[str, Dict[str, Any]] = {}
+        self.partition_key_definition_cache: Dict[str, Any] = {}
         self.default_headers: Dict[str, Any] = {
             http_constants.HttpHeaders.CacheControl: "no-cache",
             http_constants.HttpHeaders.Version: http_constants.Versions.CurrentVersion,
@@ -231,26 +230,6 @@ 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,
@@ -1282,6 +1261,7 @@ 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(
@@ -3290,14 +3270,13 @@ 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: Optional[Dict[str, Any]]
+        partition_key_definition = None
         # If the document collection link is present in the cache, then use the cached partitionkey definition
-        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")
+        if collection_link in self.partition_key_definition_cache:
+            partition_key_definition = self.partition_key_definition_cache.get(collection_link)
         # Else read the collection from backend and add it to the cache
         else:
-            container = self.ReadContainer(collection_link)
-            partition_key_definition = container.get("partitionKey")
-            self.__container_properties_cache[collection_link] = _set_properties_cache(container)
+            collection = self.ReadContainer(collection_link)
+            partition_key_definition = collection.get("partitionKey")
+            self.partition_key_definition_cache[collection_link] = partition_key_definition
         return partition_key_definition

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

@@ -36,8 +36,7 @@
     validate_cache_staleness_value,
     _deserialize_throughput,
     _replace_throughput,
-    GenerateGuidId,
-    _set_properties_cache
+    GenerateGuidId
 )
 from ..exceptions import CosmosResourceNotFoundError
 from ..http_constants import StatusCodes
@@ -82,21 +81,19 @@ 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.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
+        if self._properties is None:
+            self._properties = await self.read()
+        return self._properties
 
     @property
     async def is_system_key(self) -> bool:
@@ -170,10 +167,12 @@ async def read(
             request_options["populatePartitionKeyRangeStatistics"] = populate_partition_key_range_statistics
         if populate_quota_info is not None:
             request_options["populateQuotaInfo"] = populate_quota_info
-        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
+
+        collection_link = self.container_link
+        self._properties = await self.client_connection.ReadContainer(
+            collection_link, options=request_options, **kwargs
+        )
+        return self._properties
 
     @distributed_trace_async
     async def create_item(

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

@@ -47,7 +47,6 @@
     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
@@ -145,7 +144,7 @@ def __init__(
 
         self.connection_policy = connection_policy or ConnectionPolicy()
         self.partition_resolvers: Dict[str, RangePartitionResolver] = {}
-        self.__container_properties_cache: Dict[str, Dict[str, Any]] = {}
+        self.partition_key_definition_cache: Dict[str, Any] = {}
         self.default_headers: Dict[str, Any] = {
             http_constants.HttpHeaders.CacheControl: "no-cache",
             http_constants.HttpHeaders.Version: http_constants.Versions.CurrentVersion,
@@ -227,26 +226,6 @@ 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.
@@ -3055,14 +3034,13 @@ 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.__container_properties_cache:
-            cached_container = self.__container_properties_cache.get(collection_link)
-            partitionKeyDefinition = cached_container.get("partitionKey")
+        if collection_link in self.partition_key_definition_cache:
+            partitionKeyDefinition = self.partition_key_definition_cache.get(collection_link)
         # Else read the collection from backend and add it to the cache
         else:
-            container = await self.ReadContainer(collection_link)
-            partitionKeyDefinition = container.get("partitionKey")
-            self.__container_properties_cache[collection_link] = _set_properties_cache(container)
+            collection = await self.ReadContainer(collection_link)
+            partitionKeyDefinition = collection.get("partitionKey")
+            self.partition_key_definition_cache[collection_link] = partitionKeyDefinition
 
         # 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,8 +36,7 @@
     validate_cache_staleness_value,
     _deserialize_throughput,
     _replace_throughput,
-    GenerateGuidId,
-    _set_properties_cache
+    GenerateGuidId
 )
 from .exceptions import CosmosResourceNotFoundError
 from .http_constants import StatusCodes
@@ -84,19 +83,17 @@ 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.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
+        if self._properties is None:
+            self._properties = self.read()
+        return self._properties
 
     @property
     def is_system_key(self) -> bool:
@@ -176,10 +173,11 @@ 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
-        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
+        collection_link = self.container_link
+        self._properties = self.client_connection.ReadContainer(
+            collection_link, options=request_options, **kwargs
+        )
+        return self._properties
 
     @distributed_trace
     def read_item(  # pylint:disable=docstring-missing-param