Adds Correlated Activity Id (Azure#27820)

* correlated act id

* Update _base.py

* Create test_correlated_activity_id.py

* Update CHANGELOG.md

* Update container.py

* Update _container.py

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

Co-authored-by: Simon Moreno <30335873+simorenoh@users.noreply.github.com>

* Update container.py

* Update test_correlated_activity_id.py

* cspell

* Update http_constants.py

* Update container.py

* Update test_correlated_activity_id.py

* Update container.py

* Update _container.py

* Update _base.py

* minor changes

* Update _container.py

Co-authored-by: Simon Moreno <30335873+simorenoh@users.noreply.github.com>

Author: JericHunter

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -1,9 +1,8 @@
 ## Release History
 
 ### 4.3.1b2 (Unreleased)
-
 #### Features Added
-
+ - Added `correlated_activity_id` for query operations
 #### Breaking Changes
 
 #### Bugs Fixed

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

@@ -299,6 +299,9 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
     if options.get("autoUpgradePolicy"):
         headers[http_constants.HttpHeaders.AutoscaleSettings] = options["autoUpgradePolicy"]
 
+    if options.get("correlatedActivityId"):
+        headers[http_constants.HttpHeaders.CorrelatedActivityId] = options["correlatedActivityId"]
+
     return headers
 
 

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

@@ -30,7 +30,7 @@
 
 from ._cosmos_client_connection_async import CosmosClientConnection
 from .._base import build_options as _build_options, validate_cache_staleness_value, _deserialize_throughput, \
-    _replace_throughput
+    _replace_throughput, GenerateGuidId
 from ..exceptions import CosmosResourceNotFoundError
 from ..http_constants import StatusCodes
 from ..offer import ThroughputProperties
@@ -364,7 +364,8 @@ def query_items(
         if max_integrated_cache_staleness_in_ms:
             validate_cache_staleness_value(max_integrated_cache_staleness_in_ms)
             feed_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms
-
+        correlated_activity_id = GenerateGuidId()
+        feed_options["correlatedActivityId"] = correlated_activity_id
         if hasattr(response_hook, "clear"):
             response_hook.clear()
 

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

@@ -29,7 +29,8 @@
 from azure.core.tracing.decorator import distributed_trace  # type: ignore
 
 from ._cosmos_client_connection import CosmosClientConnection
-from ._base import build_options, validate_cache_staleness_value, _deserialize_throughput, _replace_throughput
+from ._base import build_options, validate_cache_staleness_value, _deserialize_throughput, _replace_throughput, \
+    GenerateGuidId
 from .exceptions import CosmosResourceNotFoundError
 from .http_constants import StatusCodes
 from .offer import ThroughputProperties
@@ -396,7 +397,8 @@ def query_items(
         if max_integrated_cache_staleness_in_ms:
             validate_cache_staleness_value(max_integrated_cache_staleness_in_ms)
             feed_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms
-
+        correlated_activity_id = GenerateGuidId()
+        feed_options["correlatedActivityId"] = correlated_activity_id
         if hasattr(response_hook, "clear"):
             response_hook.clear()
 

sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py

@@ -99,6 +99,7 @@ class HttpHeaders(object):
 
     # Request sender generated. Simply echoed by backend.
     ActivityId = "x-ms-activity-id"
+    CorrelatedActivityId = "x-ms-cosmos-correlated-activityid"  # cspell:disable-line
     PreTriggerInclude = "x-ms-documentdb-pre-trigger-include"
     PreTriggerExclude = "x-ms-documentdb-pre-trigger-exclude"
     PostTriggerInclude = "x-ms-documentdb-post-trigger-include"

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

@@ -0,0 +1,65 @@
+# The MIT License (MIT)
+# Copyright (c) 2022 Microsoft Corporation
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import unittest
+from unittest.mock import MagicMock
+
+import pytest
+
+import azure.cosmos.cosmos_client as cosmos_client
+import test_config
+from azure.cosmos.partition_key import PartitionKey
+
+pytestmark = pytest.mark.cosmosEmulator
+
+
+@pytest.mark.usefixtures("teardown")
+class CorrelatedActivityIdTest(unittest.TestCase):
+    configs = test_config._test_config
+    host = configs.host
+    masterKey = configs.masterKey
+
+    @classmethod
+    def setUpClass(cls):
+        cls.client = cosmos_client.CosmosClient(cls.host, cls.masterKey)
+        cls.database = cls.client.create_database_if_not_exists(test_config._test_config.TEST_DATABASE_ID)
+        cls.container = cls.database.create_container(id=test_config._test_config.TEST_COLLECTION_MULTI_PARTITION_ID,
+                                                      partition_key=PartitionKey(path="/id"))
+
+    def side_effect_correlated_activity_id(self, *args):
+        # Extract request headers from args
+        assert args[2]["x-ms-cosmos-correlated-activityid"]  # cspell:disable-line
+        raise StopIteration
+
+    def test_correlated_activity_id(self):
+        query = 'SELECT * from c ORDER BY c._ts'
+
+        cosmos_client_connection = self.container.client_connection
+        cosmos_client_connection._CosmosClientConnection__Get = MagicMock(
+            side_effect=self.side_effect_correlated_activity_id)
+        try:
+            self.container.query_items(query=query, partition_key="pk-1")
+        except StopIteration:
+            pass
+
+
+if __name__ == "__main__":
+    unittest.main()