Cosmos Diagnostics

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,6 +3,7 @@
 ### 4.3.1 (Unreleased)
 
 #### Features Added
+- Added `CosmosHttpLoggingPolicy` to replace `HttpLoggingPolicy` for logging HTTP sessions. 
 
 #### Breaking Changes
 

sdk/cosmos/azure-cosmos/README.md

@@ -602,6 +602,31 @@ even when it isn't enabled for the client:
 database = client.create_database(DATABASE_NAME, logging_enable=True)
 ```
 
+Alternatively, you can log using the CosmosHttpLoggingPolicy, which extends from the azure core HttpLoggingPolicy, by passing in your logger to the `logger` argument.
+By default, it will use the behaviour from HttpLoggingPolicy. The `enable_diagnostics_logging` argument will add additional diagnostic information to the logger.
+```python
+import logging
+from azure.cosmos import CosmosClient
+
+#Create a logger for the 'azure' SDK
+logger = logging.getLogger('azure')
+logger.setLevel(logging.DEBUG)
+
+# Configure a file output
+handler = logging.FileHandler(filename="azure")
+logger.addHandler(handler)
+
+# This client will log diagnostic information from the HTTP session by using the CosmosHttpLoggingPolicy
+client = CosmosClient(URL, credential=KEY, logger=logger, enable_diagnostics_logging=True)
+```
+| Name                    | Policy Flavor    | Parameters                 | Accepted in Init? | Accepted in Request? | Description                                                                   |
+|-------------------------|------------------|----------------------------|----------------|------------------|-------------------------------------------------------------------------------|
+| CosmosHttpLoggingPolicy | SansIOHTTPPolicy |                    |                |                  |                                                                               |
+|       |   |  logger  | X              | X                | If specified, it will be used to log information.                             |
+|       |   | enable_diagnostics_logging | X              | X                | Used to enable logging additional diagnostic information. Defaults to `False` |
+You can learn more about the different policies and how they work [here](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/core/azure-core/CLIENT_LIBRARY_DEVELOPER.md#available-policies)
+
+
 ## Next steps
 
 For more extensive documentation on the Cosmos DB service, see the [Azure Cosmos DB documentation][cosmos_docs] on docs.microsoft.com.

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

@@ -59,6 +59,7 @@
 from . import _utils
 from .partition_key import _Undefined, _Empty
 from ._auth_policy import CosmosBearerTokenCredentialPolicy
+from ._cosmos_http_logging_policy import CosmosHttpLoggingPolicy
 
 ClassType = TypeVar("ClassType")
 # pylint: disable=protected-access
@@ -192,7 +193,7 @@ def __init__(
             CustomHookPolicy(**kwargs),
             NetworkTraceLoggingPolicy(**kwargs),
             DistributedTracingPolicy(**kwargs),
-            HttpLoggingPolicy(**kwargs),
+            CosmosHttpLoggingPolicy(**kwargs),
         ]
 
         transport = kwargs.pop("transport", None)
@@ -211,6 +212,9 @@ def __init__(
 
         # Use database_account if no consistency passed in to verify consistency level to be used
         self._set_client_consistency_level(database_account, consistency_level)
+        # initialize diagnostics
+
+
 
     def _set_client_consistency_level(
             self,
@@ -232,7 +236,6 @@ def _set_client_consistency_level(
         else:
             # Set consistency level header to be used for the client
             self.default_headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level
-
         if consistency_level == documents.ConsistencyLevel.Session:
             # create a session - this is maintained only if the default consistency level
             # on the client is set to session, or if the user explicitly sets it as a property

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

@@ -0,0 +1,154 @@
+import os
+import urllib
+
+from azure.core.pipeline.policies import HttpLoggingPolicy
+import logging
+import types
+
+
+
+class CosmosHttpLoggingPolicy(HttpLoggingPolicy):
+
+    def __init__(self, logger=None, **kwargs):
+        self._enable_diagnostics_logging = kwargs.pop("enable_diagnostics_logging", None)
+        super().__init__(logger, **kwargs)
+        if self._enable_diagnostics_logging:
+            self.logger = logger or logging.getLogger(__name__)
+
+    def on_request(self, request):  # type: (PipelineRequest) -> None
+        http_request = request.http_request
+        options = request.context.options
+        self._enable_diagnostics_logging = request.context.setdefault(
+            "enable_diagnostics_logging",
+            options.pop("enable_diagnostics_logging", self._enable_diagnostics_logging))
+        if self._enable_diagnostics_logging:
+            # 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 not logger.isEnabledFor(logging.INFO):
+                return
+
+            try:
+                parsed_url = list(urllib.parse.urlparse(http_request.url))
+                parsed_qp = urllib.parse.parse_qsl(parsed_url[4], keep_blank_values=True)
+                filtered_qp = [(key, self._redact_query_param(key, value)) for key, value in parsed_qp]
+                # 4 is query
+                parsed_url[4] = "&".join(["=".join(part) for part in filtered_qp])
+                redacted_url = urllib.parse.urlunparse(parsed_url)
+
+                multi_record = os.environ.get(HttpLoggingPolicy.MULTI_RECORD_LOG, False)
+                if multi_record:
+                    logger.info("Request URL: %r", redacted_url)
+                    logger.info("Request method: %r", http_request.method)
+                    logger.info("Request headers:")
+                    for header, value in http_request.headers.items():
+                        #value = self._redact_header(header, value)
+                        logger.info("    %r: %r", header, value)
+                    if isinstance(http_request.body, types.GeneratorType):
+                        logger.info("File upload")
+                        return
+                    try:
+                        if isinstance(http_request.body, types.AsyncGeneratorType):
+                            logger.info("File upload")
+                            return
+                    except AttributeError:
+                        pass
+                    if http_request.body:
+                        logger.info("A body is sent with the request")
+                        return
+                    logger.info("No body was attached to the request")
+                    return
+                log_string = "Request URL: '{}'".format(redacted_url)
+                log_string += "\nRequest method: '{}'".format(http_request.method)
+                log_string += "\nRequest headers:"
+                for header, value in http_request.headers.items():
+                    #value = self._redact_header(header, value)
+                    log_string += "\n    '{}': '{}'".format(header, value)
+                if isinstance(http_request.body, types.GeneratorType):
+                    log_string += "\nFile upload"
+                    logger.info(log_string)
+                    return
+                try:
+                    if isinstance(http_request.body, types.AsyncGeneratorType):
+                        log_string += "\nFile upload"
+                        logger.info(log_string)
+                        return
+                except AttributeError:
+                    pass
+                if http_request.body:
+                    log_string += "\nA body is sent with the request"
+                    logger.info(log_string)
+                    return
+                log_string += "\nNo body was attached to the request"
+                logger.info(log_string)
+
+            except Exception as err:  # pylint: disable=broad-except
+                logger.warning("Failed to log request: %s", repr(err))
+        else:
+            super().on_request(request)
+
+
+    def on_response(self, request, response):  # type: (PipelineRequest, PipelineResponse) -> None
+        if self._enable_diagnostics_logging:
+            http_response = response.http_response
+            ir = http_response.internal_response
+            try:
+                logger = response.context["logger"]
+
+                if not logger.isEnabledFor(logging.INFO):
+                    return
+
+                multi_record = os.environ.get(HttpLoggingPolicy.MULTI_RECORD_LOG, False)
+                if multi_record:
+                    logger.info("Response status: %r", http_response.status_code)
+                    logger.info("Response status reason: %r", ir.reason)
+                    try:
+                        logger.info("Elapsed Time: %r", ir.elapsed)
+                    except AttributeError:
+                        logger.info("Elapsed Time: %r", None)
+                    if http_response.status_code >= 400:
+                        sm = ir.text
+                        sm.replace("true", "True")
+                        sm.replace("false", "False")
+                        temp_sm = eval(sm)
+                        temp_sm['message'] = temp_sm['message'].replace("\r", " ")
+                        logger.into("Response status error message: %r", temp_sm['message'])
+                    logger.info("Response headers:")
+                    for res_header, value in http_response.headers.items():
+                        logger.info("    %r: %r", res_header, value)
+                    return
+                log_string = "Response status: {}".format(http_response.status_code)
+                log_string += "\nResponse status reason: {}".format(ir.reason)
+                try:
+                    log_string += "\nElapsed Time: {}".format(ir.elapsed)
+                except AttributeError:
+                    log_string += "\nElapsed Time: {}".format(None)
+                if http_response.status_code >= 400:
+                    sm = ir.text
+                    sm.replace("true", "True")
+                    sm.replace("false", "False")
+                    temp_sm = eval(sm)
+                    temp_sm['message'] = temp_sm['message'].replace("\r", " ")
+                    log_string += "\nResponse status error message: {}".format(temp_sm['message'])
+                log_string += "\nResponse headers:"
+                for res_header, value in http_response.headers.items():
+                    log_string += "\n    '{}': '{}'".format(res_header, value)
+                logger.info(log_string)
+            except Exception as err:  # pylint: disable=broad-except
+                logger.warning("Failed to log response: %s", repr(err))
+        else:
+            super().on_response(request, response)
+
+        # # Current System info
+        # def get_system_info(self):
+        #     ret = {}
+        #     ret["system"] = platform.system()
+        #     ret["python version"] = platform.python_version()
+        #     ret["architecture"] = platform.architecture()
+        #     ret["cpu"] = platform.processor()
+        #     ret["cpu count"] = os.cpu_count()
+        #     ret["machine"] = platform.machine()
+        #
+        #     return ret

sdk/cosmos/azure-cosmos/azure/cosmos/_retry_utility.py

@@ -69,6 +69,9 @@ def Execute(client, global_endpoint_manager, function, *args, **kwargs):
     )
     partition_key_range_gone_retry_policy = _gone_retry_policy.PartitionKeyRangeGoneRetryPolicy(client, *args)
 
+    #on a new execution, erase last exception from client, currently always ovverides even if there is exception
+    # client.last_exceptions = None
+
     while True:
         try:
             client_timeout = kwargs.get('timeout')
@@ -87,10 +90,10 @@ def Execute(client, global_endpoint_manager, function, *args, **kwargs):
             client.last_response_headers[
                 HttpHeaders.ThrottleRetryWaitTimeInMs
             ] = resourceThrottle_retry_policy.cummulative_wait_time_in_milliseconds
-
             return result
         except exceptions.CosmosHttpResponseError as e:
             retry_policy = None
+
             if e.status_code == StatusCodes.FORBIDDEN and e.sub_status == SubStatusCodes.WRITE_FORBIDDEN:
                 retry_policy = endpointDiscovery_retry_policy
             elif e.status_code == StatusCodes.TOO_MANY_REQUESTS:

sdk/cosmos/azure-cosmos/azure/cosmos/_synchronized_request.py

@@ -140,10 +140,10 @@ def _Request(global_endpoint_manager, request_params, connection_policy, pipelin
     response = response.http_response
     headers = dict(response.headers)
 
+
     data = response.body()
     if data:
         data = data.decode("utf-8")
-
     if response.status_code == 404:
         raise exceptions.CosmosResourceNotFoundError(message=data, response=response)
     if response.status_code == 409:

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

@@ -105,14 +105,12 @@ async def _Request(global_endpoint_manager, request_params, connection_policy, p
             connection_verify=kwargs.pop("connection_verify", is_ssl_enabled),
             **kwargs
         )
-
     response = response.http_response
     headers = dict(response.headers)
 
     data = response.body()
     if data:
         data = data.decode("utf-8")
-
     if response.status_code == 404:
         raise exceptions.CosmosResourceNotFoundError(message=data, response=response)
     if response.status_code == 409:

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

@@ -394,3 +394,4 @@ async def _get_database_account(self, **kwargs: Any) -> DatabaseAccount:
         if response_hook:
             response_hook(self.client_connection.last_response_headers)
         return result
+

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

@@ -59,6 +59,7 @@
 from .. import _utils
 from ..partition_key import _Undefined, _Empty
 from ._auth_policy_async import AsyncCosmosBearerTokenCredentialPolicy
+from .._cosmos_http_logging_policy import CosmosHttpLoggingPolicy
 
 ClassType = TypeVar("ClassType")
 # pylint: disable=protected-access
@@ -145,6 +146,7 @@ def __init__(
         if consistency_level is not None:
             self.default_headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level
 
+
         # Keeps the latest response headers from the server.
         self.last_response_headers = None
 
@@ -195,7 +197,7 @@ def __init__(
             CustomHookPolicy(**kwargs),
             NetworkTraceLoggingPolicy(**kwargs),
             DistributedTracingPolicy(**kwargs),
-            HttpLoggingPolicy(**kwargs),
+            CosmosHttpLoggingPolicy(**kwargs),
             ]
 
         transport = kwargs.pop("transport", None)
@@ -251,7 +253,6 @@ async def _setup(self):
         else:
             # Use database_account if no consistency passed in to verify consistency level to be used
             user_defined_consistency = self._check_if_account_session_consistency(database_account)
-
         if user_defined_consistency == documents.ConsistencyLevel.Session:
             # create a Session if the user wants Session consistency
             self.session = _session.Session(self.url_connection)
@@ -276,7 +277,6 @@ def _check_if_account_session_consistency(
             # We only set the header if we're using session consistency in the account in order to keep
             # the current update_session logic which uses the header
             self.default_headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level
-
         return consistency_level
 
 

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

@@ -783,3 +783,4 @@ async def replace_throughput(self, throughput: int, **kwargs: Any) -> Throughput
         if response_hook:
             response_hook(self.client_connection.last_response_headers, data)
         return ThroughputProperties(offer_throughput=data["content"]["offerThroughput"], properties=data)
+

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

@@ -316,3 +316,4 @@ async def delete_permission(self, permission: Union[str, Dict[str, Any], Permiss
 
         if response_hook:
             response_hook(self.client_connection.last_response_headers, result)
+

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

@@ -704,6 +704,8 @@ def replace_throughput(self, throughput, **kwargs):
 
         if response_hook:
             response_hook(self.client_connection.last_response_headers, data)
+        # self.diagnostics.update_diagnostics(self.client_connection.last_response_headers, data,
+        #                                     self.client_connection.last_exceptions)
 
         return ThroughputProperties(offer_throughput=data["content"]["offerThroughput"], properties=data)
 
@@ -727,6 +729,8 @@ def list_conflicts(self, max_item_count=None, **kwargs):
         )
         if response_hook:
             response_hook(self.client_connection.last_response_headers, result)
+        # self.diagnostics.update_diagnostics(self.client_connection.last_response_headers, result,
+        #                                     self.client_connection.last_exceptions)
         return result
 
     @distributed_trace
@@ -794,6 +798,8 @@ def get_conflict(self, conflict, partition_key, **kwargs):
         )
         if response_hook:
             response_hook(self.client_connection.last_response_headers, result)
+        # self.diagnostics.update_diagnostics(self.client_connection.last_response_headers, result,
+        #                                     self.client_connection.last_exceptions)
         return result
 
     @distributed_trace
@@ -820,3 +826,4 @@ def delete_conflict(self, conflict, partition_key, **kwargs):
         )
         if response_hook:
             response_hook(self.client_connection.last_response_headers, result)
+