Fix for formatting empty, malformed, or non json payload in error messages (#40889)

* Accomdate diagnsotis format error for empty or malformed payload

In some cases when we get an exception the payload is not in a json format or is empty in which case we accommodate for that when formatting the error message.

* Add fix for async client with logging, and setting defaults that are not none for open telemetry

* Update _cosmos_http_logging_policy.py

* Update _cosmos_http_logging_policy.py

* Update _cosmos_http_logging_policy.py

* Update CHANGELOG.md

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

Co-authored-by: Simon Moreno <[email protected]>

* Update _cosmos_http_logging_policy.py

* Update test_cosmos_http_logging_policy_async.py

fixes test that were broken for async tests of http logging policy

* Update test_cosmos_http_logging_policy_async.py

* Update test_cosmos_http_logging_policy.py

* Update test_cosmos_http_logging_policy_async.py

---------

Co-authored-by: tvaron3 <[email protected]>
Co-authored-by: Simon Moreno <[email protected]>

Author: 3 people

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed Diagnostics Error Log Formatting to handle error messages from non-CosmosHttpResponseExceptions. See [PR 40889](https://github.com/Azure/azure-sdk-for-python/pull/40889/files)
 * Fixed bug where `multiple_write_locations` option in client was not being honored. See [PR 40999](https://github.com/Azure/azure-sdk-for-python/pull/40999).
 
 #### Other Changes

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_http_logging_policy.py

@@ -37,7 +37,7 @@
 from ._location_cache import LocationCache
 from .http_constants import HttpHeaders
 from ._global_endpoint_manager import _GlobalEndpointManager
-from .documents import DatabaseAccount, ConnectionPolicy
+from .documents import DatabaseAccount
 
 if TYPE_CHECKING:
     from azure.core.rest import HttpRequest, HttpResponse, AsyncHttpResponse
@@ -55,10 +55,16 @@
 
 
 def _format_error(payload: str) -> str:
-    output = json.loads(payload)
-    ret_str = "\n\t" + "Code: " + output['code'] + "\n"
-    message = output["message"].replace("\r\n", "\n\t\t").replace(",", ",\n\t\t")
-    ret_str += "\t" + message + "\n"
+    try:
+        output = json.loads(payload)
+        ret_str = "\n\t" + "Code: " + output['code'] + "\n"
+        message = output["message"].replace("\r\n", "\n\t\t").replace(",", ",\n\t\t")
+        ret_str += "\t" + message + "\n"
+    except (json.JSONDecodeError, KeyError):
+        try:
+            ret_str = "\t" + payload.replace("\r\n", "\n\t\t").replace(",", ",\n\t\t") + "\n"
+        except AttributeError:
+            ret_str = str(payload)
     return ret_str
 
 
@@ -109,15 +115,16 @@ def on_request(
         if self._enable_diagnostics_logging:
 
             http_request = request.http_request
-            request.context["start_time"] = time.time()
+            if "start_time" not in request.context:
+                request.context["start_time"] = time.time()
             options = request.context.options
             # Get logger in my context first (request has been retried)
             # then read from kwargs (pop if that's the case)
             # then use my instance logger
             logger = request.context.setdefault("logger", options.pop("logger", self.logger))
             # If filtered is applied, and we are not calling on request from on response, just return to avoid logging
             # the request again
-            filter_applied = (logger.filters) or any(bool(h.filters) for h in logger.handlers)
+            filter_applied = bool(logger.filters) or any(bool(h.filters) for h in logger.handlers)
             if filter_applied and 'logger_attributes' not in request.context:
                 return
             operation_type = http_request.headers.get('x-ms-thinclient-proxy-operation-type')
@@ -236,7 +243,7 @@ def on_response(  # pylint: disable=too-many-statements, too-many-branches, too-
             http_response = response.http_response
             headers = request.http_request.headers
             sub_status_str = http_response.headers.get("x-ms-substatus")
-            sub_status_code: Optional[int] = int(sub_status_str) if sub_status_str else None
+            sub_status_code: Optional[int] = int(sub_status_str) if sub_status_str else 0
             url_obj = request.http_request.url  # type: ignore[attr-defined, union-attr]
             try:
                 duration: Optional[float] = float(http_response.headers.get("x-ms-request-duration-ms"))  # type: ignore[union-attr, arg-type]  # pylint: disable=line-too-long
@@ -248,7 +255,7 @@ def on_response(  # pylint: disable=too-many-statements, too-many-branches, too-
                         "status_code": http_response.status_code, "sub_status_code": sub_status_code,
                         "verb": request.http_request.method,
                         "operation_type": headers.get('x-ms-thinclient-proxy-operation-type'),
-                        "url": str(url_obj), "database_name": None, "collection_name": None,
+                        "url": str(url_obj), "database_name": "", "collection_name": "",
                         "resource_type": headers.get('x-ms-thinclient-proxy-resource-type'), "is_request": False}  # type: ignore[assignment]  # pylint: disable=line-too-long
 
             if log_data["url"]:
@@ -316,10 +323,10 @@ def __get_client_settings(self) -> Optional[Dict[str, Any]]:
         client_excluded_regions = []
         client_account_read_regions = []
         client_account_write_regions = []
-
-        if self.__global_endpoint_manager:
-            if self.__global_endpoint_manager.Client and self.__global_endpoint_manager.Client.connection_policy:
-                connection_policy: ConnectionPolicy = self.__global_endpoint_manager.Client.connection_policy
+        if self.__global_endpoint_manager and hasattr(self.__global_endpoint_manager, 'client'):
+            gem_client = self.__global_endpoint_manager.client
+            if gem_client and gem_client.connection_policy:
+                connection_policy = gem_client.connection_policy
                 client_preferred_regions = connection_policy.PreferredLocations
                 client_excluded_regions = connection_policy.ExcludedLocations
 
@@ -328,10 +335,10 @@ def __get_client_settings(self) -> Optional[Dict[str, Any]]:
                 client_account_read_regions = location_cache.account_read_locations
                 client_account_write_regions = location_cache.account_write_locations
 
-        return {"Client Preferred Regions": client_preferred_regions,
-                "Client Excluded Regions": client_excluded_regions,
-                "Client Account Read Regions": client_account_read_regions,
-                "Client Account Write Regions": client_account_write_regions}
+        return {"Preferred Regions": client_preferred_regions,
+                "Excluded Regions": client_excluded_regions,
+                "Account Read Regions": client_account_read_regions,
+                "Account Write Regions": client_account_write_regions}
 
     def __get_database_account_settings(self) -> Optional[DatabaseAccount]:
         if self.__global_endpoint_manager and hasattr(self.__global_endpoint_manager, '_database_account_cache'):

sdk/cosmos/azure-cosmos/azure/cosmos/_global_endpoint_manager.py

@@ -44,7 +44,7 @@ class _GlobalEndpointManager(object): # pylint: disable=too-many-instance-attrib
     """
 
     def __init__(self, client):
-        self.Client = client
+        self.client = client
         self.EnableEndpointDiscovery = client.connection_policy.EnableEndpointDiscovery
         self.PreferredLocations = client.connection_policy.PreferredLocations
         self.DefaultEndpoint = client.url_connection
@@ -170,28 +170,28 @@ def _endpoints_health_check(self, **kwargs):
                 if success_count >= 4:
                     break
                 # save current dba timeouts
-                previous_dba_read_timeout = self.Client.connection_policy.DBAReadTimeout
-                previous_dba_connection_timeout = self.Client.connection_policy.DBAConnectionTimeout
+                previous_dba_read_timeout = self.client.connection_policy.DBAReadTimeout
+                previous_dba_connection_timeout = self.client.connection_policy.DBAConnectionTimeout
                 try:
                     if (endpoint in
                             self.location_cache.location_unavailability_info_by_endpoint):
                         # if the endpoint is unavailable, we need to lower the timeouts to be more aggressive in the
                         # health check. This helps reduce the time the health check is blocking all requests.
-                        self.Client.connection_policy.override_dba_timeouts(constants._Constants
+                        self.client.connection_policy.override_dba_timeouts(constants._Constants
                                                                             .UnavailableEndpointDBATimeouts,
                                                                             constants._Constants
                                                                             .UnavailableEndpointDBATimeouts)
-                        self.Client._GetDatabaseAccountCheck(endpoint, **kwargs)
+                        self.client._GetDatabaseAccountCheck(endpoint, **kwargs)
                     else:
-                        self.Client._GetDatabaseAccountCheck(endpoint, **kwargs)
+                        self.client._GetDatabaseAccountCheck(endpoint, **kwargs)
                     success_count += 1
                     self.location_cache.mark_endpoint_available(endpoint)
                 except (exceptions.CosmosHttpResponseError, AzureError):
                     self.mark_endpoint_unavailable_for_read(endpoint, False)
                     self.mark_endpoint_unavailable_for_write(endpoint, False)
                 finally:
                     # after the health check for that endpoint setting the timeouts back to their original values
-                    self.Client.connection_policy.override_dba_timeouts(previous_dba_read_timeout,
+                    self.client.connection_policy.override_dba_timeouts(previous_dba_read_timeout,
                                                                         previous_dba_connection_timeout)
         self.location_cache.update_location_cache()
 
@@ -204,20 +204,20 @@ def _GetDatabaseAccountStub(self, endpoint, **kwargs):
         :rtype: ~azure.cosmos.DatabaseAccount
         """
         if endpoint in self.location_cache.location_unavailability_info_by_endpoint:
-            previous_dba_read_timeout = self.Client.connection_policy.DBAReadTimeout
-            previous_dba_connection_timeout = self.Client.connection_policy.DBAConnectionTimeout
+            previous_dba_read_timeout = self.client.connection_policy.DBAReadTimeout
+            previous_dba_connection_timeout = self.client.connection_policy.DBAConnectionTimeout
             try:
                 # if the endpoint is unavailable, we need to lower the timeouts to be more aggressive in the
                 # health check. This helps reduce the time the health check is blocking all requests.
-                self.Client.connection_policy.override_dba_timeouts(constants._Constants
+                self.client.connection_policy.override_dba_timeouts(constants._Constants
                                                                     .UnavailableEndpointDBATimeouts,
                                                                     constants._Constants
                                                                     .UnavailableEndpointDBATimeouts)
-                database_account = self.Client.GetDatabaseAccount(endpoint, **kwargs)
+                database_account = self.client.GetDatabaseAccount(endpoint, **kwargs)
             finally:
                 # after the health check for that endpoint setting the timeouts back to their original values
-                self.Client.connection_policy.override_dba_timeouts(previous_dba_read_timeout,
+                self.client.connection_policy.override_dba_timeouts(previous_dba_read_timeout,
                                                                     previous_dba_connection_timeout)
         else:
-            database_account = self.Client.GetDatabaseAccount(endpoint, **kwargs)
+            database_account = self.client.GetDatabaseAccount(endpoint, **kwargs)
         return database_account

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

@@ -132,7 +132,7 @@ def test_default_http_logging_policy(self):
         self.client_default.delete_database(database_id)
 
     def test_cosmos_http_logging_policy(self):
-        # Test if we can log into from reading a database
+        # Test if we can log info from reading a database
         database_id = "database_test-" + str(uuid.uuid4())
         self.client_diagnostic.create_database(id=database_id)
         assert all(m.levelname == 'INFO' for m in self.mock_handler_diagnostic.messages)
@@ -185,8 +185,8 @@ def test_cosmos_http_logging_policy(self):
 
     def test_filtered_diagnostics_logging_policy(self):
         # Test if we can log errors with the filtered diagnostics logger
-        database_id = "database_test-" + str(uuid.uuid4())
-        container_id = "container_test-" + str(uuid.uuid4())
+        database_id = "database_test_" + str(uuid.uuid4())
+        container_id = "diagnostics_container_test_" + str(uuid.uuid4())
         self.client_filtered_diagnostic.create_database(id=database_id)
         database = self.client_filtered_diagnostic.get_database_client(database_id)
         database.create_container(id=container_id, partition_key=PartitionKey(path="/pk"))
@@ -273,16 +273,16 @@ def test_client_settings(self):
         # Verify endpoint locations
         messages_split = mock_handler.messages[1].message.split("\n")
         for message in messages_split:
-            if "Client Preferred Regions:" in message:
+            if "Preferred Regions:" in message:
                 locations = get_locations_list(message)
                 assert all_locations == locations
-            elif "Client Excluded Regions:" in message:
+            elif "Excluded Regions:" in message:
                 locations = get_locations_list(message)
                 assert client_excluded_locations == locations
-            elif "Client Account Read Regions:" in message:
+            elif "Account Read Regions:" in message:
                 locations = get_locations_list(message)
                 assert all_locations == locations
-            elif "Client Account Write Regions:" in message:
+            elif "Account Write Regions:" in message:
                 locations = get_locations_list(message)
                 assert all_locations == locations