[Cosmos] Add max_integrated_cache_staleness param option to item methods

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,13 +3,17 @@
 ### 4.3.0b3 (Unreleased)
 
 #### Features Added
+- Added new **provisional** `max_integrated_cache_staleness_in_ms` parameter to read item and query items APIs in order
+  to make use of the **preview** CosmosDB integrated cache functionality.
+  Please see [Azure Cosmos DB integrated cache](https://docs.microsoft.com/azure/cosmos-db/integrated-cache) for more details.
 - Added support for split-proof queries for the async client
 
 ### Bugs fixed
-- Default consistency level for the sync and async clients is no longer "Session" and will instead be set to the 
+- Default consistency level for the sync and async clients is no longer `Session` and will instead be set to the 
   consistency level of the user's cosmos account setting on initialization if not passed during client initialization. 
-  This change will impact client application in terms of RUs and latency. Users relying on default `Session` consistency will need to pass it explicitly if their account consistency is different than `Session`.
-  Please see [Consistency Levels in Azure Cosmos DB](https://docs.microsoft.com/azure/cosmos-db/consistency-levels) for more details.
+  This change will impact client application in terms of RUs and latency. Users relying on default `Session` consistency
+  will need to pass it explicitly if their account consistency is different than `Session`.
+  Please see [Consistency Levels in Azure Cosmos DB](https://docs.microsoft.com/azure/cosmos-db/consistency-levels) for more details.  
 - Fixed invalid request body being sent when passing in `serverScript` body parameter to replace operations for trigger, sproc and udf resources.
 
 ### 4.3.0b2 (2022-01-25)

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

@@ -71,6 +71,7 @@
     'query_version': 'queryVersion'
 }
 
+
 def _get_match_headers(kwargs):
     # type: (Dict[str, Any]) -> Tuple(Optional[str], Optional[str])
     if_match = kwargs.pop('if_match', None)
@@ -112,14 +113,14 @@ def build_options(kwargs):
 
 
 def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
-    cosmos_client_connection,
-    default_headers,
-    verb,
-    path,
-    resource_id,
-    resource_type,
-    options,
-    partition_key_range_id=None,
+        cosmos_client_connection,
+        default_headers,
+        verb,
+        path,
+        resource_id,
+        resource_type,
+        options,
+        partition_key_range_id=None,
 ):
     """Gets HTTP request headers.
 
@@ -292,6 +293,9 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
     if options.get("populateQuotaInfo"):
         headers[http_constants.HttpHeaders.PopulateQuotaInfo] = options["populateQuotaInfo"]
 
+    if options.get("maxIntegratedCacheStaleness"):
+        headers[http_constants.HttpHeaders.DedicatedGatewayCacheStaleness] = options["maxIntegratedCacheStaleness"]
+
     return headers
 
 
@@ -638,7 +642,7 @@ def ParsePaths(paths):
                 newIndex += 1
 
             # This will extract the token excluding the quote chars
-            token = path[currentIndex + 1 : newIndex]
+            token = path[currentIndex + 1: newIndex]
             tokens.append(token)
             currentIndex = newIndex + 1
         else:
@@ -657,3 +661,9 @@ def ParsePaths(paths):
             tokens.append(token)
 
     return tokens
+
+
+def validate_cache_staleness_value(max_integrated_cache_staleness):
+    int(max_integrated_cache_staleness)  # Will throw error if data type cant be converted to int
+    if max_integrated_cache_staleness <= 0:
+        raise ValueError("Parameter 'max_integrated_cache_staleness_in_ms' can only be a positive integer.")

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

@@ -141,7 +141,10 @@ def __init__(
             http_constants.HttpHeaders.IsContinuationExpected: False,
         }
 
-        # Keeps the latest response headers from server.
+        # Keeps the latest request headers from the client.
+        self.last_request_headers = None
+
+        # Keeps the latest response headers from the server.
         self.last_response_headers = None
 
         self._useMultipleWriteLocations = False
@@ -1874,6 +1877,7 @@ def ExecuteStoredProcedure(self, sproc_link, params, options=None, **kwargs):
 
         # ExecuteStoredProcedure will use WriteEndpoint since it uses POST operation
         request_params = _request_object.RequestObject("sprocs", documents._OperationType.ExecuteJavaScript)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = self.__Post(path, request_params, params, headers, **kwargs)
         return result
 
@@ -2043,6 +2047,7 @@ def GetDatabaseAccount(self, url_connection=None, **kwargs):
         headers = base.GetHeaders(self, initial_headers, "get", "", "", "", {})  # path  # id  # type
 
         request_params = _request_object.RequestObject("databaseaccount", documents._OperationType.Read, url_connection)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = self.__Get("", request_params, headers, **kwargs)
         database_account = documents.DatabaseAccount()
         database_account.DatabasesLink = "/dbs/"
@@ -2097,6 +2102,7 @@ def Create(self, body, path, typ, id, initial_headers, options=None, **kwargs):
         # Create will use WriteEndpoint since it uses POST operation
 
         request_params = _request_object.RequestObject(typ, documents._OperationType.Create)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = self.__Post(path, request_params, body, headers, **kwargs)
 
         # update session for write request
@@ -2130,6 +2136,7 @@ def Upsert(self, body, path, typ, id, initial_headers, options=None, **kwargs):
 
         # Upsert will use WriteEndpoint since it uses POST operation
         request_params = _request_object.RequestObject(typ, documents._OperationType.Upsert)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = self.__Post(path, request_params, body, headers, **kwargs)
         # update session for write request
         self._UpdateSessionIfRequired(headers, result, self.last_response_headers)
@@ -2160,6 +2167,7 @@ def Replace(self, resource, path, typ, id, initial_headers, options=None,
         headers = base.GetHeaders(self, initial_headers, "put", path, id, typ, options)
         # Replace will use WriteEndpoint since it uses PUT operation
         request_params = _request_object.RequestObject(typ, documents._OperationType.Replace)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = self.__Put(path, request_params, resource, headers, **kwargs)
 
         # update session for request mutates data on server side
@@ -2189,6 +2197,7 @@ def Read(self, path, typ, id, initial_headers, options=None, **kwargs):  # pylin
         headers = base.GetHeaders(self, initial_headers, "get", path, id, typ, options)
         # Read will use ReadEndpoint since it uses GET operation
         request_params = _request_object.RequestObject(typ, documents._OperationType.Read)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = self.__Get(path, request_params, headers, **kwargs)
         return result
 
@@ -2216,6 +2225,7 @@ def DeleteResource(self, path, typ, id, initial_headers, options=None,
         headers = base.GetHeaders(self, initial_headers, "delete", path, id, typ, options)
         # Delete will use WriteEndpoint since it uses DELETE operation
         request_params = _request_object.RequestObject(typ, documents._OperationType.Delete)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = self.__Delete(path, request_params, headers, **kwargs)
 
         # update session for request mutates data on server side
@@ -2413,6 +2423,7 @@ def __GetBodiesFromQueryResult(result):
             request_params = _request_object.RequestObject(
                 typ, documents._OperationType.QueryPlan if is_query_plan else documents._OperationType.ReadFeed)
             headers = base.GetHeaders(self, initial_headers, "get", path, id_, typ, options, partition_key_range_id)
+            self.last_request_headers = headers.copy()
             result, self.last_response_headers = self.__Get(path, request_params, headers, **kwargs)
             if response_hook:
                 response_hook(self.last_response_headers, result)
@@ -2436,8 +2447,9 @@ def __GetBodiesFromQueryResult(result):
 
         # Query operations will use ReadEndpoint even though it uses POST(for regular query operations)
         request_params = _request_object.RequestObject(typ, documents._OperationType.SqlQuery)
-        req_headers = base.GetHeaders(self, initial_headers, "post", path, id_, typ, options, partition_key_range_id)
-        result, self.last_response_headers = self.__Post(path, request_params, query, req_headers, **kwargs)
+        headers = base.GetHeaders(self, initial_headers, "post", path, id_, typ, options, partition_key_range_id)
+        self.last_request_headers = headers.copy()
+        result, self.last_response_headers = self.__Post(path, request_params, query, headers, **kwargs)
 
         if response_hook:
             response_hook(self.last_response_headers, result)

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

@@ -141,7 +141,10 @@ def __init__(
         if consistency_level is not None:
             self.default_headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level
 
-        # Keeps the latest response headers from server.
+        # Keeps the latest request headers from the client.
+        self.last_request_headers = None
+
+        # Keeps the latest response headers from the server.
         self.last_response_headers = None
 
         self._useMultipleWriteLocations = False
@@ -326,6 +329,7 @@ async def GetDatabaseAccount(self, url_connection=None, **kwargs):
         headers = base.GetHeaders(self, initial_headers, "get", "", "", "", {})  # path  # id  # type
 
         request_params = _request_object.RequestObject("databaseaccount", documents._OperationType.Read, url_connection)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = await self.__Get("", request_params, headers, **kwargs)
         database_account = documents.DatabaseAccount()
         database_account.DatabasesLink = "/dbs/"
@@ -575,6 +579,7 @@ async def ExecuteStoredProcedure(self, sproc_link, params, options=None, **kwarg
 
         # ExecuteStoredProcedure will use WriteEndpoint since it uses POST operation
         request_params = _request_object.RequestObject("sprocs", documents._OperationType.ExecuteJavaScript)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = await self.__Post(path, request_params, params, headers, **kwargs)
         return result
 
@@ -603,6 +608,7 @@ async def Create(self, body, path, typ, id, initial_headers, options=None, **kwa
         # Create will use WriteEndpoint since it uses POST operation
 
         request_params = _request_object.RequestObject(typ, documents._OperationType.Create)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = await self.__Post(path, request_params, body, headers, **kwargs)
 
         # update session for write request
@@ -717,6 +723,7 @@ async def Upsert(self, body, path, typ, id, initial_headers, options=None, **kwa
 
         # Upsert will use WriteEndpoint since it uses POST operation
         request_params = _request_object.RequestObject(typ, documents._OperationType.Upsert)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = await self.__Post(path, request_params, body, headers, **kwargs)
         # update session for write request
         self._UpdateSessionIfRequired(headers, result, self.last_response_headers)
@@ -957,6 +964,7 @@ async def Read(self, path, typ, id, initial_headers, options=None, **kwargs):  #
         headers = base.GetHeaders(self, initial_headers, "get", path, id, typ, options)
         # Read will use ReadEndpoint since it uses GET operation
         request_params = _request_object.RequestObject(typ, documents._OperationType.Read)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = await self.__Get(path, request_params, headers, **kwargs)
         return result
 
@@ -1219,6 +1227,7 @@ async def Replace(self, resource, path, typ, id, initial_headers, options=None,
         headers = base.GetHeaders(self, initial_headers, "put", path, id, typ, options)
         # Replace will use WriteEndpoint since it uses PUT operation
         request_params = _request_object.RequestObject(typ, documents._OperationType.Replace)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = await self.__Put(path, request_params, resource, headers, **kwargs)
 
         # update session for request mutates data on server side
@@ -1463,6 +1472,7 @@ async def DeleteResource(self, path, typ, id, initial_headers, options=None, **k
         headers = base.GetHeaders(self, initial_headers, "delete", path, id, typ, options)
         # Delete will use WriteEndpoint since it uses DELETE operation
         request_params = _request_object.RequestObject(typ, documents._OperationType.Delete)
+        self.last_request_headers = headers.copy()
         result, self.last_response_headers = await self.__Delete(path, request_params, headers, **kwargs)
 
         # update session for request mutates data on server side
@@ -2253,6 +2263,7 @@ def __GetBodiesFromQueryResult(result):
             request_params = _request_object.RequestObject(typ,
                         documents._OperationType.QueryPlan if is_query_plan else documents._OperationType.ReadFeed)
             headers = base.GetHeaders(self, initial_headers, "get", path, id_, typ, options, partition_key_range_id)
+            self.last_request_headers = headers.copy()
             result, self.last_response_headers = await self.__Get(path, request_params, headers, **kwargs)
             if response_hook:
                 response_hook(self.last_response_headers, result)
@@ -2276,8 +2287,9 @@ def __GetBodiesFromQueryResult(result):
 
         # Query operations will use ReadEndpoint even though it uses POST(for regular query operations)
         request_params = _request_object.RequestObject(typ, documents._OperationType.SqlQuery)
-        req_headers = base.GetHeaders(self, initial_headers, "post", path, id_, typ, options, partition_key_range_id)
-        result, self.last_response_headers = await self.__Post(path, request_params, query, req_headers, **kwargs)
+        headers = base.GetHeaders(self, initial_headers, "post", path, id_, typ, options, partition_key_range_id)
+        self.last_request_headers = headers.copy()
+        result, self.last_response_headers = await self.__Post(path, request_params, query, headers, **kwargs)
 
         if response_hook:
             response_hook(self.last_response_headers, result)