Allowing CorrelationActivityId to be sent to BE and retrieved from FeedResponse (Azure#26908)

* Allowing CorrelationActivityId to be sent to BE and retrieved from FeedResponse

* Fixing build issue

* Adding an end-to-end unit test from Spark to verify that only one correlationActivityId is used

* Fix typo in comment

* Adding more error details

* Adding additional check for query diagnostics assertion

* Fixed UUID encoding (only used for ActivityId and CorrelatedActivityId - so safe)

* Fixing test issue

* Small logging changes

* Reverting new public ctor in CosmosBulkExecutionOptions

Author: FabianMeiswinkel

eng/versioning/version_client.txt

@@ -82,8 +82,8 @@ com.azure:azure-cosmos;4.25.0;4.26.0-beta.2
 com.azure:azure-cosmos-benchmark;4.0.1-beta.1;4.0.1-beta.1
 com.azure:azure-cosmos-dotnet-benchmark;4.0.1-beta.1;4.0.1-beta.1
 com.azure.cosmos.spark:azure-cosmos-spark_3_2-12;1.0.0-beta.1;1.0.0-beta.1
-com.azure.cosmos.spark:azure-cosmos-spark_3-1_2-12;4.6.0;4.7.0-beta.1
-com.azure.cosmos.spark:azure-cosmos-spark_3-2_2-12;4.6.0;4.7.0-beta.1
+com.azure.cosmos.spark:azure-cosmos-spark_3-1_2-12;4.6.0;4.6.1-beta.1
+com.azure.cosmos.spark:azure-cosmos-spark_3-2_2-12;4.6.0;4.6.1-beta.1
 com.azure:azure-cosmos-encryption;1.0.0-beta.9;1.0.0-beta.10
 com.azure:azure-data-appconfiguration;1.2.5;1.3.0-beta.1
 com.azure:azure-data-appconfiguration-perf;1.0.0-beta.1;1.0.0-beta.1

sdk/cosmos/azure-cosmos-spark_3-1_2-12/CHANGELOG.md

@@ -1,5 +1,5 @@
 ## Release History
-### 4.7.0-beta.1 (Unreleased)
+### 4.6.1-beta.1 (Unreleased)
 
 ### 4.6.0 (2022-01-25)
 #### Key Bug Fixes

sdk/cosmos/azure-cosmos-spark_3-1_2-12/pom.xml

@@ -11,7 +11,7 @@
   </parent>
   <groupId>com.azure.cosmos.spark</groupId>
   <artifactId>azure-cosmos-spark_3-1_2-12</artifactId>
-  <version>4.7.0-beta.1</version> <!-- {x-version-update;com.azure.cosmos.spark:azure-cosmos-spark_3-1_2-12;current} -->
+  <version>4.6.1-beta.1</version> <!-- {x-version-update;com.azure.cosmos.spark:azure-cosmos-spark_3-1_2-12;current} -->
   <packaging>jar</packaging>
   <url>https://github.com/Azure/azure-sdk-for-java/tree/main/sdk/cosmos/azure-cosmos-spark_3-1_2-12</url>
   <name>OLTP Spark 3.1 Connector for Azure Cosmos DB SQL API</name>

sdk/cosmos/azure-cosmos-spark_3-2_2-12/CHANGELOG.md

@@ -1,5 +1,5 @@
 ## Release History
-### 4.7.0-beta.1 (Unreleased)
+### 4.6.1-beta.1 (Unreleased)
 
 ### 4.6.0 (2022-01-25)
 #### Key Bug Fixes

sdk/cosmos/azure-cosmos-spark_3-2_2-12/pom.xml

@@ -11,7 +11,7 @@
   </parent>
   <groupId>com.azure.cosmos.spark</groupId>
   <artifactId>azure-cosmos-spark_3-2_2-12</artifactId>
-  <version>4.7.0-beta.1</version> <!-- {x-version-update;com.azure.cosmos.spark:azure-cosmos-spark_3-2_2-12;current} -->
+  <version>4.6.1-beta.1</version> <!-- {x-version-update;com.azure.cosmos.spark:azure-cosmos-spark_3-2_2-12;current} -->
   <packaging>jar</packaging>
   <url>https://github.com/Azure/azure-sdk-for-java/tree/main/sdk/cosmos/azure-cosmos-spark_3-2_2-12</url>
   <name>OLTP Spark 3.2 Connector for Azure Cosmos DB SQL API</name>

sdk/cosmos/azure-cosmos-spark_3_2-12/src/main/scala/com/azure/cosmos/spark/BulkWriter.scala

@@ -95,7 +95,7 @@ class BulkWriter(container: CosmosAsyncContainer,
   private def initializeOperationContext(): SparkTaskContext = {
     val taskContext = TaskContext.get
 
-    val diagnosticsContext: DiagnosticsContext = DiagnosticsContext(UUID.randomUUID().toString, "BulkWriter")
+    val diagnosticsContext: DiagnosticsContext = DiagnosticsContext(UUID.randomUUID(), "BulkWriter")
 
     if (taskContext != null) {
       val taskDiagnosticsContext = SparkTaskContext(diagnosticsContext.correlationActivityId,

sdk/cosmos/azure-cosmos-spark_3_2-12/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReader.scala

@@ -36,7 +36,10 @@ private case class ItemsPartitionReader
 
   private val containerTargetConfig = CosmosContainerConfig.parseCosmosContainerConfig(config)
   log.logInfo(s"Reading from feed range $feedRange of " +
-    s"container ${containerTargetConfig.database}.${containerTargetConfig.container}")
+    s"container ${containerTargetConfig.database}.${containerTargetConfig.container} - " +
+    s"correlationActivityId ${diagnosticsContext.correlationActivityId}, " +
+    s"query: ${cosmosQuery.toString}")
+
   private val readConfig = CosmosReadConfig.parseCosmosReadConfig(config)
   private val clientCacheItem = CosmosClientCache(
     CosmosClientConfiguration(config, readConfig.forceEventualConsistency),
@@ -63,11 +66,12 @@ private case class ItemsPartitionReader
       val taskContext = TaskContext.get
       assert(taskContext != null)
 
-      val taskDiagnosticsContext = SparkTaskContext(diagnosticsContext.correlationActivityId,
+      val taskDiagnosticsContext = SparkTaskContext(
+        diagnosticsContext.correlationActivityId,
         taskContext.stageId(),
         taskContext.partitionId(),
         taskContext.taskAttemptId(),
-        feedRange.toString + " " + cosmosQuery.toSqlQuerySpec.getQueryText)
+        feedRange.toString + " " + cosmosQuery.toString)
 
       val listener: OperationListener =
         DiagnosticsLoader.getDiagnosticsProvider(diagnosticsConfig).getLogger(this.getClass)
@@ -91,6 +95,14 @@ private case class ItemsPartitionReader
           queryOptions, null, readConfig.maxItemCount)
         // scalastyle:on null
       }
+
+      ImplementationBridgeHelpers
+        .CosmosQueryRequestOptionsHelper
+        .getCosmosQueryRequestOptionsAccessor
+        .setCorrelationActivityId(
+          queryOptions,
+          diagnosticsContext.correlationActivityId)
+
       cosmosAsyncContainer.queryItems(cosmosQuery.toSqlQuerySpec, queryOptions, classOf[ObjectNode])
     },
     readConfig.maxItemCount

sdk/cosmos/azure-cosmos-spark_3_2-12/src/main/scala/com/azure/cosmos/spark/ItemsScan.scala

@@ -103,10 +103,12 @@ private case class ItemsScan(session: SparkSession,
   }
 
   override def createReaderFactory(): PartitionReaderFactory = {
+    val correlationActivityId = UUID.randomUUID()
+    log.logInfo(s"Creating ItemsScan with CorrelationActivityId '${correlationActivityId.toString}' for query '${cosmosQuery.queryText}'")
     ItemsScanPartitionReaderFactory(config,
       schema,
       cosmosQuery,
-      DiagnosticsContext(UUID.randomUUID().toString, cosmosQuery.queryText),
+      DiagnosticsContext(correlationActivityId, cosmosQuery.queryText),
       cosmosClientStateHandle,
       DiagnosticsConfig.parseDiagnosticsConfig(config))
   }

sdk/cosmos/azure-cosmos-spark_3_2-12/src/main/scala/com/azure/cosmos/spark/PointWriter.scala

@@ -53,7 +53,7 @@ class PointWriter(container: CosmosAsyncContainer, cosmosWriteConfig: CosmosWrit
   private val pendingPointWrites = new TrieMap[Future[Unit], Boolean]()
   private val closed = new AtomicBoolean(false)
 
-  private val diagnosticsContext: DiagnosticsContext = DiagnosticsContext(UUID.randomUUID().toString, "PointWriter")
+  private val diagnosticsContext: DiagnosticsContext = DiagnosticsContext(UUID.randomUUID(), "PointWriter")
 
   private  val taskDiagnosticsContext = SparkTaskContext(diagnosticsContext.correlationActivityId,
     taskContext.stageId(),

sdk/cosmos/azure-cosmos-spark_3_2-12/src/main/scala/com/azure/cosmos/spark/diagnostics/DiagnosticsContext.scala

@@ -3,4 +3,6 @@
 
 package com.azure.cosmos.spark.diagnostics
 
-private[spark] case class DiagnosticsContext(correlationActivityId: String, details: String)
+import java.util.UUID
+
+private[spark] case class DiagnosticsContext(correlationActivityId: UUID, details: String)