fix for cosmos spark preferred regions issue (Azure#27084)

* added a comment on query plan caching. * fix for cosmos spark preferred regions issue * fix for cosmos spark preferred regions issue
xinlian12 · Mar 3, 2022 · 72a7a12 · 72a7a12
1 parent fbdaf90
commit 72a7a12
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 10 deletions.
diff --git a/...cosmos/azure-cosmos-spark_3_2-12/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/...cosmos/azure-cosmos-spark_3_2-12/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala
@@ -262,7 +262,7 @@ private object CosmosAccountConfig {
     helpMessage = "Cosmos DB Account Name")
 
 
-  private val PreferredRegionRegex = "^[a-z0-9]+$"r // this is for the final form after lower-casing and trimming the whitespaces
+  private val PreferredRegionRegex = "^[a-z0-9\\d]+(?: [a-z0-9\\d]+)*$".r
   private val PreferredRegionsList = CosmosConfigEntry[Array[String]](key = CosmosConfigNames.PreferredRegionsList,
     Option.apply(CosmosConfigNames.PreferredRegions),
     mandatory = false,
@@ -277,7 +277,7 @@ private object CosmosAccountConfig {
       } else {
         trimmedInput.split(",")
           .toStream
-          .map(preferredRegion => preferredRegion.toLowerCase(Locale.ROOT).replace(" ", ""))
+          .map(preferredRegion => preferredRegion.toLowerCase(Locale.ROOT).trim)
           .map(preferredRegion => {
             if (!PreferredRegionRegex.findFirstIn(preferredRegion).isDefined) {
               throw new IllegalArgumentException(s"$preferredRegionsListAsString is invalid")

diff --git a/...os/azure-cosmos-spark_3_2-12/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala b/...os/azure-cosmos-spark_3_2-12/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala
@@ -27,7 +27,7 @@ class CosmosConfigSpec extends UnitSpec {
     endpointConfig.applicationName.get shouldEqual "myapp"
     endpointConfig.useGatewayMode shouldEqual true
     endpointConfig.preferredRegionsList.isDefined shouldEqual true
-    endpointConfig.preferredRegionsList.get should contain theSameElementsAs Array("westus", "eastus1")
+    endpointConfig.preferredRegionsList.get should contain theSameElementsAs Array("west us", "eastus1")
   }
 
   "Config Parser" should "parse account credentials with spark.cosmos.preferredRegions" in {
@@ -46,7 +46,7 @@ class CosmosConfigSpec extends UnitSpec {
     endpointConfig.applicationName.get shouldEqual "myapp"
     endpointConfig.useGatewayMode shouldEqual true
     endpointConfig.preferredRegionsList.isDefined shouldEqual true
-    endpointConfig.preferredRegionsList.get should contain theSameElementsAs Array("westus", "eastus1")
+    endpointConfig.preferredRegionsList.get should contain theSameElementsAs Array("west us", "eastus1")
   }
 
   "Config Parser" should "parse account credentials with spark.cosmos.preferredRegions and spark.cosmos.preferredRegionsList" in {
@@ -116,6 +116,22 @@ class CosmosConfigSpec extends UnitSpec {
           "invalid configuration for spark.cosmos.preferredRegionsList:[westus, eastus. Config description: Preferred Region List"
       }
     }
+
+      userConfig = Map(
+          "spark.cosmos.accountEndpoint" -> sampleProdEndpoint,
+          "spark.cosmos.accountKey" -> "xyz",
+          "spark.cosmos.preferredRegionsList" -> "[west  us, eastus]"
+      )
+
+      try {
+          CosmosAccountConfig.parseCosmosAccountConfig(userConfig)
+          fail("invalid preferred region list")
+      } catch {
+          case e: Exception => {
+              e.getMessage shouldEqual
+                  "invalid configuration for spark.cosmos.preferredRegionsList:[west  us, eastus]. Config description: Preferred Region List"
+          }
+      }
   }
 
   it should "preferred regions parsing" in {
@@ -126,7 +142,7 @@ class CosmosConfigSpec extends UnitSpec {
     )
 
     var config = CosmosAccountConfig.parseCosmosAccountConfig(userConfig)
-    config.preferredRegionsList.get should contain theSameElementsAs Array("eastus", "westus1")
+    config.preferredRegionsList.get should contain theSameElementsAs Array("eastus", "west us1")
 
     userConfig = Map(
       "spark.cosmos.accountEndpoint" -> sampleProdEndpoint,
@@ -135,7 +151,7 @@ class CosmosConfigSpec extends UnitSpec {
     )
 
     config = CosmosAccountConfig.parseCosmosAccountConfig(userConfig)
-    config.preferredRegionsList.get should contain theSameElementsAs Array("eastus", "westus1")
+    config.preferredRegionsList.get should contain theSameElementsAs Array("eastus", "west us1")
 
     userConfig = Map(
       "spark.cosmos.accountEndpoint" -> sampleProdEndpoint,
@@ -144,7 +160,7 @@ class CosmosConfigSpec extends UnitSpec {
     )
 
     config = CosmosAccountConfig.parseCosmosAccountConfig(userConfig)
-    config.preferredRegionsList.get should contain theSameElementsAs Array("eastus", "westus1")
+    config.preferredRegionsList.get should contain theSameElementsAs Array("eastus", "west us1")
 
     userConfig = Map(
       "spark.cosmos.accountEndpoint" -> sampleProdEndpoint,
@@ -153,7 +169,7 @@ class CosmosConfigSpec extends UnitSpec {
     )
 
     config = CosmosAccountConfig.parseCosmosAccountConfig(userConfig)
-    config.preferredRegionsList.get should contain theSameElementsAs Array("eastus", "westus1")
+    config.preferredRegionsList.get should contain theSameElementsAs Array("eastus", "west us1")
 
     userConfig = Map(
       "spark.cosmos.accountEndpoint" -> sampleProdEndpoint,
@@ -162,7 +178,7 @@ class CosmosConfigSpec extends UnitSpec {
     )
 
     config = CosmosAccountConfig.parseCosmosAccountConfig(userConfig)
-    config.preferredRegionsList.get should contain theSameElementsAs Array("westus1")
+    config.preferredRegionsList.get should contain theSameElementsAs Array("west us1")
 
     userConfig = Map(
       "spark.cosmos.accountEndpoint" -> sampleProdEndpoint,
@@ -171,7 +187,7 @@ class CosmosConfigSpec extends UnitSpec {
     )
 
     config = CosmosAccountConfig.parseCosmosAccountConfig(userConfig)
-    config.preferredRegionsList.get should contain theSameElementsAs Array("westus1")
+    config.preferredRegionsList.get should contain theSameElementsAs Array("west us1")
 
     userConfig = Map(
       "spark.cosmos.accountEndpoint" -> sampleProdEndpoint,
@@ -182,6 +198,15 @@ class CosmosConfigSpec extends UnitSpec {
     config = CosmosAccountConfig.parseCosmosAccountConfig(userConfig)
     config.preferredRegionsList.get should contain theSameElementsAs Array[String]()
 
+    userConfig = Map(
+      "spark.cosmos.accountEndpoint" -> sampleProdEndpoint,
+      "spark.cosmos.accountKey" -> "xyz",
+      "spark.cosmos.preferredRegionsList" -> "[west us 1, east us 2]"
+    )
+
+    config = CosmosAccountConfig.parseCosmosAccountConfig(userConfig)
+    config.preferredRegionsList.get should contain theSameElementsAs Array("west us 1","east us 2")
+
   }