Fix transactional batch diagnostics prefix, simplify validation, rename test fields

Author: TheovanKraay

sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransactionalBulkWriter.scala

@@ -84,11 +84,8 @@ private class TransactionalBulkWriter
     case None => 2 * ContainerFeedRangesCache.getFeedRanges(container, containerConfig.feedRangeRefreshIntervalInSecondsOpt).block().size
   }
   // Validate write strategy for transactional batches
-  if (writeConfig.itemWriteStrategy != ItemWriteStrategy.ItemOverwrite) {
-    throw new IllegalArgumentException(
-      s"Transactional batches only support ItemOverwrite (upsert) write strategy. " +
-      s"Requested strategy: ${writeConfig.itemWriteStrategy}")
-  }
+  require(writeConfig.itemWriteStrategy == ItemWriteStrategy.ItemOverwrite,
+    s"Transactional batches only support ItemOverwrite (upsert) write strategy. Requested: ${writeConfig.itemWriteStrategy}")
 
   log.logInfo(
     s"TransactionalBulkWriter instantiated (Host CPU count: $cpuCount, maxPendingBatches: $maxPendingBatches, " +
@@ -186,7 +183,7 @@ private class TransactionalBulkWriter
   private def initializeOperationContext(): SparkTaskContext = {
     val taskContext = TaskContext.get
 
-    val diagnosticsContext: DiagnosticsContext = DiagnosticsContext(UUIDs.nonBlockingRandomUUID(), "BulkWriter")
+    val diagnosticsContext: DiagnosticsContext = DiagnosticsContext(UUIDs.nonBlockingRandomUUID(), "transactional-BulkWriter")
 
     if (taskContext != null) {
       val taskDiagnosticsContext = SparkTaskContext(diagnosticsContext.correlationActivityId,
@@ -406,7 +403,7 @@ private class TransactionalBulkWriter
                                     operationContext: OperationContext): Unit = {
       activeTasks.incrementAndGet()
       if (operationContext.attemptNumber > 1) {
-        logInfoOrWarning(s"bulk scheduleWrite attemptCnt: ${operationContext.attemptNumber}, " +
+        logInfoOrWarning(s"TransactionalBulkWriter scheduleWrite attemptCnt: ${operationContext.attemptNumber}, " +
               s"Context: ${operationContext.toString} $getThreadInfo")
       }
 

sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransactionalBatchITest.scala

@@ -435,8 +435,8 @@ class TransactionalBatchITest extends IntegrationSpec
       new com.azure.cosmos.models.PartitionKeyDefinition()
     )
     val paths = new java.util.ArrayList[String]()
-    paths.add("/PermId")
-    paths.add("/SourceId")
+    paths.add("/testPrimaryKey")
+    paths.add("/testSecondaryKey")
     containerProperties.getPartitionKeyDefinition.setPaths(paths)
     containerProperties.getPartitionKeyDefinition.setKind(com.azure.cosmos.models.PartitionKind.MULTI_HASH)
     containerProperties.getPartitionKeyDefinition.setVersion(com.azure.cosmos.models.PartitionKeyDefinitionVersion.V2)
@@ -452,9 +452,9 @@ class TransactionalBatchITest extends IntegrationSpec
       // Create batch operations with hierarchical partition keys
       val schema = StructType(Seq(
         StructField("id", StringType, nullable = false),
-        StructField("PermId", StringType, nullable = false),
-        StructField("SourceId", StringType, nullable = false),
-        StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false)
+        StructField("testPrimaryKey", StringType, nullable = false),
+        StructField("testSecondaryKey", StringType, nullable = false),
+        StructField("testPrice", org.apache.spark.sql.types.DoubleType, nullable = false)
       ))
 
       val batchOperations = Seq(
@@ -479,11 +479,11 @@ class TransactionalBatchITest extends IntegrationSpec
       val pk = new PartitionKeyBuilder().add(permId).add(sourceId).build()
       val item1 = container.readItem(item1Id, pk, classOf[ObjectNode]).block()
       item1 should not be null
-      item1.getItem.get("price").asDouble() shouldEqual 100.5
+      item1.getItem.get("testPrice").asDouble() shouldEqual 100.5
 
       val item2 = container.readItem(item2Id, pk, classOf[ObjectNode]).block()
       item2 should not be null
-      item2.getItem.get("price").asDouble() shouldEqual 101.25
+      item2.getItem.get("testPrice").asDouble() shouldEqual 101.25
     } finally {
       // Clean up container
       container.delete().block()
@@ -501,8 +501,8 @@ class TransactionalBatchITest extends IntegrationSpec
       new com.azure.cosmos.models.PartitionKeyDefinition()
     )
     val paths = new java.util.ArrayList[String]()
-    paths.add("/PermId")
-    paths.add("/SourceId")
+    paths.add("/testPrimaryKey")
+    paths.add("/testSecondaryKey")
     containerProperties.getPartitionKeyDefinition.setPaths(paths)
     containerProperties.getPartitionKeyDefinition.setKind(com.azure.cosmos.models.PartitionKind.MULTI_HASH)
     containerProperties.getPartitionKeyDefinition.setVersion(com.azure.cosmos.models.PartitionKeyDefinitionVersion.V2)
@@ -518,18 +518,18 @@ class TransactionalBatchITest extends IntegrationSpec
       // First, create initial record
       val initialDoc = Utils.getSimpleObjectMapper.createObjectNode()
       initialDoc.put("id", oldRecordId)
-      initialDoc.put("PermId", permId)
-      initialDoc.put("SourceId", sourceId)
-      initialDoc.put("price", 100.0)
+      initialDoc.put("testPrimaryKey", permId)
+      initialDoc.put("testSecondaryKey", sourceId)
+      initialDoc.put("testPrice", 100.0)
       initialDoc.putNull("valid_to")
       container.createItem(initialDoc).block()
 
       // Now perform atomic temporal update: close old record + create new record
       val schema = StructType(Seq(
         StructField("id", StringType, nullable = false),
-        StructField("PermId", StringType, nullable = false),
-        StructField("SourceId", StringType, nullable = false),
-        StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false),
+        StructField("testPrimaryKey", StringType, nullable = false),
+        StructField("testSecondaryKey", StringType, nullable = false),
+        StructField("testPrice", org.apache.spark.sql.types.DoubleType, nullable = false),
         StructField("valid_to", StringType, nullable = true)
       ))
 
@@ -562,7 +562,7 @@ class TransactionalBatchITest extends IntegrationSpec
       // Verify new record was created
       val newRecord = container.readItem(newRecordId, pk, classOf[ObjectNode]).block()
       newRecord should not be null
-      newRecord.getItem.get("price").asDouble() shouldEqual 150.0
+      newRecord.getItem.get("testPrice").asDouble() shouldEqual 150.0
       newRecord.getItem.get("valid_to").isNull shouldBe true
     } finally {
       // Clean up container
@@ -581,8 +581,8 @@ class TransactionalBatchITest extends IntegrationSpec
       new com.azure.cosmos.models.PartitionKeyDefinition()
     )
     val paths = new java.util.ArrayList[String]()
-    paths.add("/PermId")
-    paths.add("/SourceId")
+    paths.add("/testPrimaryKey")
+    paths.add("/testSecondaryKey")
     containerProperties.getPartitionKeyDefinition.setPaths(paths)
     containerProperties.getPartitionKeyDefinition.setKind(com.azure.cosmos.models.PartitionKind.MULTI_HASH)
     containerProperties.getPartitionKeyDefinition.setVersion(com.azure.cosmos.models.PartitionKeyDefinitionVersion.V2)
@@ -593,9 +593,9 @@ class TransactionalBatchITest extends IntegrationSpec
       // Create operations for different PermId/SourceId combinations
       val schema = StructType(Seq(
         StructField("id", StringType, nullable = false),
-        StructField("PermId", StringType, nullable = false),
-        StructField("SourceId", StringType, nullable = false),
-        StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false)
+        StructField("testPrimaryKey", StringType, nullable = false),
+        StructField("testSecondaryKey", StringType, nullable = false),
+        StructField("testPrice", org.apache.spark.sql.types.DoubleType, nullable = false)
       ))
 
       val batchOperations = Seq(
@@ -640,8 +640,8 @@ class TransactionalBatchITest extends IntegrationSpec
       new com.azure.cosmos.models.PartitionKeyDefinition()
     )
     val paths = new java.util.ArrayList[String]()
-    paths.add("/PermId")
-    paths.add("/SourceId")
+    paths.add("/testPrimaryKey")
+    paths.add("/testSecondaryKey")
     containerProperties.getPartitionKeyDefinition.setPaths(paths)
     containerProperties.getPartitionKeyDefinition.setKind(com.azure.cosmos.models.PartitionKind.MULTI_HASH)
     containerProperties.getPartitionKeyDefinition.setVersion(com.azure.cosmos.models.PartitionKeyDefinitionVersion.V2)
@@ -654,9 +654,9 @@ class TransactionalBatchITest extends IntegrationSpec
     // Create 101 operations for the same hierarchical partition key
     val schema = StructType(Seq(
       StructField("id", StringType, nullable = false),
-      StructField("PermId", StringType, nullable = false),
-      StructField("SourceId", StringType, nullable = false),
-      StructField("price", org.apache.spark.sql.types.DoubleType, nullable = false)
+      StructField("testPrimaryKey", StringType, nullable = false),
+      StructField("testSecondaryKey", StringType, nullable = false),
+      StructField("testPrice", org.apache.spark.sql.types.DoubleType, nullable = false)
     ))
 
     val batchOperations = (1 to 101).map { i =>