Request activity (Azure#41013)

* Update _base.py

* Update _cosmos_http_logging_policy.py

* Update CHANGELOG.md

* Update CHANGELOG.md

* update activity id with tests

Added tests to verify Activity id is generated client side. Proof of client generated activity id can be seen from seeing it in the request headers.

* update test

Author: bambriz

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -10,6 +10,7 @@
 * Fixed issue where Query Change Feed did not return items if the container uses legacy Hash V1 Partition Keys. This also fixes issues with not being able to change feed query for Specific Partition Key Values for HPK. See [PR 41270](https://github.com/Azure/azure-sdk-for-python/pull/41270/)
 
 #### Other Changes
+* Added Client Generated Activity IDs to all Requests. Cosmos Diagnostics Logs will more clearly show the Activity ID for each request and response. [PR 41013](https://github.com/Azure/azure-sdk-for-python/pull/41013)
 
 ### 4.12.0b1 (2025-05-19)
 

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

@@ -142,6 +142,8 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
     headers = dict(default_headers)
     options = options or {}
 
+    # Generate a new activity ID for each request client side.
+    headers[http_constants.HttpHeaders.ActivityId] = GenerateGuidId()
     if cosmos_client_connection.UseMultipleWriteLocations:
         headers[http_constants.HttpHeaders.AllowTentativeWrites] = "true"
 

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

@@ -127,14 +127,14 @@ def on_request(
             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')
+            operation_type = http_request.headers.get('x-ms-thinclient-proxy-operation-type', "")
             try:
                 url = request.http_request.url
             except AttributeError:
                 url = None
             database_name = None
             collection_name = None
-            resource_type = http_request.headers.get('x-ms-thinclient-proxy-resource-type')
+            resource_type = http_request.headers.get('x-ms-thinclient-proxy-resource-type', "")
             if url:
                 url_parts = url.split('/')
                 if 'dbs' in url_parts:
@@ -160,9 +160,11 @@ def on_request(
 
                 if filter_applied and 'logger_attributes' in request.context:
                     cosmos_logger_attributes = request.context['logger_attributes']
+                    cosmos_logger_attributes['activity_id'] = http_request.headers.get(HttpHeaders.ActivityId, "")
                     cosmos_logger_attributes['is_request'] = True
                 else:
                     cosmos_logger_attributes = {
+                        'activity_id': http_request.headers.get(HttpHeaders.ActivityId, ""),
                         'duration': None,
                         'status_code': None,
                         'sub_status_code': None,
@@ -181,6 +183,8 @@ def on_request(
                     logger.info(db_settings, extra=cosmos_logger_attributes)
                     logger.info("Request URL: %r", redacted_url, extra=cosmos_logger_attributes)
                     logger.info("Request method: %r", http_request.method, extra=cosmos_logger_attributes)
+                    logger.info("Request Activity ID: %r", http_request.headers.get(HttpHeaders.ActivityId),
+                                extra=cosmos_logger_attributes)
                     logger.info("Request headers:", extra=cosmos_logger_attributes)
                     for header, value in http_request.headers.items():
                         value = self._redact_header(header, value)
@@ -204,6 +208,7 @@ def on_request(
                 log_string += db_settings
                 log_string += "\nRequest URL: '{}'".format(redacted_url)
                 log_string += "\nRequest method: '{}'".format(http_request.method)
+                log_string += "\nRequest Activity ID: {}".format(http_request.headers.get(HttpHeaders.ActivityId))
                 log_string += "\nRequest headers:"
                 for header, value in http_request.headers.items():
                     value = self._redact_header(header, value)
@@ -251,12 +256,13 @@ def on_response(  # pylint: disable=too-many-statements, too-many-branches, too-
                 duration = (time.time() - context["start_time"]) * 1000 \
                     if "start_time" in context else None  # type: ignore[union-attr, arg-type]
 
-            log_data = {"duration": duration,
+            log_data = {"activity_id": http_response.headers.get(HttpHeaders.ActivityId, ""),
+                        "duration": duration,
                         "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'),
+                        "operation_type": headers.get('x-ms-thinclient-proxy-operation-type', ""),
                         "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
+                        "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"]:
                 url_parts: List[str] = log_data["url"].split('/')  # type: ignore[union-attr]

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

@@ -122,7 +122,7 @@ def test_default_http_logging_policy(self):
         messages_response = self.mock_handler_default.messages[1].message.split("\n")
         assert messages_request[1] == "Request method: 'GET'"
         assert 'Request headers:' in messages_request[2]
-        assert messages_request[14] == 'No body was attached to the request'
+        assert messages_request[-1] == 'No body was attached to the request'
         assert messages_response[0] == 'Response status: 200'
         assert 'Response headers:' in messages_response[1]
 
@@ -286,5 +286,59 @@ def test_client_settings(self):
                 locations = get_locations_list(message)
                 assert all_locations == locations
 
+    def test_activity_id_logging_policy(self):
+        # Create a mock handler and logger for the new client
+        self.mock_handler_activity_id = MockHandler()
+        self.logger_activity_id = create_logger("test_logger_activity_id", self.mock_handler_activity_id)
+
+        # Create a new client with the logger and enable diagnostics logging
+        self.client_activity_id = cosmos_client.CosmosClient(
+            self.host,
+            self.masterKey,
+            consistency_level="Session",
+            connection_policy=self.connectionPolicy,
+            logger=self.logger_activity_id,
+            enable_diagnostics_logging=True
+        )
+
+        # Generate a custom activity ID
+        custom_activity_id = str(uuid.uuid4())
+
+        # Create a database and container for the test
+        database_id = "database_test_activity_id_" + str(uuid.uuid4())
+        container_id = "container_test_activity_id_" + str(uuid.uuid4())
+        database = self.client_activity_id.create_database(id=database_id)
+        container = database.create_container(id=container_id, partition_key=PartitionKey(path="/pk"))
+        # Reset the mock handler to clear previous messages
+        self.mock_handler_activity_id.reset()
+
+        # Upsert an item and verify the request and response activity IDs match
+        item_id = str(uuid.uuid4())
+        item_body = {"id": item_id, "pk": item_id}
+        container.upsert_item(body=item_body)
+
+        # Verify that the request activity ID matches the response activity ID
+        # Having the Request Activity confirms we generated one from SDK
+        # Having it match the response means it was passed successfully
+        log_record_request = self.mock_handler_activity_id.messages[0]
+        log_record_response = self.mock_handler_activity_id.messages[1]
+        assert log_record_request.activity_id == log_record_response.activity_id
+
+        # Upsert another item with the custom activity ID in the initial headers
+        headers = {"x-ms-activity-id": custom_activity_id}
+        item_id_2 = str(uuid.uuid4())
+        item_body_2 = {"id": item_id_2, "pk": item_id_2}
+        container.upsert_item(body=item_body_2, initial_headers=headers)
+
+        # Verify that the custom activity ID does not match the request activity ID from the log record
+        # Users should not be able to pass in their own activity ID.
+        log_record_request_2 = self.mock_handler_activity_id.messages[2]
+        assert log_record_request_2.activity_id != custom_activity_id
+
+        # Clean up by deleting the database
+        self.client_activity_id.delete_database(database_id)
+        self.mock_handler_activity_id.reset()
+
+
 if __name__ == "__main__":
     unittest.main()

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

@@ -27,8 +27,6 @@
 class TestCosmosHttpLoggerAsync(unittest.IsolatedAsyncioTestCase):
     mock_handler_diagnostic = None
     mock_handler_default = None
-    logger_diagnostic = None
-    logger_default = None
     config = test_config.TestConfig
     host = config.host
     masterKey = config.masterKey
@@ -47,6 +45,7 @@ async def asyncSetUp(self):
         self.mock_handler_default = MockHandler()
         self.mock_handler_diagnostic = MockHandler()
         self.mock_handler_filtered_diagnostic = MockHandler()
+        self.mock_handler_activity_id = MockHandler()
         # Add filter to the filtered diagnostics handler
 
         self.mock_handler_filtered_diagnostic.addFilter(FilterStatusCode())
@@ -57,6 +56,7 @@ async def asyncSetUp(self):
         self.client_default = None
         self.client_diagnostic = None
         self.client_filtered_diagnostic = None
+        self.client_activity_id = None
 
     async def asyncTearDown(self):
         if self.client_default:
@@ -65,6 +65,8 @@ async def asyncTearDown(self):
             await self.client_diagnostic.close()
         if self.client_filtered_diagnostic:
             await self.client_filtered_diagnostic.close()
+        if self.client_activity_id:
+            await self.client_activity_id.close()
 
     async def test_default_http_logging_policy_async(self):
         self.logger.addHandler(self.mock_handler_default)
@@ -81,7 +83,7 @@ async def test_default_http_logging_policy_async(self):
         messages_response = self.mock_handler_default.messages[1].message.split("\n")
         assert messages_request[1] == "Request method: 'GET'"
         assert 'Request headers:' in messages_request[2]
-        assert messages_request[15] == 'No body was attached to the request'
+        assert messages_request[-1] == 'No body was attached to the request'
         assert messages_response[0] == 'Response status: 200'
         assert 'Response headers:' in messages_response[1]
 
@@ -257,6 +259,58 @@ async def test_client_settings_async(self):
                 assert all_locations == locations
         await initialized_objects["client"].close()
 
+    async def test_activity_id_logging_policy_async(self):
+        # Create a mock handler and logger for the new client
+        self.logger.addHandler(self.mock_handler_activity_id)
+
+        # Create a new client with the logger and enable diagnostics logging
+        self.client_activity_id = cosmos_client.CosmosClient(
+            self.host,
+            self.masterKey,
+            consistency_level="Session",
+            logger=self.logger,
+            enable_diagnostics_logging=True
+        )
+        # Generate a custom activity ID
+        custom_activity_id = str(uuid.uuid4())
+
+        # Create a database and container for the test
+        database_id = "database_test_activity_id_" + str(uuid.uuid4())
+        container_id = "container_test_activity_id_" + str(uuid.uuid4())
+        try:
+
+            database = await self.client_activity_id.create_database(id=database_id)
+            container = await database.create_container(id=container_id, partition_key=PartitionKey(path="/pk"))
+
+            # Reset the mock handler to clear previous messages
+            self.mock_handler_activity_id.reset()
+
+            # Upsert an item and verify the request and response activity IDs match
+            item_id = str(uuid.uuid4())
+            item_body = {"id": item_id, "pk": item_id}
+            await container.upsert_item(body=item_body)
+
+            # Verify that the request activity ID matches the response activity ID
+            log_record_request = self.mock_handler_activity_id.messages[0]
+            log_record_response = self.mock_handler_activity_id.messages[1]
+            assert log_record_request.activity_id == log_record_response.activity_id
+
+            # Upsert another item with the custom activity ID in the initial headers
+            headers = {"x-ms-activity-id": custom_activity_id}
+            item_id_2 = str(uuid.uuid4())
+            item_body_2 = {"id": item_id_2, "pk": item_id_2}
+            await container.upsert_item(body=item_body_2, initial_headers=headers)
+
+            # Verify that the custom activity ID does not match the request activity ID from the log record
+            log_record_request_2 = self.mock_handler_activity_id.messages[2]
+            assert log_record_request_2.activity_id != custom_activity_id
+
+        finally:
+            # Clean up by deleting the database
+            await self.client_activity_id.delete_database(database_id)
+            self.mock_handler_activity_id.reset()
+            self.logger.removeHandler(self.mock_handler_activity_id)
+
 
 if __name__ == "__main__":
     unittest.main()