Fix 'NoneType' object has no attribute when querying during Session n…

…ot available retry (#37578)

* Add fixes to 'NoneType' object has no attribute when querying

Fixes 'NoneType' object has no attribute when session retry is triggered and does not succesfully retry while querying.

* Update CHANGELOG.md

* update test files

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -10,6 +10,7 @@
 #### Bugs Fixed
 * Consolidated Container Properties Cache to be in the Client to cache partition key definition and container rid to avoid unnecessary container reads. See [PR 35731](https://github.com/Azure/azure-sdk-for-python/pull/35731)
 * Fixed SDK regex validation that would not allow for item ids to be longer than 255 characters. See [PR 36569](https://github.com/Azure/azure-sdk-for-python/pull/36569).
+* Fixed issue where 'NoneType' object has no attribute error was raised when a session retry happened during a query. See [PR 37578](https://github.com/Azure/azure-sdk-for-python/pull/37578).
 
 #### Other Changes
 * Getting offer thoughput when it has not been defined in a container will now give a 404/10004 instead of just a 404. See [PR 36043](https://github.com/Azure/azure-sdk-for-python/pull/36043)

sdk/cosmos/azure-cosmos/azure/cosmos/_endpoint_discovery_retry_policy.py

@@ -66,6 +66,9 @@ def ShouldRetry(self, exception):  # pylint: disable=unused-argument
         :returns: a boolean stating whether the request should be retried
         :rtype: bool
         """
+        if not self.request:
+            return False
+
         if not self.connection_policy.EnableEndpointDiscovery:
             return False
 

sdk/cosmos/azure-cosmos/azure/cosmos/_session_retry_policy.py

@@ -67,6 +67,8 @@ def ShouldRetry(self, _exception):
         :returns: a boolean stating whether the request should be retried
         :rtype: bool
         """
+        if not self.request:
+            return False
         self.session_token_retry_count += 1
         # clear previous location-based routing directive
         self.request.clear_route_to_location()

sdk/cosmos/azure-cosmos/azure/cosmos/_timeout_failover_retry_policy.py

@@ -23,7 +23,7 @@ def __init__(self, connection_policy, global_endpoint_manager, *args):
             self.location_endpoint = self.global_endpoint_manager.resolve_service_endpoint(self.request)
 
     def needsRetry(self):
-        if self.args:
+        if self.args and self.request:
             if (self.args[3].method == "GET") \
                     or http_constants.HttpHeaders.IsQueryPlanRequest in self.args[3].headers\
                     or http_constants.HttpHeaders.IsQuery in self.args[3].headers:

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

@@ -177,7 +177,7 @@ async def ExecuteAsync(client, global_endpoint_manager, function, *args, **kwarg
                 client.last_response_headers[
                     HttpHeaders.ThrottleRetryWaitTimeInMs
                 ] = resourceThrottle_retry_policy.cumulative_wait_time_in_milliseconds
-                if args and args[0].should_clear_session_token_on_session_read_failure:
+                if args and args[0].should_clear_session_token_on_session_read_failure and client.session:
                     client.session.clear_session_token(client.last_response_headers)
                 raise
 

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

@@ -12,7 +12,7 @@
 import azure.cosmos.cosmos_client as cosmos_client
 import azure.cosmos.exceptions as exceptions
 import test_config
-from azure.cosmos import http_constants, DatabaseProxy
+from azure.cosmos import http_constants, DatabaseProxy, _endpoint_discovery_retry_policy
 from azure.cosmos._execution_context.base_execution_context import _QueryExecutionContextBase
 from azure.cosmos._execution_context.query_execution_info import _PartitionedQueryExecutionInfo
 from azure.cosmos.documents import _DistinctType
@@ -862,6 +862,90 @@ def test_computed_properties_query(self):
         self.assertEqual(len(queried_items), 0)
         self.created_db.delete_container(created_collection.id)
 
+    def test_query_request_params_none_retry_policy(self):
+        created_collection = self.created_db.create_container(
+            "query_request_params_none_retry_policy_" + str(uuid.uuid4()), PartitionKey(path="/pk"))
+        items = [
+            {'id': str(uuid.uuid4()), 'pk': 'test', 'val': 5},
+            {'id': str(uuid.uuid4()), 'pk': 'test', 'val': 5},
+            {'id': str(uuid.uuid4()), 'pk': 'test', 'val': 5}]
+
+        for item in items:
+            created_collection.create_item(body=item)
+
+        self.OriginalExecuteFunction = retry_utility.ExecuteFunction
+        # Test session retry will properly push the exception when retries run out
+        retry_utility.ExecuteFunction = self._MockExecuteFunctionSessionRetry
+        try:
+            query = "SELECT * FROM c"
+            items = created_collection.query_items(
+                query=query,
+                enable_cross_partition_query=True
+            )
+            fetch_results = list(items)
+        except exceptions.CosmosHttpResponseError as e:
+            self.assertEqual(e.status_code, 404)
+            self.assertEqual(e.sub_status, 1002)
+
+        # Test endpoint discovery retry
+        retry_utility.ExecuteFunction = self._MockExecuteFunctionEndPointRetry
+        _endpoint_discovery_retry_policy.EndpointDiscoveryRetryPolicy.Max_retry_attempt_count = 3
+        _endpoint_discovery_retry_policy.EndpointDiscoveryRetryPolicy.Retry_after_in_milliseconds = 10
+        try:
+            query = "SELECT * FROM c"
+            items = created_collection.query_items(
+                query=query,
+                enable_cross_partition_query=True
+            )
+            fetch_results = list(items)
+        except exceptions.CosmosHttpResponseError as e:
+            self.assertEqual(e.status_code, http_constants.StatusCodes.FORBIDDEN)
+            self.assertEqual(e.sub_status, http_constants.SubStatusCodes.WRITE_FORBIDDEN)
+        _endpoint_discovery_retry_policy.EndpointDiscoveryRetryPolicy.Max_retry_attempt_count = 120
+        _endpoint_discovery_retry_policy.EndpointDiscoveryRetryPolicy.Retry_after_in_milliseconds = 1000
+
+        # Finally lets test timeout failover retry
+        retry_utility.ExecuteFunction = self._MockExecuteFunctionTimeoutFailoverRetry
+        try:
+            query = "SELECT * FROM c"
+            items = created_collection.query_items(
+                query=query,
+                enable_cross_partition_query=True
+            )
+            fetch_results = list(items)
+        except exceptions.CosmosHttpResponseError as e:
+            self.assertEqual(e.status_code, http_constants.StatusCodes.REQUEST_TIMEOUT)
+            retry_utility.ExecuteFunction = self.OriginalExecuteFunction
+        retry_utility.ExecuteFunction = self.OriginalExecuteFunction
+        self.created_db.delete_container(created_collection.id)
+
+
+    def _MockExecuteFunctionSessionRetry(self, function, *args, **kwargs):
+        if args:
+            if args[1].operation_type == 'SqlQuery':
+                ex_to_raise = exceptions.CosmosHttpResponseError(status_code=http_constants.StatusCodes.NOT_FOUND,
+                                                                 message="Read Session is Not Available")
+                ex_to_raise.sub_status = http_constants.SubStatusCodes.READ_SESSION_NOTAVAILABLE
+                raise ex_to_raise
+        return self.OriginalExecuteFunction(function, *args, **kwargs)
+
+    def _MockExecuteFunctionEndPointRetry(self, function, *args, **kwargs):
+        if args:
+            if args[1].operation_type == 'SqlQuery':
+                ex_to_raise = exceptions.CosmosHttpResponseError(status_code=http_constants.StatusCodes.FORBIDDEN,
+                                                                 message="End Point Discovery")
+                ex_to_raise.sub_status = http_constants.SubStatusCodes.WRITE_FORBIDDEN
+                raise ex_to_raise
+        return self.OriginalExecuteFunction(function, *args, **kwargs)
+
+    def _MockExecuteFunctionTimeoutFailoverRetry(self, function, *args, **kwargs):
+        if args:
+            if args[1].operation_type == 'SqlQuery':
+                ex_to_raise = exceptions.CosmosHttpResponseError(status_code=http_constants.StatusCodes.REQUEST_TIMEOUT,
+                                                                 message="Timeout Failover")
+                raise ex_to_raise
+        return self.OriginalExecuteFunction(function, *args, **kwargs)
+
     def _MockNextFunction(self):
         if self.count < len(self.payloads):
             item, result = self.get_mock_result(self.payloads, self.count)

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

@@ -12,7 +12,7 @@
 import azure.cosmos.aio._retry_utility_async as retry_utility
 import azure.cosmos.exceptions as exceptions
 import test_config
-from azure.cosmos import http_constants
+from azure.cosmos import http_constants, _endpoint_discovery_retry_policy
 from azure.cosmos._execution_context.query_execution_info import _PartitionedQueryExecutionInfo
 from azure.cosmos.aio import CosmosClient, DatabaseProxy, ContainerProxy
 from azure.cosmos.documents import _DistinctType
@@ -936,6 +936,91 @@ async def query_items(database):
         self.client.client_connection.connection_policy.RetryOptions = old_retry
         await self.created_db.delete_container(created_collection.id)
 
+    async def test_query_request_params_none_retry_policy(self):
+        created_collection = await self.created_db.create_container_if_not_exists(
+            id="query_request_params_none_retry_policy_" + str(uuid.uuid4()),
+            partition_key=PartitionKey(path="/pk")
+        )
+        items = [
+            {'id': str(uuid.uuid4()), 'pk': 'test', 'val': 5},
+            {'id': str(uuid.uuid4()), 'pk': 'test', 'val': 5},
+            {'id': str(uuid.uuid4()), 'pk': 'test', 'val': 5}
+        ]
+
+        for item in items:
+            await created_collection.create_item(body=item)
+
+        self.OriginalExecuteFunction = retry_utility.ExecuteFunctionAsync
+        # Test session retry will properly push the exception when retries run out
+        retry_utility.ExecuteFunctionAsync = self._MockExecuteFunctionSessionRetry
+        try:
+            query = "SELECT * FROM c"
+            items = created_collection.query_items(
+                query=query,
+                enable_cross_partition_query=True
+            )
+            fetch_results = [item async for item in items]
+        except exceptions.CosmosHttpResponseError as e:
+            assert e.status_code == 404
+            assert e.sub_status == 1002
+
+        # Test endpoint discovery retry
+        retry_utility.ExecuteFunctionAsync = self._MockExecuteFunctionEndPointRetry
+        _endpoint_discovery_retry_policy.EndpointDiscoveryRetryPolicy.Max_retry_attempt_count = 3
+        _endpoint_discovery_retry_policy.EndpointDiscoveryRetryPolicy.Retry_after_in_milliseconds = 10
+        try:
+            query = "SELECT * FROM c"
+            items = created_collection.query_items(
+                query=query,
+                enable_cross_partition_query=True
+            )
+            fetch_results = [item async for item in items]
+        except exceptions.CosmosHttpResponseError as e:
+            assert e.status_code == http_constants.StatusCodes.FORBIDDEN
+            assert e.sub_status == http_constants.SubStatusCodes.WRITE_FORBIDDEN
+        _endpoint_discovery_retry_policy.EndpointDiscoveryRetryPolicy.Max_retry_attempt_count = 120
+        _endpoint_discovery_retry_policy.EndpointDiscoveryRetryPolicy.Retry_after_in_milliseconds = 1000
+
+        # Finally lets test timeout failover retry
+        retry_utility.ExecuteFunctionAsync = self._MockExecuteFunctionTimeoutFailoverRetry
+        try:
+            query = "SELECT * FROM c"
+            items = created_collection.query_items(
+                query=query,
+                enable_cross_partition_query=True
+            )
+            fetch_results = [item async for item in items]
+        except exceptions.CosmosHttpResponseError as e:
+            assert e.status_code == http_constants.StatusCodes.REQUEST_TIMEOUT
+        retry_utility.ExecuteFunctionAsync = self.OriginalExecuteFunction
+        await self.created_db.delete_container(created_collection.id)
+
+    async def _MockExecuteFunctionSessionRetry(self, function, *args, **kwargs):
+        if args:
+            if args[1].operation_type == 'SqlQuery':
+                ex_to_raise = exceptions.CosmosHttpResponseError(status_code=http_constants.StatusCodes.NOT_FOUND,
+                                                                 message="Read Session is Not Available")
+                ex_to_raise.sub_status = http_constants.SubStatusCodes.READ_SESSION_NOTAVAILABLE
+                raise ex_to_raise
+        return await self.OriginalExecuteFunction(function, *args, **kwargs)
+
+    async def _MockExecuteFunctionEndPointRetry(self, function, *args, **kwargs):
+        if args:
+            if args[1].operation_type == 'SqlQuery':
+                ex_to_raise = exceptions.CosmosHttpResponseError(status_code=http_constants.StatusCodes.FORBIDDEN,
+                                                                 message="End Point Discovery")
+                ex_to_raise.sub_status = http_constants.SubStatusCodes.WRITE_FORBIDDEN
+                raise ex_to_raise
+        return await self.OriginalExecuteFunction(function, *args, **kwargs)
+
+    async def _MockExecuteFunctionTimeoutFailoverRetry(self, function, *args, **kwargs):
+        if args:
+            if args[1].operation_type == 'SqlQuery':
+                ex_to_raise = exceptions.CosmosHttpResponseError(status_code=http_constants.StatusCodes.REQUEST_TIMEOUT,
+                                                                 message="Timeout Failover")
+                raise ex_to_raise
+        return await self.OriginalExecuteFunction(function, *args, **kwargs)
+
 
 if __name__ == '__main__':
     unittest.main()