fix for hpk container creation on aad auth

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

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed an issue where creating containers with hierarchical partition keys (multi-hash) through the Spark catalog on the AAD path would fail. - See [PR 48548](https://github.com/Azure/azure-sdk-for-java/pull/48548)
 
 #### Other Changes
 

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

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed an issue where creating containers with hierarchical partition keys (multi-hash) through the Spark catalog on the AAD path would fail. - See [PR 48548](https://github.com/Azure/azure-sdk-for-java/pull/48548)
 
 #### Other Changes
 

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

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed an issue where creating containers with hierarchical partition keys (multi-hash) through the Spark catalog on the AAD path would fail. - See [PR 48548](https://github.com/Azure/azure-sdk-for-java/pull/48548)
 
 #### Other Changes
 

sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed an issue where creating containers with hierarchical partition keys (multi-hash) through the Spark catalog on the AAD path would fail. - See [PR 48548](https://github.com/Azure/azure-sdk-for-java/pull/48548)
 
 #### Other Changes
 

sdk/cosmos/azure-cosmos-spark_3/pom.xml

@@ -172,7 +172,27 @@
     <dependency>
       <groupId>org.mockito</groupId>
       <artifactId>mockito-core</artifactId>
-      <version>4.8.1</version> <!-- {x-version-update;cosmos_org.mockito:mockito-core;external_dependency} -->
+      <version>4.11.0</version> <!-- {x-version-update;org.mockito:mockito-core;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-inline</artifactId>
+      <version>4.11.0</version> <!-- {x-version-update;org.mockito:mockito-inline;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+    <!-- bytebuddy dependencies are required for mockito 4.11.0 to work with Java 21. Mockito 4.11.0 is the last release -->
+    <!-- of Mockito supporting Java 8 as a baseline. -->
+    <dependency>
+      <groupId>net.bytebuddy</groupId>
+      <artifactId>byte-buddy</artifactId>
+      <version>1.17.7</version> <!-- {x-version-update;testdep_net.bytebuddy:byte-buddy;external_dependency} -->
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>net.bytebuddy</groupId>
+      <artifactId>byte-buddy-agent</artifactId>
+      <version>1.17.7</version> <!-- {x-version-update;testdep_net.bytebuddy:byte-buddy-agent;external_dependency} -->
       <scope>test</scope>
     </dependency>
     <dependency>

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

@@ -8,7 +8,7 @@ import com.azure.cosmos.models.{FeedRange, PartitionKeyDefinitionVersion, SparkM
 import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait
 import com.azure.cosmos.spark.{ContainerFeedRangesCache, CosmosConstants}
 import com.azure.resourcemanager.cosmos.CosmosManager
-import com.azure.resourcemanager.cosmos.models.{AutoscaleSettings, AutoscaleSettingsResource, ContainerPartitionKey, CreateUpdateOptions, ExcludedPath, IncludedPath, IndexingMode, IndexingPolicy, SqlContainerCreateUpdateParameters, SqlContainerGetPropertiesResource, SqlContainerResource, SqlDatabaseCreateUpdateParameters, SqlDatabaseResource, ThroughputSettingsGetPropertiesResource, ThroughputSettingsResource, ThroughputSettingsUpdateParameters, VectorEmbeddingPolicy => MgmtVectorEmbeddingPolicy, VectorIndex, VectorIndexType}
+import com.azure.resourcemanager.cosmos.models.{AutoscaleSettings, AutoscaleSettingsResource, ContainerPartitionKey, CreateUpdateOptions, ExcludedPath, IncludedPath, IndexingMode, IndexingPolicy, PartitionKind, SqlContainerCreateUpdateParameters, SqlContainerGetPropertiesResource, SqlContainerResource, SqlDatabaseCreateUpdateParameters, SqlDatabaseResource, ThroughputSettingsGetPropertiesResource, ThroughputSettingsResource, ThroughputSettingsUpdateParameters, VectorEmbeddingPolicy => MgmtVectorEmbeddingPolicy, VectorIndex, VectorIndexType}
 import com.fasterxml.jackson.annotation.JsonInclude.Include
 import com.fasterxml.jackson.databind.ObjectMapper
 import org.apache.spark.sql.connector.catalog.{NamespaceChange, TableChange}
@@ -229,13 +229,42 @@ private[spark] case class CosmosCatalogManagementSDKClient(resourceGroupName: St
         val partitionKeyPath = CosmosContainerProperties.getPartitionKeyPath(containerProperties)
         val pkVersion = CosmosContainerProperties.getPartitionKeyVersion(containerProperties)
 
-        containerPartitionKey.withPaths(util.Arrays.asList(partitionKeyPath))
-        if (pkVersion.isDefined) {
-            val partitionKeyDefinitionVersion = PartitionKeyDefinitionVersion.valueOf(pkVersion.get)
-            partitionKeyDefinitionVersion match {
-                case PartitionKeyDefinitionVersion.V1 => containerPartitionKey.withVersion(1)
-                case PartitionKeyDefinitionVersion.V2 => containerPartitionKey.withVersion(2)
-                case version: PartitionKeyDefinitionVersion => throw new IllegalStateException(s"Partition key version $version is not supported.")
+        val pathList = partitionKeyPath.split(",").toList
+        if (pathList.size >= 2) {
+            containerPartitionKey.withKind(CosmosContainerProperties.getPartitionKeyKind(containerProperties) match {
+                case Some(pkKind) =>
+                    if (pkKind == PartitionKind.HASH.toString) {
+                        throw new IllegalArgumentException("PartitionKind HASH is not supported for multi-hash partition key")
+                    }
+                    PartitionKind.MULTI_HASH
+                case None => PartitionKind.MULTI_HASH
+            })
+            containerPartitionKey.withVersion(pkVersion match {
+                case Some(version) =>
+                    if (version == PartitionKeyDefinitionVersion.V1.toString) {
+                        throw new IllegalArgumentException("PartitionKeyVersion V1 is not supported for multi-hash partition key")
+                    }
+                    2
+                case None => 2
+            })
+            containerPartitionKey.withPaths(pathList.map(_.trim).asJava)
+        } else {
+            containerPartitionKey.withKind(CosmosContainerProperties.getPartitionKeyKind(containerProperties) match {
+                case Some(pkKind) =>
+                    if (pkKind == PartitionKind.MULTI_HASH.toString) {
+                        throw new IllegalArgumentException("PartitionKind MULTI_HASH is not supported for single-hash partition key")
+                    }
+                    PartitionKind.HASH
+                case None => PartitionKind.HASH
+            })
+            containerPartitionKey.withPaths(util.Arrays.asList(partitionKeyPath))
+            if (pkVersion.isDefined) {
+                val partitionKeyDefinitionVersion = PartitionKeyDefinitionVersion.valueOf(pkVersion.get)
+                partitionKeyDefinitionVersion match {
+                    case PartitionKeyDefinitionVersion.V1 => containerPartitionKey.withVersion(1)
+                    case PartitionKeyDefinitionVersion.V2 => containerPartitionKey.withVersion(2)
+                    case version: PartitionKeyDefinitionVersion => throw new IllegalStateException(s"Partition key version $version is not supported.")
+                }
             }
         }
 

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

@@ -246,7 +246,7 @@ class CosmosPartitionPlannerITest
     val initialFeedRangesCount = pagesRetrievedCounterMap.size()
 
     // mocking container to count number of pages retrieved
-    val mockContainer = Mockito.spy(container)
+    val mockContainer = Mockito.spy[CosmosAsyncContainer](container)
     Mockito.doAnswer(new Answer[CosmosPagedFlux[ObjectNode]]() {
       override def answer(invocationOnMock: InvocationOnMock): CosmosPagedFlux[ObjectNode] = {
         val requestOptions: CosmosChangeFeedRequestOptions =
@@ -285,7 +285,7 @@ class CosmosPartitionPlannerITest
 
     val container = this.cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainer)
 
-    val mockContainer = Mockito.spy(container)
+    val mockContainer = Mockito.spy[CosmosAsyncContainer](container)
     val invocationCount = new AtomicInteger(0)
     Mockito.doAnswer(new Answer[Mono[util.List[FeedRange]]]() {
       override def answer(invocationOnMock: InvocationOnMock): Mono[util.List[FeedRange]] = {

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

@@ -0,0 +1,149 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.spark.catalog
+
+import com.azure.cosmos.spark.UnitSpec
+import com.azure.resourcemanager.cosmos.CosmosManager
+import com.azure.resourcemanager.cosmos.fluent.CosmosDBManagementClient
+import com.azure.resourcemanager.cosmos.fluent.SqlResourcesClient
+import com.azure.resourcemanager.cosmos.models.{ContainerPartitionKey, PartitionKind}
+import org.mockito.{MockMakers, Mockito}
+import org.mockito.Mockito.{mock, when}
+
+import java.lang.reflect.InvocationTargetException
+import scala.collection.JavaConverters._
+
+class CosmosCatalogManagementSDKClientSpec extends UnitSpec {
+  //scalastyle:off multiple.string.literals
+
+  private val mockCosmosManager = mock(classOf[CosmosManager], Mockito.withSettings().mockMaker(MockMakers.INLINE))
+  private val mockServiceClient = mock(classOf[CosmosDBManagementClient], Mockito.withSettings().mockMaker(MockMakers.INLINE))
+  private val mockSqlResources = mock(classOf[SqlResourcesClient])
+  when(mockCosmosManager.serviceClient()).thenReturn(mockServiceClient)
+  when(mockServiceClient.getSqlResources()).thenReturn(mockSqlResources)
+
+  private val client = CosmosCatalogManagementSDKClient(
+    resourceGroupName = "rg",
+    databaseAccountName = "account",
+    cosmosManager = mockCosmosManager,
+    cosmosAsyncClient = null
+  )
+
+  private def invokeGetPartitionKeyDefinition(containerProperties: Map[String, String]): ContainerPartitionKey = {
+    val method = client.getClass.getDeclaredMethod("getPartitionKeyDefinition", classOf[Map[_, _]])
+    method.setAccessible(true)
+    try {
+      method.invoke(client, containerProperties).asInstanceOf[ContainerPartitionKey]
+    } catch {
+      case e: InvocationTargetException => throw e.getCause
+    }
+  }
+
+  // --- Multi-path partition key tests ---
+
+  "getPartitionKeyDefinition" should "default to MULTI_HASH with version 2 for multi-path without explicit kind or version" in {
+    val props = Map("partitionKeyPath" -> "/tenantId,/userId,/sessionId")
+    val result = invokeGetPartitionKeyDefinition(props)
+
+    result.kind() shouldEqual PartitionKind.MULTI_HASH
+    result.version() shouldEqual 2
+    result.paths().asScala should contain theSameElementsInOrderAs List("/tenantId", "/userId", "/sessionId")
+  }
+
+  it should "succeed for multi-path with explicit MultiHash kind" in {
+    val props = Map(
+      "partitionKeyPath" -> "/a,/b",
+      "partitionKeyKind" -> "MultiHash"
+    )
+    val result = invokeGetPartitionKeyDefinition(props)
+
+    result.kind() shouldEqual PartitionKind.MULTI_HASH
+    result.version() shouldEqual 2
+    result.paths().asScala should contain theSameElementsInOrderAs List("/a", "/b")
+  }
+
+  it should "throw IllegalArgumentException for multi-path with Hash kind" in {
+    val props = Map(
+      "partitionKeyPath" -> "/a,/b",
+      "partitionKeyKind" -> "Hash"
+    )
+
+    an[IllegalArgumentException] should be thrownBy invokeGetPartitionKeyDefinition(props)
+  }
+
+  it should "throw IllegalArgumentException for multi-path with V1 version" in {
+    val props = Map(
+      "partitionKeyPath" -> "/a,/b",
+      "partitionKeyVersion" -> "V1"
+    )
+
+    an[IllegalArgumentException] should be thrownBy invokeGetPartitionKeyDefinition(props)
+  }
+
+  it should "succeed for multi-path with explicit V2 version" in {
+    val props = Map(
+      "partitionKeyPath" -> "/a,/b",
+      "partitionKeyVersion" -> "V2"
+    )
+    val result = invokeGetPartitionKeyDefinition(props)
+
+    result.kind() shouldEqual PartitionKind.MULTI_HASH
+    result.version() shouldEqual 2
+    result.paths().asScala should contain theSameElementsInOrderAs List("/a", "/b")
+  }
+
+  // --- Single-path partition key tests ---
+
+  it should "default to HASH for single-path without explicit kind" in {
+    val props = Map("partitionKeyPath" -> "/id")
+    val result = invokeGetPartitionKeyDefinition(props)
+
+    result.kind() shouldEqual PartitionKind.HASH
+    result.paths().asScala should contain theSameElementsInOrderAs List("/id")
+  }
+
+  it should "throw IllegalArgumentException for single-path with MULTI_HASH kind" in {
+    val props = Map(
+      "partitionKeyPath" -> "/id",
+      "partitionKeyKind" -> "MultiHash"
+    )
+
+    an[IllegalArgumentException] should be thrownBy invokeGetPartitionKeyDefinition(props)
+  }
+
+  it should "set version 1 for single-path with explicit V1 version" in {
+    val props = Map(
+      "partitionKeyPath" -> "/id",
+      "partitionKeyVersion" -> "V1"
+    )
+    val result = invokeGetPartitionKeyDefinition(props)
+
+    result.kind() shouldEqual PartitionKind.HASH
+    result.version() shouldEqual 1
+    result.paths().asScala should contain theSameElementsInOrderAs List("/id")
+  }
+
+  it should "set version 2 for single-path with explicit V2 version" in {
+    val props = Map(
+      "partitionKeyPath" -> "/id",
+      "partitionKeyVersion" -> "V2"
+    )
+    val result = invokeGetPartitionKeyDefinition(props)
+
+    result.kind() shouldEqual PartitionKind.HASH
+    result.version() shouldEqual 2
+    result.paths().asScala should contain theSameElementsInOrderAs List("/id")
+  }
+
+  // --- Default path test ---
+
+  it should "use default partition key path /id when partitionKeyPath is not specified" in {
+    val props = Map.empty[String, String]
+    val result = invokeGetPartitionKeyDefinition(props)
+
+    result.kind() shouldEqual PartitionKind.HASH
+    result.paths().asScala should contain theSameElementsInOrderAs List("/id")
+  }
+
+  //scalastyle:on multiple.string.literals
+}

sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed an issue where creating containers with hierarchical partition keys (multi-hash) through the Spark catalog on the AAD path would fail. - See [PR 48548](https://github.com/Azure/azure-sdk-for-java/pull/48548)
 
 #### Other Changes