diff --git a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md index 31ae31b1b013..74ac4c379930 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md +++ b/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 diff --git a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md index f4e90bc17b63..02250a25b340 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md +++ b/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 diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md index 39254b96c705..648fcd56c889 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md +++ b/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 diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md index 5d677995de71..cf431fe12e03 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md +++ b/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 diff --git a/sdk/cosmos/azure-cosmos-spark_3/pom.xml b/sdk/cosmos/azure-cosmos-spark_3/pom.xml index c78e81556e8c..7323fc20eb8b 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/pom.xml +++ b/sdk/cosmos/azure-cosmos-spark_3/pom.xml @@ -172,7 +172,27 @@ org.mockito mockito-core - 4.8.1 + 4.11.0 + test + + + org.mockito + mockito-inline + 4.11.0 + test + + + + + net.bytebuddy + byte-buddy + 1.17.7 + test + + + net.bytebuddy + byte-buddy-agent + 1.17.7 test diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/catalog/CosmosCatalogManagementSDKClient.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/catalog/CosmosCatalogManagementSDKClient.scala index 8ed1b3cb45c3..b6ee7330a595 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/catalog/CosmosCatalogManagementSDKClient.scala +++ b/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.") + } } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionPlannerITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionPlannerITest.scala index 434ca8e74b37..392d36456138 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionPlannerITest.scala +++ b/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]] = { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/catalog/CosmosCatalogManagementSDKClientSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/catalog/CosmosCatalogManagementSDKClientSpec.scala new file mode 100644 index 000000000000..67ce549fb7a3 --- /dev/null +++ b/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 +} diff --git a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md index 9a438a3f1c53..3540c1054dda 100644 --- a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md +++ b/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