diff --git a/sdk/cosmos/.gitignore b/sdk/cosmos/.gitignore index 1ea74182f6bb..81d278c1728f 100644 --- a/sdk/cosmos/.gitignore +++ b/sdk/cosmos/.gitignore @@ -2,3 +2,5 @@ metastore_db/* spark-warehouse/* + +.temp/ 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 46cb9d880f87..87600e4128c1 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 @@ -3,6 +3,8 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added Spark config `spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch` (default `1`) to bound the per-task prefetch parallelism the SDK uses inside `readManyByPartitionKeys`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking 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 d1f06f57031a..56b546f71d87 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 @@ -3,6 +3,8 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added Spark config `spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch` (default `1`) to bound the per-task prefetch parallelism the SDK uses inside `readManyByPartitionKeys`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala b/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala index 5f9cb1dbdbc8..c5bce27bcda9 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala +++ b/sdk/cosmos/azure-cosmos-spark_3-5/src/test/scala/com/azure/cosmos/spark/SparkE2EQueryITest.scala @@ -4,13 +4,20 @@ package com.azure.cosmos.spark import com.azure.cosmos.implementation.TestConfigurations +import com.azure.cosmos.models.{CosmosContainerProperties, CosmosItemRequestOptions, PartitionKey, PartitionKeyBuilder, PartitionKeyDefinition, PartitionKeyDefinitionVersion, PartitionKind, ThroughputProperties} +import com.azure.cosmos.spark.udf.GetCosmosPartitionKeyValue import com.fasterxml.jackson.databind.node.ObjectNode +import org.apache.spark.sql.functions.expr +import org.apache.spark.sql.types.StringType -import java.util.UUID +import java.util.{ArrayList, UUID} + +import scala.collection.JavaConverters._ class SparkE2EQueryITest extends SparkE2EQueryITestBase { + // scalastyle:off multiple.string.literals "spark query" can "return proper Cosmos specific query plan on explain with nullable properties" in { val cosmosEndpoint = TestConfigurations.HOST val cosmosMasterKey = TestConfigurations.MASTER_KEY @@ -67,4 +74,194 @@ class SparkE2EQueryITest val item = rowsArray(0) item.getAs[String]("id") shouldEqual id } + + "spark readManyByPartitionKeys" can "use a matching top-level partition key column without the UDF" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + val requestOptions = new CosmosItemRequestOptions() + + Seq("pkA", "pkB").foreach { pkValue => + val item = objectMapper.createObjectNode() + item.put("id", s"item-$pkValue") + item.put("pk", pkValue) + item.put("payload", s"value-$pkValue") + + container.createItem(item, new PartitionKey(pkValue), requestOptions).block() + } + + val cfg = Map( + "spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.accountKey" -> cosmosMasterKey, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, + "spark.cosmos.read.inferSchema.enabled" -> "true" + ) + + val sparkSession = spark + import sparkSession.implicits._ + + val rows = CosmosItemsDataSource + .readManyByPartitionKeys(Seq("pkA", "pkB").toDF("pk"), cfg.asJava) + .selectExpr("id", "pk", "payload") + .collect() + + rows should have size 2 + rows.map(_.getAs[String]("id")).toSet shouldEqual Set("item-pkA", "item-pkB") + rows.map(_.getAs[String]("pk")).toSet shouldEqual Set("pkA", "pkB") + rows.map(_.getAs[String]("payload")).toSet shouldEqual Set("value-pkA", "value-pkB") + } + + "spark readManyByPartitionKeys" can "require the UDF for nested partition key paths and succeed with it" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val containerName = s"nested-pk-${UUID.randomUUID()}" + + val pkPaths = new ArrayList[String]() + pkPaths.add("/tenant/id") + + val pkDefinition = new PartitionKeyDefinition() + pkDefinition.setPaths(pkPaths) + pkDefinition.setKind(PartitionKind.HASH) + pkDefinition.setVersion(PartitionKeyDefinitionVersion.V2) + + val containerProperties = new CosmosContainerProperties(containerName, pkDefinition) + cosmosClient + .getDatabase(cosmosDatabase) + .createContainerIfNotExists(containerProperties, ThroughputProperties.createManualThroughput(400)) + .block() + + try { + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(containerName) + val requestOptions = new CosmosItemRequestOptions() + + Seq("tenantA", "tenantB").foreach { tenantId => + val item = objectMapper.createObjectNode() + item.put("id", s"item-$tenantId") + item.put("payload", s"value-$tenantId") + item.putObject("tenant").put("id", tenantId) + + container.createItem(item, new PartitionKey(tenantId), requestOptions).block() + } + + val cfg = Map( + "spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.accountKey" -> cosmosMasterKey, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> containerName, + "spark.cosmos.read.inferSchema.enabled" -> "true" + ) + + val sparkSession = spark + import sparkSession.implicits._ + + val missingUdfError = the[IllegalArgumentException] thrownBy { + CosmosItemsDataSource.readManyByPartitionKeys(Seq("tenantA").toDF("tenantId"), cfg.asJava) + } + + missingUdfError.getMessage should include("Nested paths cannot be resolved from DataFrame columns automatically") + missingUdfError.getMessage should include("_partitionKeyIdentity") + + spark.udf.register("GetCosmosPartitionKeyValue", new GetCosmosPartitionKeyValue(), StringType) + + val inputDf = Seq("tenantA", "tenantB") + .toDF("tenantId") + .withColumn("_partitionKeyIdentity", expr("GetCosmosPartitionKeyValue(tenantId)")) + + val rows = CosmosItemsDataSource + .readManyByPartitionKeys(inputDf, cfg.asJava) + .selectExpr("id", "tenant.id as tenantId") + .collect() + + rows should have size 2 + rows.map(_.getAs[String]("id")).toSet shouldEqual Set("item-tenantA", "item-tenantB") + rows.map(_.getAs[String]("tenantId")).toSet shouldEqual Set("tenantA", "tenantB") + } finally { + cosmosClient + .getDatabase(cosmosDatabase) + .getContainer(containerName) + .delete() + .block() + } + } + + + "spark readManyByPartitionKeys" can "support partial top-level hierarchical partition keys from DataFrame columns without the UDF" in { + val cosmosEndpoint = TestConfigurations.HOST + val cosmosMasterKey = TestConfigurations.MASTER_KEY + val containerName = s"top-level-hpk-${UUID.randomUUID()}" + + val pkPaths = new ArrayList[String]() + pkPaths.add("/tenant") + pkPaths.add("/region") + pkPaths.add("/team") + + val pkDefinition = new PartitionKeyDefinition() + pkDefinition.setPaths(pkPaths) + pkDefinition.setKind(PartitionKind.MULTI_HASH) + pkDefinition.setVersion(PartitionKeyDefinitionVersion.V2) + + val containerProperties = new CosmosContainerProperties(containerName, pkDefinition) + cosmosClient + .getDatabase(cosmosDatabase) + .createContainerIfNotExists(containerProperties, ThroughputProperties.createManualThroughput(400)) + .block() + + try { + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(containerName) + val requestOptions = new CosmosItemRequestOptions() + + Seq( + ("tenantA", "east", "sales", "item-a1"), + ("tenantA", "west", "hr", "item-a2"), + ("tenantB", "east", "sales", "item-b1") + ).foreach { case (tenant, region, team, id) => + val item = objectMapper.createObjectNode() + item.put("id", id) + item.put("tenant", tenant) + item.put("region", region) + item.put("team", team) + item.put("payload", s"$tenant-$region-$team") + + val pk = new PartitionKeyBuilder().add(tenant).add(region).add(team).build() + container.createItem(item, pk, requestOptions).block() + } + + val cfg = Map( + "spark.cosmos.accountEndpoint" -> cosmosEndpoint, + "spark.cosmos.accountKey" -> cosmosMasterKey, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> containerName, + "spark.cosmos.read.inferSchema.enabled" -> "true" + ) + + val sparkSession = spark + import sparkSession.implicits._ + + val tenantRows = CosmosItemsDataSource + .readManyByPartitionKeys(Seq("tenantA").toDF("tenant"), cfg.asJava) + .selectExpr("id", "tenant", "region", "team") + .collect() + + tenantRows should have size 2 + tenantRows.map(_.getAs[String]("id")).toSet shouldEqual Set("item-a1", "item-a2") + tenantRows.map(_.getAs[String]("tenant")).toSet shouldEqual Set("tenantA") + + val tenantRegionRows = CosmosItemsDataSource + .readManyByPartitionKeys(Seq(("tenantA", "east")).toDF("tenant", "region"), cfg.asJava) + .selectExpr("id", "tenant", "region", "team") + .collect() + + tenantRegionRows should have size 1 + tenantRegionRows.head.getAs[String]("id") shouldEqual "item-a1" + } finally { + cosmosClient + .getDatabase(cosmosDatabase) + .getContainer(containerName) + .delete() + .block() + } + } + + // scalastyle:on multiple.string.literals } 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 93bb3fa96a75..acc6d23dbe2b 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 @@ -3,6 +3,8 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added Spark config `spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch` (default `1`) to bound the per-task prefetch parallelism the SDK uses inside `readManyByPartitionKeys`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking 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 e073483e335c..ea6312673520 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 @@ -3,6 +3,8 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added Spark config `spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch` (default `1`) to bound the per-task prefetch parallelism the SDK uses inside `readManyByPartitionKeys`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/CosmosItemSerializerNoExceptionWrapping.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/CosmosItemSerializerNoExceptionWrapping.scala index a354cee66efc..9dabaf82c8f6 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/CosmosItemSerializerNoExceptionWrapping.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/CosmosItemSerializerNoExceptionWrapping.scala @@ -14,4 +14,12 @@ private[cosmos] abstract class CosmosItemSerializerNoExceptionWrapping extends C .CosmosItemSerializerHelper .getCosmosItemSerializerAccessor .setCanSerialize(this, false) + + // canSerialize is set to false above, so the SDK will never call serialize(). + // This default implementation satisfies the abstract method contract and throws if + // called unexpectedly. + override def serialize[T](item: T): java.util.Map[String, AnyRef] = { + throw new UnsupportedOperationException( + "serialize() is not supported on CosmosItemSerializerNoExceptionWrapping (canSerialize = false)") + } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala index 951f4735444d..990d18d0b487 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConfig.scala @@ -92,6 +92,8 @@ private[spark] object CosmosConfigNames { val ReadPartitioningFeedRangeFilter = "spark.cosmos.partitioning.feedRangeFilter" val ReadRuntimeFilteringEnabled = "spark.cosmos.read.runtimeFiltering.enabled" val ReadManyFilteringEnabled = "spark.cosmos.read.readManyFiltering.enabled" + val ReadManyByPkNullHandling = "spark.cosmos.read.readManyByPk.nullHandling" + val ReadManyByPkMaxConcurrentBatchPrefetch = "spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch" val ViewsRepositoryPath = "spark.cosmos.views.repositoryPath" val DiagnosticsMode = "spark.cosmos.diagnostics" val DiagnosticsSamplingMaxCount = "spark.cosmos.diagnostics.sampling.maxCount" @@ -226,6 +228,8 @@ private[spark] object CosmosConfigNames { ReadPartitioningFeedRangeFilter, ReadRuntimeFilteringEnabled, ReadManyFilteringEnabled, + ReadManyByPkNullHandling, + ReadManyByPkMaxConcurrentBatchPrefetch, ViewsRepositoryPath, DiagnosticsMode, DiagnosticsSamplingIntervalInSeconds, @@ -1042,7 +1046,9 @@ private case class CosmosReadConfig(readConsistencyStrategy: ReadConsistencyStra throughputControlConfig: Option[CosmosThroughputControlConfig] = None, runtimeFilteringEnabled: Boolean, readManyFilteringConfig: CosmosReadManyFilteringConfig, - responseContinuationTokenLimitInKb: Option[Int] = None) + responseContinuationTokenLimitInKb: Option[Int] = None, + readManyByPkTreatNullAsNone: Boolean = false, + readManyByPkMaxConcurrentBatchPrefetch: Int = 1) private object SchemaConversionModes extends Enumeration { type SchemaConversionMode = Value @@ -1136,6 +1142,31 @@ private object CosmosReadConfig { helpMessage = " Indicates whether dynamic partition pruning filters will be pushed down when applicable." ) + private val ReadManyByPkNullHandling = CosmosConfigEntry[String]( + key = CosmosConfigNames.ReadManyByPkNullHandling, + mandatory = false, + defaultValue = Some("Null"), + parseFromStringFunction = value => value, + helpMessage = "Determines how null values in partition key columns are treated for " + + "readManyByPartitionKeys. 'Null' (default) maps null to a JSON null via addNullValue(), which " + + "is appropriate when the document field exists with an explicit null value. 'None' maps null " + + "to PartitionKey.NONE via addNoneValue(), which is only supported for single-path partition keys " + + "and should only be used when the partition key path does not exist at all in the document. " + + "Hierarchical partition keys reject this mode. These two semantics hash to DIFFERENT physical " + + "partitions - picking the wrong mode for your data will silently return zero rows." + ) + + private val ReadManyByPkMaxConcurrentBatchPrefetch = CosmosConfigEntry[Int]( + key = CosmosConfigNames.ReadManyByPkMaxConcurrentBatchPrefetch, + mandatory = false, + defaultValue = Some(1), + parseFromStringFunction = value => Math.min(64, Math.max(1, value.toInt)), + helpMessage = "The maximum number of per-physical-partition batches whose first page is prefetched " + + "concurrently inside a single Spark task by the SDK's readManyByPartitionKeys execution. The " + + "default is `1` - max is `64`, because Spark already parallelises across tasks - increase this when individual " + + "tasks span many physical partitions and additional intra-task prefetch is desired." + ) + def parseCosmosReadConfig(cfg: Map[String, String]): CosmosReadConfig = { val forceEventualConsistency = CosmosConfigEntry.parse(cfg, ForceEventualConsistency) val readConsistencyStrategyOverride = CosmosConfigEntry.parse(cfg, ReadConsistencyStrategyOverride) @@ -1158,6 +1189,9 @@ private object CosmosReadConfig { val throughputControlConfigOpt = CosmosThroughputControlConfig.parseThroughputControlConfig(cfg) val runtimeFilteringEnabled = CosmosConfigEntry.parse(cfg, ReadRuntimeFilteringEnabled) val readManyFilteringConfig = CosmosReadManyFilteringConfig.parseCosmosReadManyFilterConfig(cfg) + val readManyByPkNullHandling = CosmosConfigEntry.parse(cfg, ReadManyByPkNullHandling) + val readManyByPkTreatNullAsNone = readManyByPkNullHandling.getOrElse("Null").equalsIgnoreCase("None") + val readManyByPkMaxConcurrentBatchPrefetch = CosmosConfigEntry.parse(cfg, ReadManyByPkMaxConcurrentBatchPrefetch).getOrElse(1) val effectiveReadConsistencyStrategy = if (readConsistencyStrategyOverride.getOrElse(ReadConsistencyStrategy.DEFAULT) != ReadConsistencyStrategy.DEFAULT) { readConsistencyStrategyOverride.get @@ -1189,7 +1223,9 @@ private object CosmosReadConfig { throughputControlConfigOpt, runtimeFilteringEnabled.get, readManyFilteringConfig, - responseContinuationTokenLimitInKb) + responseContinuationTokenLimitInKb, + readManyByPkTreatNullAsNone, + readManyByPkMaxConcurrentBatchPrefetch) } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala index 9ece47416526..00761f23d399 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala @@ -45,6 +45,7 @@ private[cosmos] object CosmosConstants { val Id = "id" val ETag = "_etag" val ItemIdentity = "_itemIdentity" + val PartitionKeyIdentity = "_partitionKeyIdentity" } object StatusCodes { diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala index a35cff27af68..7723b566db12 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala @@ -2,7 +2,7 @@ // Licensed under the MIT License. package com.azure.cosmos.spark -import com.azure.cosmos.models.{CosmosItemIdentity, PartitionKey} +import com.azure.cosmos.models.{CosmosItemIdentity, PartitionKey, PartitionKeyBuilder} import com.azure.cosmos.spark.CosmosPredicates.assertOnSparkDriver import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait import org.apache.spark.sql.{DataFrame, Row, SparkSession} @@ -112,4 +112,124 @@ object CosmosItemsDataSource { readManyReader.readMany(df.rdd, readManyFilterExtraction) } + + def readManyByPartitionKeys(df: DataFrame, userConfig: java.util.Map[String, String]): DataFrame = { + readManyByPartitionKeys(df, userConfig, null) + } + + def readManyByPartitionKeys( + df: DataFrame, + userConfig: java.util.Map[String, String], + userProvidedSchema: StructType): DataFrame = { + + val readManyReader = new CosmosReadManyByPartitionKeyReader( + userProvidedSchema, + userConfig.asScala.toMap) + + // Initialize reader state once: resolves PK paths, infers schema, broadcasts client caches, + // and returns the resolved treatNullAsNone flag from the reader's parsed config - avoiding + // duplicate config parsing between the data source and the reader. + val readerState = readManyReader.initializeReaderState() + val (pkPaths, _, _, sharedTreatNullAsNone) = readerState + + // Option 1: Look for the _partitionKeyIdentity column (produced by GetCosmosPartitionKeyValue UDF) + val pkIdentityFieldExtraction = df + .schema + .find(field => field.name.equals(CosmosConstants.Properties.PartitionKeyIdentity) && field.dataType.equals(StringType)) + .map(field => (row: Row) => { + val rawValue = row.getString(row.fieldIndex(field.name)) + CosmosPartitionKeyHelper.tryParsePartitionKey(rawValue, sharedTreatNullAsNone) + .getOrElse(throw new IllegalArgumentException( + s"Invalid _partitionKeyIdentity value in row: '$rawValue'. " + + "Expected format: pk([...json...])")) + }) + + // Option 2: Detect PK columns by matching the container's partition key paths against the DataFrame schema + val pkColumnExtraction: Option[Row => PartitionKey] = if (pkIdentityFieldExtraction.isDefined) { + None // no need to resolve PK paths - _partitionKeyIdentity column takes precedence + } else { + val treatNullAsNone = sharedTreatNullAsNone + + // Nested PK paths (containing /) cannot be resolved from top-level DataFrame columns. + if (pkPaths.exists(_.contains("/"))) { + throw new IllegalArgumentException( + "Container has nested partition key path(s) " + pkPaths.mkString("[", ",", "]") + ". " + + "Nested paths cannot be resolved from DataFrame columns automatically - add a " + + "'_partitionKeyIdentity' column produced by the GetCosmosPartitionKeyValue UDF.") + } + + // Allow DataFrames to provide a contiguous top-level prefix of the container's + // hierarchical partition key paths. For example: tenant, or tenant + region. + val dfFieldNames = df.schema.fieldNames.toSet + val matchedPrefix = pkPaths.takeWhile(path => dfFieldNames.contains(path)) + val hasNonPrefixMatch = pkPaths.drop(matchedPrefix.size).exists(path => dfFieldNames.contains(path)) + + if (hasNonPrefixMatch) { + throw new IllegalArgumentException( + "DataFrame columns matching the container's partition key paths must form a contiguous top-level prefix " + + "(for example: tenant, or tenant + region). " + + "For nested or non-prefix partition key extraction, add a '_partitionKeyIdentity' column produced " + + "by the GetCosmosPartitionKeyValue UDF.") + } + + if (matchedPrefix.nonEmpty) { + Some((row: Row) => { + if (matchedPrefix.size == 1) { + buildPartitionKey(row.getAs[Any](matchedPrefix.head), treatNullAsNone) + } else { + val builder = new PartitionKeyBuilder() + for (path <- matchedPrefix) { + addPartitionKeyComponent(builder, row.getAs[Any](path), treatNullAsNone, matchedPrefix.size) + } + builder.build() + } + }) + } else { + None + } + } + + val pkExtraction = pkIdentityFieldExtraction + .orElse(pkColumnExtraction) + .getOrElse( + throw new IllegalArgumentException( + "Cannot determine partition key extraction from the input DataFrame. " + + "Either add a '_partitionKeyIdentity' column (using the GetCosmosPartitionKeyValue UDF) " + + "or ensure the DataFrame contains columns matching a top-level prefix of the container's partition key paths.")) + + readManyReader.readManyByPartitionKeys(df.rdd, pkExtraction, readerState) + } + + private def addPartitionKeyComponent( + builder: PartitionKeyBuilder, + value: Any, + treatNullAsNone: Boolean, + partitionKeyComponentCount: Int): Unit = { + value match { + case s: String => builder.add(s) + case n: Number => builder.add(n.doubleValue()) + case b: Boolean => builder.add(b) + case null => + CosmosPartitionKeyHelper.validateNoneHandlingForPartitionKeyComponentCount( + partitionKeyComponentCount, + treatNullAsNone) + if (treatNullAsNone) builder.addNoneValue() + else builder.addNullValue() + case other => + // Reject unknown types rather than silently .toString-ing them - the document field + // was stored with its original type and a stringified value will never match. + // Supported types: String, Number (Byte/Short/Int/Long/Float/Double/BigDecimal), Boolean, null. + throw new IllegalArgumentException( + s"Unsupported partition key column type '${other.getClass.getName}' with value '$other'. " + + "Supported types are String, Number (integral or floating-point), Boolean, and null. " + + "For other source types, convert the column before calling readManyByPartitionKeys or use " + + "the GetCosmosPartitionKeyValue UDF to produce a '_partitionKeyIdentity' column.") + } + } + + private def buildPartitionKey(value: Any, treatNullAsNone: Boolean): PartitionKey = { + val builder = new PartitionKeyBuilder() + addPartitionKeyComponent(builder, value, treatNullAsNone, 1) + builder.build() + } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala new file mode 100644 index 000000000000..234974212983 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala @@ -0,0 +1,93 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.spark + +import com.azure.cosmos.implementation.routing.PartitionKeyInternal +import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, Utils} +import com.azure.cosmos.models.{PartitionKey, PartitionKeyBuilder} +import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait + +import java.util + +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import + +private[spark] object CosmosPartitionKeyHelper extends BasicLoggingTrait { + private[spark] val HierarchicalPartitionKeyNoneHandlingErrorMessage = + s"The configuration '${CosmosConfigNames.ReadManyByPkNullHandling}=None' is not supported for " + + "hierarchical partition keys because PartitionKey.NONE can't be used with multiple paths. " + + "Use 'Null' for explicit JSON null values, filter out rows with missing partition key " + + "components, or provide fully-defined hierarchical partition keys." + + private[spark] def validateNoneHandlingForPartitionKeyComponentCount( + componentCount: Int, + treatNullAsNone: Boolean): Unit = { + if (treatNullAsNone && componentCount > 1) { + throw new IllegalArgumentException(HierarchicalPartitionKeyNoneHandlingErrorMessage) + } + } + + // pattern will be recognized + // pk(partitionKeyValue) + // + // (?i) : The whole matching is case-insensitive + // pk[(](.*)[)]: partitionKey Value + private val cosmosPartitionKeyStringRegx = """(?i)^pk\((.*)\)$""".r + private val objectMapper = Utils.getSimpleObjectMapper + + def getCosmosPartitionKeyValueString(partitionKeyValue: List[Object]): String = { + s"pk(${objectMapper.writeValueAsString(partitionKeyValue.asJava)})" + } + + def tryParsePartitionKey(cosmosPartitionKeyString: String): Option[PartitionKey] = + tryParsePartitionKey(cosmosPartitionKeyString, treatNullAsNone = false) + + /** + * Parses a pk(...) string into a [[PartitionKey]]. + * + * When treatNullAsNone is true, any JSON null components in the serialized array are mapped to + * [[PartitionKeyBuilder.addNoneValue()]] (meaning the document field is absent/undefined). + * When false, they are mapped to [[PartitionKeyBuilder.addNullValue()]] (JSON null value). + * This matches the spark.cosmos.read.readManyByPk.nullHandling config for the non-UDF column path. + */ + def tryParsePartitionKey( + cosmosPartitionKeyString: String, + treatNullAsNone: Boolean): Option[PartitionKey] = { + require(cosmosPartitionKeyString != null, "Argument 'cosmosPartitionKeyString' must not be null.") + cosmosPartitionKeyString match { + case cosmosPartitionKeyStringRegx(pkValue) => + scala.util.Try(Utils.parse(pkValue, classOf[Object])).toOption.flatMap { + case arrayList: util.ArrayList[Object @unchecked] => + val components = arrayList.toArray + if (components.exists(_ == null)) { + validateNoneHandlingForPartitionKeyComponentCount(components.length, treatNullAsNone) + // Build via PartitionKeyBuilder so nulls can be disambiguated between + // JSON-null (addNullValue) and undefined (addNoneValue) based on config. + val builder = new PartitionKeyBuilder() + components.foreach { + case null => + if (treatNullAsNone) builder.addNoneValue() else builder.addNullValue() + case s: String => builder.add(s) + case n: java.lang.Number => builder.add(n.doubleValue()) + case b: java.lang.Boolean => builder.add(b.booleanValue()) + case other => + throw new IllegalArgumentException( + s"Unsupported partition key component type '${other.getClass.getName}' with value '$other'. " + + "Supported types are String, Number (integral or floating-point), Boolean, and null.") + } + Some(builder.build()) + } else { + Some( + ImplementationBridgeHelpers + .PartitionKeyHelper + .getPartitionKeyAccessor + .toPartitionKey(PartitionKeyInternal.fromObjectArray(components, false))) + } + case other => Some(new PartitionKey(other)) + } + case _ => None + } + } +} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala new file mode 100644 index 000000000000..e3ea4298b3f3 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala @@ -0,0 +1,199 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.spark + +import com.azure.cosmos.{CosmosException, ReadConsistencyStrategy, SparkBridgeInternal} +import com.azure.cosmos.implementation.{CosmosClientMetadataCachesSnapshot, UUIDs} +import com.azure.cosmos.models.PartitionKey +import com.azure.cosmos.spark.CosmosPredicates.assertOnSparkDriver +import com.azure.cosmos.spark.diagnostics.{BasicLoggingTrait, DiagnosticsContext} +import com.fasterxml.jackson.databind.node.ObjectNode +import org.apache.spark.TaskContext +import org.apache.spark.broadcast.Broadcast +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.types.StructType + +import java.util.UUID +import java.util.concurrent.atomic.AtomicBoolean + +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import + +private[spark] class CosmosReadManyByPartitionKeyReader( + val userProvidedSchema: StructType, + val userConfig: Map[String, String] + ) extends BasicLoggingTrait with Serializable { + val effectiveUserConfig: Map[String, String] = CosmosConfig.getEffectiveConfig( + databaseName = None, + containerName = None, + userConfig) + + val clientConfig: CosmosAccountConfig = CosmosAccountConfig.parseCosmosAccountConfig(effectiveUserConfig) + val readConfig: CosmosReadConfig = CosmosReadConfig.parseCosmosReadConfig(effectiveUserConfig) + val cosmosContainerConfig: CosmosContainerConfig = + CosmosContainerConfig.parseCosmosContainerConfig(effectiveUserConfig) + //scalastyle:off multiple.string.literals + val tableName: String = s"com.azure.cosmos.spark.items.${clientConfig.accountName}." + + s"${cosmosContainerConfig.database}.${cosmosContainerConfig.container}" + private lazy val sparkSession = { + assertOnSparkDriver() + SparkSession.active + } + val sparkEnvironmentInfo: String = CosmosClientConfiguration.getSparkEnvironmentInfo(Some(sparkSession)) + logTrace(s"Instantiated ${this.getClass.getSimpleName} for $tableName") + + /** + * Resolves the partition key paths for the target container. + * Uses a single Loan block that also infers the schema (if needed) and warms + * the client metadata caches for broadcast — avoiding three separate Loan round-trips. + * + * @return (pkPaths, schema, broadcastClientStates) + */ + private[spark] def initializeReaderState(): (List[String], StructType, Broadcast[CosmosClientMetadataCachesSnapshots], Boolean) = { + val calledFrom = s"CosmosReadManyByPartitionKeyReader($tableName).initializeReaderState" + Loan( + List[Option[CosmosClientCacheItem]]( + Some( + CosmosClientCache( + CosmosClientConfiguration( + effectiveUserConfig, + readConsistencyStrategy = readConfig.readConsistencyStrategy, + sparkEnvironmentInfo), + None, + calledFrom)), + ThroughputControlHelper.getThroughputControlClientCacheItem( + effectiveUserConfig, + calledFrom, + None, + sparkEnvironmentInfo) + )) + .to(clientCacheItems => { + val container = + ThroughputControlHelper.getContainer( + effectiveUserConfig, + cosmosContainerConfig, + clientCacheItems(0).get, + clientCacheItems(1)) + + // 1. Resolve PK paths from the collection cache (with transient-error retry) + val pkPaths = TransientErrorsRetryPolicy.executeWithRetry(() => { + SparkBridgeInternal + .getContainerPropertiesFromCollectionCache(container) + .getPartitionKeyDefinition + .getPaths.asScala.map(_.stripPrefix("/")).toList + }) + + // 2. Infer schema if not user-provided + val schema = Option.apply(userProvidedSchema).getOrElse( + CosmosTableSchemaInferrer.inferSchema( + clientCacheItems(0).get, + clientCacheItems(1), + effectiveUserConfig, + ItemsTable.defaultSchemaForInferenceDisabled)) + + // 3. Warm-up readItem so collection/routing-map caches are populated before broadcast + try { + container.readItem( + UUIDs.nonBlockingRandomUUID().toString, + new PartitionKey(UUIDs.nonBlockingRandomUUID().toString), + classOf[ObjectNode]) + .block() + } catch { + case _: CosmosException => + // Expected when the random read targets a non-existent item; we only need the + // routing map / collection cache populated as a side-effect of the call. + () + } + + // 4. Serialize and broadcast client state + val state = new CosmosClientMetadataCachesSnapshot() + state.serialize(clientCacheItems(0).get.cosmosClient) + + var throughputControlState: Option[CosmosClientMetadataCachesSnapshot] = None + if (clientCacheItems(1).isDefined) { + throughputControlState = Some(new CosmosClientMetadataCachesSnapshot()) + throughputControlState.get.serialize(clientCacheItems(1).get.cosmosClient) + } + + val metadataSnapshots = CosmosClientMetadataCachesSnapshots(state, throughputControlState) + val broadcastStates = sparkSession.sparkContext.broadcast(metadataSnapshots) + + (pkPaths, schema, broadcastStates, readConfig.readManyByPkTreatNullAsNone) + }) + } + + def readManyByPartitionKeys( + inputRdd: RDD[Row], + pkExtraction: Row => PartitionKey, + readerState: (List[String], StructType, Broadcast[CosmosClientMetadataCachesSnapshots], Boolean)): DataFrame = { + + val correlationActivityId = UUIDs.nonBlockingRandomUUID() + val (_, schema, clientStates, _) = readerState + + sparkSession.sqlContext.createDataFrame( + inputRdd.mapPartitionsWithIndex( + (partitionIndex: Int, rowIterator: Iterator[Row]) => { + val pkIterator: Iterator[PartitionKey] = rowIterator + .map(row => pkExtraction.apply(row)) + + logInfo(s"Creating an ItemsPartitionReaderWithReadManyByPartitionKey for Activity $correlationActivityId to read for " + + s"input partition [$partitionIndex] ${tableName}") + + val taskContext = TaskContext.get + val reader = new ItemsPartitionReaderWithReadManyByPartitionKey( + effectiveUserConfig, + CosmosReadManyHelper.FullRangeFeedRange, + schema, + DiagnosticsContext(correlationActivityId, partitionIndex.toString), + clientStates, + DiagnosticsConfig.parseDiagnosticsConfig(effectiveUserConfig), + sparkEnvironmentInfo, + taskContext, + pkIterator) + + new Iterator[Row] { + private val isClosed = new AtomicBoolean(false) + + private def closeReader(): Unit = { + if (isClosed.compareAndSet(false, true)) { + reader.close() + } + } + + if (taskContext != null) { + taskContext.addTaskCompletionListener[Unit](_ => closeReader()) + } + + override def hasNext: Boolean = { + try { + val hasMore = reader.next() + if (!hasMore) { + closeReader() + } + hasMore + } catch { + case error: Throwable => + closeReader() + throw error + } + } + + override def next(): Row = { + try { + reader.getCurrentRow() + } catch { + case error: Throwable => + closeReader() + throw error + } + } + } + }, + preservesPartitioning = true + ), + schema) + } +} + diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala new file mode 100644 index 000000000000..cfcb7397ec03 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala @@ -0,0 +1,318 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.spark + +import com.azure.cosmos.{CosmosAsyncContainer, CosmosEndToEndOperationLatencyPolicyConfigBuilder, CosmosItemSerializer, CosmosItemSerializerNoExceptionWrapping, SparkBridgeInternal} +import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple +import com.azure.cosmos.implementation.{ImplementationBridgeHelpers, ObjectNodeMap, SparkRowItem, Utils} +import com.azure.cosmos.models.{CosmosReadManyByPartitionKeysRequestOptions, ModelBridgeInternal, PartitionKey, PartitionKeyDefinition, SqlQuerySpec} +import com.azure.cosmos.spark.BulkWriter.getThreadInfo +import com.azure.cosmos.spark.CosmosTableSchemaInferrer.IdAttributeName +import com.azure.cosmos.spark.diagnostics.{DetailedFeedDiagnosticsProvider, DiagnosticsContext, DiagnosticsLoader, LoggerHelper, SparkTaskContext} +import com.azure.cosmos.util.CosmosPagedFlux +import com.fasterxml.jackson.databind.node.ObjectNode +import org.apache.spark.TaskContext +import org.apache.spark.broadcast.Broadcast +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.connector.read.PartitionReader +import org.apache.spark.sql.types.StructType + +import java.util +import java.util.concurrent.atomic.AtomicBoolean + +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import + +private[spark] case class ItemsPartitionReaderWithReadManyByPartitionKey +( + config: Map[String, String], + feedRange: NormalizedRange, + readSchema: StructType, + diagnosticsContext: DiagnosticsContext, + cosmosClientStateHandles: Broadcast[CosmosClientMetadataCachesSnapshots], + diagnosticsConfig: DiagnosticsConfig, + sparkEnvironmentInfo: String, + taskContext: TaskContext, + readManyPartitionKeys: Iterator[PartitionKey] +) + extends PartitionReader[InternalRow] { + + private lazy val log = LoggerHelper.getLogger(diagnosticsConfig, this.getClass) + + private val readManyOptions = new CosmosReadManyByPartitionKeysRequestOptions() + private val readManyOptionsImpl = ImplementationBridgeHelpers + .CosmosReadManyByPartitionKeysRequestOptionsHelper + .getCosmosReadManyByPartitionKeysRequestOptionsAccessor + .getImpl(readManyOptions) + + private val readConfig = CosmosReadConfig.parseCosmosReadConfig(config) + ThroughputControlHelper.populateThroughputControlGroupName(readManyOptionsImpl, readConfig.throughputControlConfig) + readManyOptions.setMaxConcurrentBatchPrefetch(readConfig.readManyByPkMaxConcurrentBatchPrefetch) + + private val operationContext = { + assert(taskContext != null) + + SparkTaskContext(diagnosticsContext.correlationActivityId, + taskContext.stageId(), + taskContext.partitionId(), + taskContext.taskAttemptId(), + feedRange.toString) + } + + private val operationContextAndListenerTuple: Option[OperationContextAndListenerTuple] = { + if (diagnosticsConfig.mode.isDefined) { + val listener = + DiagnosticsLoader.getDiagnosticsProvider(diagnosticsConfig).getLogger(this.getClass) + + val ctxAndListener = new OperationContextAndListenerTuple(operationContext, listener) + + readManyOptionsImpl + .setOperationContextAndListenerTuple(ctxAndListener) + + Some(ctxAndListener) + } else { + None + } + } + + log.logTrace(s"Instantiated ${this.getClass.getSimpleName}, Context: ${operationContext.toString} $getThreadInfo") + + private val containerTargetConfig = CosmosContainerConfig.parseCosmosContainerConfig(config) + + log.logInfo(s"Using ReadManyByPartitionKey from feed range $feedRange of " + + s"container ${containerTargetConfig.database}.${containerTargetConfig.container} - " + + s"correlationActivityId ${diagnosticsContext.correlationActivityId}, " + + s"Context: ${operationContext.toString} $getThreadInfo") + + private val clientCacheItem = CosmosClientCache( + CosmosClientConfiguration(config, readConfig.readConsistencyStrategy, sparkEnvironmentInfo), + Some(cosmosClientStateHandles.value.cosmosClientMetadataCaches), + s"ItemsPartitionReaderWithReadManyByPartitionKey($feedRange, ${containerTargetConfig.database}.${containerTargetConfig.container})" + ) + + private val throughputControlClientCacheItemOpt = + ThroughputControlHelper.getThroughputControlClientCacheItem( + config, + clientCacheItem.context, + Some(cosmosClientStateHandles), + sparkEnvironmentInfo) + + private val cosmosAsyncContainer = + ThroughputControlHelper.getContainer( + config, + containerTargetConfig, + clientCacheItem, + throughputControlClientCacheItemOpt) + + private val partitionKeyDefinition: PartitionKeyDefinition = { + TransientErrorsRetryPolicy.executeWithRetry(() => { + SparkBridgeInternal + .getContainerPropertiesFromCollectionCache(cosmosAsyncContainer).getPartitionKeyDefinition + }) + } + + private val cosmosSerializationConfig = CosmosSerializationConfig.parseSerializationConfig(config) + private val cosmosRowConverter = CosmosRowConverter.get(cosmosSerializationConfig) + + readManyOptionsImpl + .setCustomItemSerializer( + new CosmosItemSerializerNoExceptionWrapping { + // The base class (CosmosItemSerializerNoExceptionWrapping) sets canSerialize = false, + // which prevents the SDK from calling serialize() — so no override is needed here. + // Only deserialization (the hot path) is customized below. + + override def deserialize[T](jsonNodeMap: util.Map[String, AnyRef], classType: Class[T]): T = { + if (jsonNodeMap == null) { + throw new IllegalStateException("The 'jsonNodeMap' should never be null here.") + } + + if (classType != classOf[SparkRowItem]) { + throw new IllegalStateException("The 'classType' must be 'classOf[SparkRowItem])' here.") + } + + val objectNode: ObjectNode = jsonNodeMap match { + case map: ObjectNodeMap => + map.getObjectNode + case _ => + Utils.getSimpleObjectMapper.convertValue(jsonNodeMap, classOf[ObjectNode]) + } + + val partitionKey = PartitionKeyHelper.getPartitionKeyPath(objectNode, partitionKeyDefinition) + + val row = cosmosRowConverter.fromObjectNodeToRow(readSchema, + objectNode, + readConfig.schemaConversionMode) + + SparkRowItem(row, getPartitionKeyForFeedDiagnostics(partitionKey)).asInstanceOf[T] + } + } + ) + + // Collect all PK values upfront - the SDK now owns normalization and deduplication, + // so Spark preserves the caller's input as-is and relies on the SDK's set-based semantics. + // Callers should still dedupe their DataFrame input when practical to avoid extra work. + // + // NOTE on memory footprint: every PartitionKey from this iterator is materialized into a + // single ArrayList here, and the SDK in turn keeps the normalized set, the EPK->PK map, + // and the per-batch BatchDescriptor lists alive for the lifetime of the reader. For a + // single Spark partition with O(N) input rows this is O(N) memory on the executor; if + // upstream Spark partitioning sends millions of distinct PKs to one task this can become + // a noticeable allocation. Repartition upstream when N is very large. + private val PK_COUNT_LARGE_INPUT_WARN_THRESHOLD = 200000 + private lazy val pkList = { + val values = new java.util.ArrayList[PartitionKey]() + readManyPartitionKeys.foreach(values.add) + if (values.size() > PK_COUNT_LARGE_INPUT_WARN_THRESHOLD) { + log.logWarning( + s"ItemsPartitionReaderWithReadManyByPartitionKey received ${values.size()} partition " + + s"keys for a single Spark partition (feedRange=$feedRange). Large PK lists materialize " + + s"the full set in memory plus the SDK's normalized batch metadata; consider increasing " + + s"upstream Spark parallelism so each task processes <= " + + s"$PK_COUNT_LARGE_INPUT_WARN_THRESHOLD distinct partition keys.") + } + values + } + + private val endToEndTimeoutPolicy = + new CosmosEndToEndOperationLatencyPolicyConfigBuilder( + java.time.Duration.ofSeconds(CosmosConstants.readOperationEndToEndTimeoutInSeconds)) + .enable(true) + .build + + readManyOptionsImpl.setCosmosEndToEndOperationLatencyPolicyConfig(endToEndTimeoutPolicy) + + private trait CloseableSparkRowItemIterator { + def hasNext: Boolean + def next(): SparkRowItem + def close(): Unit + } + + private object EmptySparkRowItemIterator extends CloseableSparkRowItemIterator { + override def hasNext: Boolean = false + + override def next(): SparkRowItem = { + throw new java.util.NoSuchElementException("No items available for empty partition-key list.") + } + + override def close(): Unit = {} + } + + // Pass the full PK list to the SDK (which batches per physical partition internally). + // On transient I/O failures the retry iterator re-creates the underlying flux from the + // continuation token of the last fully-committed page, matching the pattern used by + // TransientIOErrorsRetryingIterator for queries and change feed. + private val isClosed = new AtomicBoolean(false) + private var iteratorOpt: Option[CloseableSparkRowItemIterator] = None + + private def getOrCreateIterator: CloseableSparkRowItemIterator = iteratorOpt match { + case Some(existing) => existing + case None => + val created = + if (pkList.isEmpty) { + EmptySparkRowItemIterator + } else { + new CloseableSparkRowItemIterator { + private val customQueryOpt = readConfig.customQuery.map(_.toSqlQuerySpec) + + // Factory that creates a CosmosPagedFlux from an optional continuation token. + // On the first call continuationToken is null (start from scratch); on retry + // it is the continuation token from the last fully-drained page. + // A fresh CosmosReadManyByPartitionKeysRequestOptions instance is created per + // call to avoid mutating the shared readManyOptions object, which would be + // fragile if the SDK ever stopped cloning options internally. + private val fluxFactory: String => CosmosPagedFlux[SparkRowItem] = { (continuationToken: String) => + val perCallOptions = new CosmosReadManyByPartitionKeysRequestOptions() + perCallOptions.setMaxConcurrentBatchPrefetch(readConfig.readManyByPkMaxConcurrentBatchPrefetch) + perCallOptions.setContinuationToken(continuationToken) + val perCallOptionsImpl = ImplementationBridgeHelpers + .CosmosReadManyByPartitionKeysRequestOptionsHelper + .getCosmosReadManyByPartitionKeysRequestOptionsAccessor + .getImpl(perCallOptions) + ThroughputControlHelper.populateThroughputControlGroupName(perCallOptionsImpl, readConfig.throughputControlConfig) + perCallOptionsImpl.setCosmosEndToEndOperationLatencyPolicyConfig(endToEndTimeoutPolicy) + if (operationContextAndListenerTuple.isDefined) { + perCallOptionsImpl.setOperationContextAndListenerTuple(operationContextAndListenerTuple.get) + } + perCallOptionsImpl + .setCustomItemSerializer(readManyOptionsImpl.getCustomItemSerializer) + customQueryOpt match { + case Some(query) => + cosmosAsyncContainer.readManyByPartitionKeys(pkList, query, perCallOptions, classOf[SparkRowItem]) + case None => + cosmosAsyncContainer.readManyByPartitionKeys(pkList, perCallOptions, classOf[SparkRowItem]) + } + } + + private val delegate = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + fluxFactory, + readConfig.maxItemCount, + readConfig.prefetchBufferSize, + operationContextAndListenerTuple, + classOf[SparkRowItem] + ) + + override def hasNext: Boolean = delegate.hasNext + + override def next(): SparkRowItem = delegate.next() + + override def close(): Unit = delegate.close() + } + } + + iteratorOpt = Some(created) + created + } + + private val rowSerializer: ExpressionEncoder.Serializer[Row] = RowSerializerPool.getOrCreateSerializer(readSchema) + + private def shouldLogDetailedFeedDiagnostics(): Boolean = { + diagnosticsConfig.mode.isDefined && + diagnosticsConfig.mode.get.equalsIgnoreCase(classOf[DetailedFeedDiagnosticsProvider].getName) + } + + private def getPartitionKeyForFeedDiagnostics(pkValue: PartitionKey): Option[PartitionKey] = { + if (shouldLogDetailedFeedDiagnostics()) { + Some(pkValue) + } else { + None + } + } + + private var currentRow: Option[Row] = None + + override def next(): Boolean = { + val hasMore = getOrCreateIterator.hasNext + if (hasMore) { + currentRow = Some(getOrCreateIterator.next().row) + } else { + currentRow = None + } + hasMore + } + + override def get(): InternalRow = { + cosmosRowConverter.fromRowToInternalRow( + currentRow.getOrElse(throw new NoSuchElementException("No current row - next() must be called first")), + rowSerializer) + } + + def getCurrentRow(): Row = { + currentRow.getOrElse(throw new NoSuchElementException("No current row - next() must be called first")) + } + + override def close(): Unit = { + if (isClosed.compareAndSet(false, true)) { + iteratorOpt.foreach(_.close()) + iteratorOpt = None + RowSerializerPool.returnSerializerToPool(readSchema, rowSerializer) + clientCacheItem.close() + if (throughputControlClientCacheItemOpt.isDefined) { + throughputControlClientCacheItemOpt.get.close() + } + } + } +} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingIterator.scala index b227d1a5ffb1..c949cd846f6c 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingIterator.scala @@ -50,7 +50,6 @@ private class TransientIOErrorsRetryingIterator[TSparkRow] 5 + CosmosConstants.readOperationEndToEndTimeoutInSeconds, scala.concurrent.duration.SECONDS) - private val rnd = Random // scalastyle:off null private val lastContinuationToken = new AtomicReference[String](null) // scalastyle:on null @@ -204,46 +203,17 @@ private class TransientIOErrorsRetryingIterator[TSparkRow] } private[spark] def executeWithRetry[T](methodName: String, func: () => T): T = { - val loop = new Breaks() - var returnValue: Option[T] = None - - loop.breakable { - while (true) { - val retryIntervalInMs = rnd.nextInt(maxRetryIntervalInMs) - - try { - returnValue = Some(func()) - retryCount.set(0) - loop.break - } - catch { - case cosmosException: CosmosException => - if (Exceptions.canBeTransientFailure(cosmosException.getStatusCode, cosmosException.getSubStatusCode)) { - val retryCountSnapshot = retryCount.incrementAndGet() - if (retryCountSnapshot > maxRetryCount) { - logError( - s"Too many transient failure retry attempts in TransientIOErrorsRetryingIterator.$methodName", - cosmosException) - throw cosmosException - } else { - logWarning( - s"Transient failure handled in TransientIOErrorsRetryingIterator.$methodName -" + - s" will be retried (attempt#$retryCountSnapshot) in ${retryIntervalInMs}ms", - cosmosException) - } - } else { - throw cosmosException - } - case other: Throwable => throw other - } - + TransientIOErrorsRetryingIterator.executeWithRetry( + "TransientIOErrorsRetryingIterator", + methodName, + func, + maxRetryCount, + maxRetryIntervalInMs, + retryCount, + () => { currentItemIterator = None currentFeedResponseIterator = None - Thread.sleep(retryIntervalInMs) - } - } - - returnValue.get + }) } private[this] def validateNextLsn(itemIterator: BufferedIterator[TSparkRow]): Boolean = { @@ -275,7 +245,7 @@ private class TransientIOErrorsRetryingIterator[TSparkRow] } } -private object TransientIOErrorsRetryingIterator extends BasicLoggingTrait { +private[spark] object TransientIOErrorsRetryingIterator extends BasicLoggingTrait { private val maxConcurrency = SparkUtils.getNumberOfHostCPUCores val executorService: ExecutorService = new ThreadPoolExecutor( @@ -293,4 +263,60 @@ private object TransientIOErrorsRetryingIterator extends BasicLoggingTrait { ) val executionContext = ExecutionContext.fromExecutorService(executorService) + + /** + * Shared retry logic for transient I/O failures. Both TransientIOErrorsRetryingIterator + * and TransientIOErrorsRetryingReadManyByPartitionKeyIterator delegate to this method + * to avoid duplicating the retry loop, backoff, and transient-failure classification. + */ + def executeWithRetry[T]( + callerName: String, + methodName: String, + func: () => T, + maxRetryCount: Long, + maxRetryIntervalInMs: Int, + retryCount: AtomicLong, + onRetry: () => Unit): T = { + + val rnd = Random + val loop = new Breaks() + var returnValue: Option[T] = None + + loop.breakable { + while (true) { + val retryIntervalInMs = rnd.nextInt(maxRetryIntervalInMs) + + try { + returnValue = Some(func()) + retryCount.set(0) + loop.break + } + catch { + case cosmosException: CosmosException => + if (Exceptions.canBeTransientFailure(cosmosException.getStatusCode, cosmosException.getSubStatusCode)) { + val retryCountSnapshot = retryCount.incrementAndGet() + if (retryCountSnapshot > maxRetryCount) { + logError( + s"Too many transient failure retry attempts in $callerName.$methodName", + cosmosException) + throw cosmosException + } else { + logWarning( + s"Transient failure handled in $callerName.$methodName -" + + s" will be retried (attempt#$retryCountSnapshot) in ${retryIntervalInMs}ms", + cosmosException) + } + } else { + throw cosmosException + } + case other: Throwable => throw other + } + + onRetry() + Thread.sleep(retryIntervalInMs) + } + } + + returnValue.get + } } diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala new file mode 100644 index 000000000000..6779a4117dd8 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIterator.scala @@ -0,0 +1,206 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.spark + +import com.azure.cosmos.implementation.OperationCancelledException +import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple +import com.azure.cosmos.models.FeedResponse +import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait +import com.azure.cosmos.util.{CosmosPagedFlux, CosmosPagedIterable} + +import java.util.concurrent.TimeoutException +import java.util.concurrent.atomic.{AtomicLong, AtomicReference} +import scala.concurrent.{Await, ExecutionContext, Future} + +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import + +/** + * Retry-safe iterator for readManyByPartitionKeys that uses continuation-token-based + * replay (matching the pattern of TransientIOErrorsRetryingIterator). + * + * On transient I/O failures the iterator re-creates the underlying CosmosPagedFlux from the + * continuation token of the last fully-committed page. A page is "committed" only after all + * its items have been drained by the caller; the continuation token is captured from each + * FeedResponse and used as the resume point on retry. This avoids the correctness issues of + * page-count-based skipping (where the server is not guaranteed to return the same page + * boundaries across requests). + * + * Note: transient I/O failures can only occur during {@code hasNext} (which fetches the next + * FeedResponse page from the network). Once a page has been fetched, iterating over its + * in-memory items ({@code next()}) performs no I/O and cannot trigger a retry. Therefore, + * partially-consumed pages are never replayed and the iterator provides exactly-once + * delivery in practice. + */ +private[spark] class TransientIOErrorsRetryingReadManyByPartitionKeyIterator[TSparkRow] +( + val cosmosPagedFluxFactory: String => CosmosPagedFlux[TSparkRow], + val pageSize: Int, + val pagePrefetchBufferSize: Int, + val operationContextAndListener: Option[OperationContextAndListenerTuple], + val classType: Class[TSparkRow] +) extends BufferedIterator[TSparkRow] with BasicLoggingTrait with AutoCloseable { + + private[spark] var maxRetryIntervalInMs = CosmosConstants.maxRetryIntervalForTransientFailuresInMs + private[spark] var maxRetryCount = CosmosConstants.maxRetryCountForTransientFailures + + private val maxPageRetrievalTimeout = scala.concurrent.duration.FiniteDuration( + 5 + CosmosConstants.readOperationEndToEndTimeoutInSeconds, + scala.concurrent.duration.SECONDS) + + // scalastyle:off null + private val lastContinuationToken = new AtomicReference[String](null) + private val pendingContinuationToken = new AtomicReference[String](null) + // scalastyle:on null + private val retryCount = new AtomicLong(0) + private lazy val operationContextString = operationContextAndListener match { + case Some(o) => if (o.getOperationContext != null) { + o.getOperationContext.toString + } else { + "n/a" + } + case None => "n/a" + } + + private[spark] var currentFeedResponseIterator: Option[BufferedIterator[FeedResponse[TSparkRow]]] = None + private[spark] var currentItemIterator: Option[BufferedIterator[TSparkRow]] = None + + override def hasNext: Boolean = { + executeWithRetry("hasNextInternal", () => hasNextInternal) + } + + private def hasNextInternal: Boolean = { + var returnValue: Option[Boolean] = None + + while (returnValue.isEmpty) { + returnValue = hasNextInternalCore + } + + returnValue.get + } + + private def hasNextInternalCore: Option[Boolean] = { + if (hasBufferedNext) { + Some(true) + } else { + // All items from the previous page have been drained — promote the pending + // continuation token so that any retry resumes from the NEXT page rather than + // replaying items the caller has already consumed. + val pending = pendingContinuationToken.getAndSet(null) + if (pending != null) { + lastContinuationToken.set(pending) + } + + val feedResponseIterator = currentFeedResponseIterator match { + case Some(existing) => existing + case None => + val newPagedFlux = cosmosPagedFluxFactory.apply(lastContinuationToken.get) + currentFeedResponseIterator = Some( + new CosmosPagedIterable[TSparkRow]( + newPagedFlux, + pageSize, + pagePrefetchBufferSize + ) + .iterableByPage() + .iterator + .asScala + .buffered + ) + + currentFeedResponseIterator.get + } + + val hasNext: Boolean = try { + Await.result( + Future { + feedResponseIterator.hasNext + }(TransientIOErrorsRetryingIterator.executionContext), + maxPageRetrievalTimeout) + } catch { + case endToEndTimeoutException: OperationCancelledException => + val message = s"End-to-end timeout hit when trying to retrieve the next page. " + + s"ContinuationToken: $lastContinuationToken, Context: $operationContextString" + logError(message, throwable = endToEndTimeoutException) + throw endToEndTimeoutException + + case timeoutException: TimeoutException => + val message = s"Attempting to retrieve the next page timed out. " + + s"ContinuationToken: $lastContinuationToken, Context: $operationContextString" + logError(message, timeoutException) + val exception = new OperationCancelledException(message, null) + exception.setStackTrace(timeoutException.getStackTrace) + throw exception + + case other: Throwable => throw other + } + + if (hasNext) { + val feedResponse = feedResponseIterator.next() + if (operationContextAndListener.isDefined) { + operationContextAndListener.get.getOperationListener.feedResponseProcessedListener( + operationContextAndListener.get.getOperationContext, + feedResponse) + } + val iteratorCandidate = feedResponse.getResults.iterator().asScala.buffered + // Store the continuation token from this page as "pending". It will only be promoted + // to lastContinuationToken (the retry resume point) after all items in this page have + // been drained by the caller. This ensures that on a transient failure mid-page the + // retry resumes from the PREVIOUS page's continuation (i.e. re-fetches the current + // page) and never skips items. + pendingContinuationToken.set(feedResponse.getContinuationToken) + + if (iteratorCandidate.hasNext) { + currentItemIterator = Some(iteratorCandidate) + Some(true) + } else { + // empty page interleaved - attempt to get next FeedResponse + None + } + } else { + // Flux exhausted + Some(false) + } + } + } + + private def hasBufferedNext: Boolean = { + currentItemIterator match { + case Some(iterator) => if (iterator.hasNext) { + true + } else { + currentItemIterator = None + false + } + case None => false + } + } + + override def next(): TSparkRow = { + currentItemIterator.get.next() + } + + override def head: TSparkRow = { + currentItemIterator.get.head + } + + private[spark] def executeWithRetry[T](methodName: String, func: () => T): T = { + TransientIOErrorsRetryingIterator.executeWithRetry( + "TransientIOErrorsRetryingReadManyByPartitionKeyIterator", + methodName, + func, + maxRetryCount, + maxRetryIntervalInMs, + retryCount, + () => { + currentItemIterator = None + currentFeedResponseIterator = None + pendingContinuationToken.set(null) + }) + } + override def close(): Unit = { + currentItemIterator = None + currentFeedResponseIterator = None + } +} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyIterator.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyIterator.scala index c51c5c1226e1..02feecd6901c 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyIterator.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyIterator.scala @@ -10,7 +10,7 @@ import com.azure.cosmos.models.{CosmosItemIdentity, CosmosReadManyRequestOptions import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait import java.time.Duration -import java.util.concurrent.{ExecutorService, SynchronousQueue, ThreadPoolExecutor, TimeUnit, TimeoutException} +import java.util.concurrent.TimeoutException import scala.concurrent.{Await, ExecutionContext, Future} // scalastyle:off underscore.import @@ -87,7 +87,7 @@ private[spark] class TransientIOErrorsRetryingReadManyIterator[TSparkRow] .getCosmosAsyncContainerAccessor .readMany(container, readManyFilterList.asJava, queryOptionsWithEnd2EndTimeout, classType) .block() - }(TransientIOErrorsRetryingReadManyIterator.executionContext), + }(TransientIOErrorsRetryingIterator.executionContext), maxPageRetrievalTimeout) } catch { case endToEndTimeoutException: OperationCancelledException => @@ -172,23 +172,3 @@ private[spark] class TransientIOErrorsRetryingReadManyIterator[TSparkRow] override def close(): Unit = {} } - -private object TransientIOErrorsRetryingReadManyIterator extends BasicLoggingTrait { - private val maxConcurrency = SparkUtils.getNumberOfHostCPUCores - - val executorService: ExecutorService = new ThreadPoolExecutor( - maxConcurrency, - maxConcurrency, - 0L, - TimeUnit.MILLISECONDS, - // A synchronous queue does not have any internal capacity, not even a capacity of one. - new SynchronousQueue(), - SparkUtils.daemonThreadFactory(), - // if all worker threads are busy, - // this policy makes the caller thread execute the task. - // This provides a simple feedback control mechanism that will slow down the rate that new tasks are submitted. - new ThreadPoolExecutor.CallerRunsPolicy() - ) - - val executionContext = ExecutionContext.fromExecutorService(executorService) -} \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala new file mode 100644 index 000000000000..3e536269f2fb --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala @@ -0,0 +1,29 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.spark.udf + +import com.azure.cosmos.spark.CosmosPartitionKeyHelper +import org.apache.spark.sql.api.java.UDF1 + +@SerialVersionUID(1L) +class GetCosmosPartitionKeyValue extends UDF1[Object, String] { + // Null is a valid partition-key value (JSON null). A null input is serialized as a + // single-level partition key with a JSON null component; parsing that string back via + // CosmosPartitionKeyHelper.tryParsePartitionKey yields a PartitionKey built with + // addNullValue(). If the caller instead wants PartitionKey.NONE semantics (absent PK + // field) they should filter the null row before calling this UDF and use the schema-matched + // readManyByPartitionKeys path with readManyByPk.nullHandling=None. That None mode is only + // supported for single-path partition keys; hierarchical partition keys reject it. + override def call(partitionKeyValue: Object): String = { + partitionKeyValue match { + case null => + CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List(null)) + // for subpartitions case - Seq covers both WrappedArray (Scala 2.12) and ArraySeq (Scala 2.13) + case seq: Seq[Any] => + CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(seq.map(_.asInstanceOf[Object]).toList) + case _ => + CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List(partitionKeyValue)) + } + } +} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala index 17f75e45a746..17a298d62131 100644 --- a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosConfigSpec.scala @@ -457,6 +457,7 @@ class CosmosConfigSpec extends UnitSpec with BasicLoggingTrait { config.runtimeFilteringEnabled shouldBe true config.readManyFilteringConfig.readManyFilteringEnabled shouldBe false config.readManyFilteringConfig.readManyFilterProperty shouldEqual "_itemIdentity" + config.readManyByPkTreatNullAsNone shouldBe false userConfig = Map( "spark.cosmos.read.forceEventualConsistency" -> "false", @@ -630,6 +631,47 @@ class CosmosConfigSpec extends UnitSpec with BasicLoggingTrait { config.customQuery.get.queryText shouldBe queryText } + it should "parse readManyByPk nullHandling configuration" in { + // Default (not specified) should treat null as JSON null (addNullValue) + var userConfig = Map( + "spark.cosmos.read.forceEventualConsistency" -> "false" + ) + var config = CosmosReadConfig.parseCosmosReadConfig(userConfig) + config.readManyByPkTreatNullAsNone shouldBe false + + // Explicit "Null" should treat null as JSON null (addNullValue) + userConfig = Map( + "spark.cosmos.read.forceEventualConsistency" -> "false", + "spark.cosmos.read.readManyByPk.nullHandling" -> "Null" + ) + config = CosmosReadConfig.parseCosmosReadConfig(userConfig) + config.readManyByPkTreatNullAsNone shouldBe false + + // Case-insensitive "null" + userConfig = Map( + "spark.cosmos.read.forceEventualConsistency" -> "false", + "spark.cosmos.read.readManyByPk.nullHandling" -> "null" + ) + config = CosmosReadConfig.parseCosmosReadConfig(userConfig) + config.readManyByPkTreatNullAsNone shouldBe false + + // "None" should treat null as PartitionKey.NONE (addNoneValue) + userConfig = Map( + "spark.cosmos.read.forceEventualConsistency" -> "false", + "spark.cosmos.read.readManyByPk.nullHandling" -> "None" + ) + config = CosmosReadConfig.parseCosmosReadConfig(userConfig) + config.readManyByPkTreatNullAsNone shouldBe true + + // Case-insensitive "none" + userConfig = Map( + "spark.cosmos.read.forceEventualConsistency" -> "false", + "spark.cosmos.read.readManyByPk.nullHandling" -> "none" + ) + config = CosmosReadConfig.parseCosmosReadConfig(userConfig) + config.readManyByPkTreatNullAsNone shouldBe true + } + it should "throw on invalid read configuration" in { val userConfig = Map( "spark.cosmos.read.schemaConversionMode" -> "not a valid value" diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala new file mode 100644 index 000000000000..d1d543b8a59f --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala @@ -0,0 +1,139 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.spark + +import com.azure.cosmos.models.{PartitionKey, PartitionKeyBuilder} + +class CosmosPartitionKeyHelperSpec extends UnitSpec { + //scalastyle:off multiple.string.literals + //scalastyle:off magic.number + + it should "return the correct partition key value string for single PK" in { + val pkString = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("pk1")) + pkString shouldEqual "pk([\"pk1\"])" + } + + it should "return the correct partition key value string for HPK" in { + val pkString = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("city1", "zip1")) + pkString shouldEqual "pk([\"city1\",\"zip1\"])" + } + + it should "return the correct partition key value string for 3-level HPK" in { + val pkString = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("a", "b", "c")) + pkString shouldEqual "pk([\"a\",\"b\",\"c\"])" + } + + it should "parse valid single PK string" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"myPkValue\"])") + pk.isDefined shouldBe true + pk.get shouldEqual new PartitionKey("myPkValue") + } + + it should "parse valid HPK string" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"city1\",\"zip1\"])") + pk.isDefined shouldBe true + val expected = new PartitionKeyBuilder().add("city1").add("zip1").build() + pk.get shouldEqual expected + } + + it should "parse valid 3-level HPK string" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"a\",\"b\",\"c\"])") + pk.isDefined shouldBe true + val expected = new PartitionKeyBuilder().add("a").add("b").add("c").build() + pk.get shouldEqual expected + } + + it should "roundtrip single PK" in { + val original = "pk([\"roundtrip\"])" + val parsed = CosmosPartitionKeyHelper.tryParsePartitionKey(original) + parsed.isDefined shouldBe true + val serialized = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("roundtrip")) + serialized shouldEqual original + } + + it should "roundtrip HPK" in { + val original = "pk([\"city\",\"zip\"])" + val parsed = CosmosPartitionKeyHelper.tryParsePartitionKey(original) + parsed.isDefined shouldBe true + val serialized = CosmosPartitionKeyHelper.getCosmosPartitionKeyValueString(List("city", "zip")) + serialized shouldEqual original + } + + it should "return None for malformed string" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("invalid_format") + pk.isDefined shouldBe false + } + + it should "return None for missing pk prefix" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("[\"value\"]") + pk.isDefined shouldBe false + } + + it should "be case-insensitive for parsing" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("PK([\"value\"])") + pk.isDefined shouldBe true + pk.get shouldEqual new PartitionKey("value") + } + + + it should "return None for malformed JSON inside pk() wrapper" in { + // Invalid JSON that would cause JsonProcessingException + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk({invalid json})") + pk.isDefined shouldBe false + } + + it should "return None for truncated JSON inside pk() wrapper" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"unterminated)") + pk.isDefined shouldBe false + } + + it should "parse single-path null as PartitionKey.NONE when treatNullAsNone is true" in { + val pk = CosmosPartitionKeyHelper.tryParsePartitionKey("pk([null])", treatNullAsNone = true) + + pk.isDefined shouldBe true + pk.get shouldEqual PartitionKey.NONE + } + + it should "throw for unsupported component types in the null-handling builder path" in { + val error = the[IllegalArgumentException] thrownBy { + CosmosPartitionKeyHelper.tryParsePartitionKey("pk([null,{\"nested\":\"value\"}])", treatNullAsNone = false) + } + + error.getMessage should include("Unsupported partition key component type") + error.getMessage should include("java.util.LinkedHashMap") + } + it should "throw a clear error when None nullHandling is used for hierarchical partition keys" in { + val error = the[IllegalArgumentException] thrownBy { + CosmosPartitionKeyHelper.tryParsePartitionKey("pk([\"Redmond\",null])", treatNullAsNone = true) + } + + error.getMessage should include(CosmosConfigNames.ReadManyByPkNullHandling) + error.getMessage should include("hierarchical partition keys") + } + + it should "reject addNoneValue in hierarchical partition keys" in { + val error = the[IllegalStateException] thrownBy { + new PartitionKeyBuilder().add("Redmond").addNoneValue().build() + } + + error.getMessage should include("PartitionKey.None can't be used with multiple paths") + } + + it should "throw IllegalArgumentException for null input" in { + val error = the[IllegalArgumentException] thrownBy { + CosmosPartitionKeyHelper.tryParsePartitionKey(null) + } + error.getMessage should include("must not be null") + } + + it should "throw IllegalArgumentException for null input with treatNullAsNone" in { + val error = the[IllegalArgumentException] thrownBy { + CosmosPartitionKeyHelper.tryParsePartitionKey(null, treatNullAsNone = true) + } + error.getMessage should include("must not be null") + } + + //scalastyle:on multiple.string.literals + //scalastyle:on magic.number +} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala new file mode 100644 index 000000000000..5c2d7b59836d --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala @@ -0,0 +1,158 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.spark + +import com.azure.cosmos.implementation.{CosmosClientMetadataCachesSnapshot, TestConfigurations, Utils} +import com.azure.cosmos.models.PartitionKey +import com.azure.cosmos.spark.diagnostics.DiagnosticsContext +import com.fasterxml.jackson.databind.node.ObjectNode +import org.apache.spark.MockTaskContext +import org.apache.spark.broadcast.Broadcast +import org.apache.spark.sql.types.{StringType, StructField, StructType} + +import java.util.UUID +import scala.collection.mutable.ListBuffer + +class ItemsPartitionReaderWithReadManyByPartitionKeyITest + extends IntegrationSpec + with Spark + with AutoCleanableCosmosContainersWithPkAsPartitionKey { + private val idProperty = "id" + private val pkProperty = "pk" + + //scalastyle:off multiple.string.literals + //scalastyle:off magic.number + + it should "be able to retrieve all items for given partition keys" in { + val container = cosmosClient.getDatabase(cosmosDatabase).getContainer(cosmosContainersWithPkAsPartitionKey) + + // Create items with known PK values + val partitionKeyDefinition = container.read().block().getProperties.getPartitionKeyDefinition + val allItemsByPk = scala.collection.mutable.Map[String, ListBuffer[ObjectNode]]() + val pkValues = List("pkA", "pkB", "pkC") + + for (pk <- pkValues) { + allItemsByPk(pk) = ListBuffer[ObjectNode]() + for (_ <- 1 to 5) { + val objectNode = Utils.getSimpleObjectMapper.createObjectNode() + objectNode.put(idProperty, UUID.randomUUID().toString) + objectNode.put(pkProperty, pk) + container.createItem(objectNode).block() + allItemsByPk(pk) += objectNode + } + } + + val config = Map( + "spark.cosmos.accountEndpoint" -> TestConfigurations.HOST, + "spark.cosmos.accountKey" -> TestConfigurations.MASTER_KEY, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, + "spark.cosmos.read.inferSchema.enabled" -> "true", + "spark.cosmos.applicationName" -> "ReadManyByPKTest" + ) + + val readSchema = StructType(Seq( + StructField(idProperty, StringType, false), + StructField(pkProperty, StringType, false) + )) + + val diagnosticsContext = DiagnosticsContext(UUID.randomUUID(), "") + val diagnosticsConfig = DiagnosticsConfig.parseDiagnosticsConfig(config) + val cosmosClientMetadataCachesSnapshots = getCosmosClientMetadataCachesSnapshots() + + // Read items for pkA and pkB (not pkC) + val targetPks = List("pkA", "pkB") + val pkIterator = targetPks.map(pk => new PartitionKey(pk)).iterator + + val reader = ItemsPartitionReaderWithReadManyByPartitionKey( + config, + NormalizedRange("", "FF"), + readSchema, + diagnosticsContext, + cosmosClientMetadataCachesSnapshots, + diagnosticsConfig, + "", + MockTaskContext.mockTaskContext(), + pkIterator + ) + + val cosmosRowConverter = CosmosRowConverter.get(CosmosSerializationConfig.parseSerializationConfig(config)) + val itemsReadFromReader = ListBuffer[ObjectNode]() + while (reader.next()) { + itemsReadFromReader += cosmosRowConverter.fromInternalRowToObjectNode(reader.get(), readSchema) + } + + // Should have 10 items (5 for pkA + 5 for pkB) + itemsReadFromReader.size shouldEqual 10 + + // All items should be from pkA or pkB + itemsReadFromReader.foreach(item => { + val pk = item.get(pkProperty).asText() + targetPks should contain(pk) + }) + + // Validate all expected IDs are present + val expectedIds = (allItemsByPk("pkA") ++ allItemsByPk("pkB")).map(_.get(idProperty).asText()).toSet + val actualIds = itemsReadFromReader.map(_.get(idProperty).asText()).toSet + actualIds shouldEqual expectedIds + + reader.close() + } + + it should "return empty results for non-existent partition keys" in { + val config = Map( + "spark.cosmos.accountEndpoint" -> TestConfigurations.HOST, + "spark.cosmos.accountKey" -> TestConfigurations.MASTER_KEY, + "spark.cosmos.database" -> cosmosDatabase, + "spark.cosmos.container" -> cosmosContainersWithPkAsPartitionKey, + "spark.cosmos.read.inferSchema.enabled" -> "true", + "spark.cosmos.applicationName" -> "ReadManyByPKEmptyTest" + ) + + val readSchema = StructType(Seq( + StructField(idProperty, StringType, false), + StructField(pkProperty, StringType, false) + )) + + val diagnosticsContext = DiagnosticsContext(UUID.randomUUID(), "") + val diagnosticsConfig = DiagnosticsConfig.parseDiagnosticsConfig(config) + val cosmosClientMetadataCachesSnapshots = getCosmosClientMetadataCachesSnapshots() + + val pkIterator = List(new PartitionKey("nonExistentPk")).iterator + + val reader = ItemsPartitionReaderWithReadManyByPartitionKey( + config, + NormalizedRange("", "FF"), + readSchema, + diagnosticsContext, + cosmosClientMetadataCachesSnapshots, + diagnosticsConfig, + "", + MockTaskContext.mockTaskContext(), + pkIterator + ) + + val itemsReadFromReader = ListBuffer[ObjectNode]() + val cosmosRowConverter = CosmosRowConverter.get(CosmosSerializationConfig.parseSerializationConfig(config)) + while (reader.next()) { + itemsReadFromReader += cosmosRowConverter.fromInternalRowToObjectNode(reader.get(), readSchema) + } + + itemsReadFromReader.size shouldEqual 0 + reader.close() + } + + private def getCosmosClientMetadataCachesSnapshots(): Broadcast[CosmosClientMetadataCachesSnapshots] = { + val cosmosClientMetadataCachesSnapshot = new CosmosClientMetadataCachesSnapshot() + cosmosClientMetadataCachesSnapshot.serialize(cosmosClient) + + spark.sparkContext.broadcast( + CosmosClientMetadataCachesSnapshots( + cosmosClientMetadataCachesSnapshot, + Option.empty[CosmosClientMetadataCachesSnapshot])) + } + + //scalastyle:on multiple.string.literals + //scalastyle:on magic.number +} diff --git a/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIteratorSpec.scala b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIteratorSpec.scala new file mode 100644 index 000000000000..64047f41634c --- /dev/null +++ b/sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/TransientIOErrorsRetryingReadManyByPartitionKeyIteratorSpec.scala @@ -0,0 +1,356 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.spark + +import com.azure.cosmos.CosmosException +import com.azure.cosmos.implementation.SparkRowItem +import com.azure.cosmos.models.{FeedResponse, ModelBridgeInternal} +import com.azure.cosmos.spark.diagnostics.BasicLoggingTrait +import com.azure.cosmos.util.UtilBridgeInternal +import com.fasterxml.jackson.core.JsonProcessingException +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.node.ObjectNode +import reactor.core.publisher.Flux + +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicLong + +// scalastyle:off underscore.import +import scala.collection.JavaConverters._ +// scalastyle:on underscore.import + +//scalastyle:off magic.number +//scalastyle:off multiple.string.literals +class TransientIOErrorsRetryingReadManyByPartitionKeyIteratorSpec extends UnitSpec with BasicLoggingTrait { + + private val rnd = scala.util.Random + private val pageSize = 2 + private val cosmosSerializationConfig = CosmosSerializationConfig( + SerializationInclusionModes.Always, + SerializationDateTimeConversionModes.Default + ) + + private val cosmosRowConverter = CosmosRowConverter.get(cosmosSerializationConfig) + private val objectMapper = new ObjectMapper + + "TransientIOErrors" should "be retried without duplicates or missing records" in { + + val pageCount = 30 + val transientErrorCount = new AtomicLong(0) + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + continuationToken => generateMockedCosmosPagedFlux( + continuationToken, pageCount, 0.2, transientErrorCount, injectEmptyPages = false), + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + iterator.maxRetryIntervalInMs = 5 + + val items = drainAll(iterator) + items.size shouldEqual (pageCount * pageSize) + + transientErrorCount.get > 0 shouldEqual true + assertNoDuplicates(items) + } + + "TransientIOErrors" should "be retried without duplicates when empty pages exist" in { + + val pageCount = 30 + val transientErrorCount = new AtomicLong(0) + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + continuationToken => generateMockedCosmosPagedFlux( + continuationToken, pageCount, 0.2, transientErrorCount, injectEmptyPages = true), + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + iterator.maxRetryIntervalInMs = 5 + + // Pages 2,4,6,8,10,12,14,16,18,20 are empty (10 empty pages out of first 20) + val items = drainAll(iterator) + items.size shouldEqual ((pageCount - 10) * pageSize) + + transientErrorCount.get > 0 shouldEqual true + assertNoDuplicates(items) + } + + "Continuation token replay" should "not re-yield already consumed items after transient error" in { + + // Inject a transient error at a deterministic point (after page 3) so we can assert + // that pages 1-3's items are not duplicated after the retry. + val pageCount = 10 + val errorInjectedAfterPage = 3 + val errorInjected = new AtomicLong(0) + val factoryCallCount = new AtomicLong(0) + + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + continuationToken => { + factoryCallCount.incrementAndGet() + generateMockedCosmosPagedFluxWithDeterministicError( + continuationToken, pageCount, errorInjectedAfterPage, errorInjected) + }, + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + iterator.maxRetryIntervalInMs = 5 + + val items = drainAll(iterator) + items.size shouldEqual (pageCount * pageSize) + + errorInjected.get shouldEqual 1 + factoryCallCount.get shouldEqual 2 // initial + 1 retry + assertNoDuplicates(items) + } + + "Continuation token" should "be passed to factory on retry" in { + + val pageCount = 30 + val capturedContinuationTokens = new java.util.concurrent.CopyOnWriteArrayList[String]() + val transientErrorCount = new AtomicLong(0) + + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + continuationToken => { + capturedContinuationTokens.add(continuationToken) + generateMockedCosmosPagedFlux( + continuationToken, pageCount, 0.15, transientErrorCount, injectEmptyPages = false) + }, + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + iterator.maxRetryIntervalInMs = 5 + + drainAll(iterator) + + transientErrorCount.get > 0 shouldEqual true + + // First call should have null continuation token (start from beginning) + capturedContinuationTokens.get(0) shouldEqual null + + // Subsequent retry calls should have non-null continuation tokens + // (resume from last committed page) + val retryTokens = capturedContinuationTokens.asScala.drop(1) + retryTokens should not be empty + retryTokens.foreach(_ should not be null) + } + + "Non-transient errors" should "not be retried and propagate immediately" in { + + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + _ => { + val flux = Flux.error[FeedResponse[SparkRowItem]](new DummyNonTransientCosmosException) + UtilBridgeInternal.createCosmosPagedFlux(_ => flux) + }, + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + iterator.maxRetryIntervalInMs = 5 + + val thrown = the[CosmosException] thrownBy { + iterator.hasNext + } + thrown.getStatusCode shouldEqual 404 + } + + "Iterator close" should "clear internal state without extra factory calls" in { + + val factoryCallCount = new AtomicLong(0) + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + continuationToken => { + factoryCallCount.incrementAndGet() + generateMockedCosmosPagedFlux( + continuationToken, 30, 0.0, new AtomicLong(0), injectEmptyPages = false) + }, + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + + iterator.hasNext shouldEqual true + iterator.next() + + factoryCallCount.get shouldEqual 1 + + iterator.currentFeedResponseIterator.isDefined shouldEqual true + + iterator.close() + + iterator.currentFeedResponseIterator shouldEqual None + iterator.currentItemIterator shouldEqual None + + factoryCallCount.get shouldEqual 1 + } + + "Iterator close" should "be safe to call multiple times" in { + + val iterator = new TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]( + continuationToken => generateMockedCosmosPagedFlux( + continuationToken, 30, 0.0, new AtomicLong(0), injectEmptyPages = false), + pageSize, + 1, + None, + classOf[SparkRowItem] + ) + + iterator.hasNext shouldEqual true + iterator.next() + + iterator.close() + iterator.close() + iterator.close() + // No exception should be thrown + } + + // --- helpers --- + + private def drainAll(iterator: TransientIOErrorsRetryingReadManyByPartitionKeyIterator[SparkRowItem]): List[SparkRowItem] = { + val buffer = scala.collection.mutable.ListBuffer[SparkRowItem]() + while (iterator.hasNext) { + buffer += iterator.next() + } + buffer.toList + } + + private def assertNoDuplicates(items: List[SparkRowItem]): Unit = { + val ids = items.map(item => item.row.getString(0)) // "id" field + ids.size shouldEqual ids.distinct.size + } + + @throws[JsonProcessingException] + private def getDocumentDefinition(documentId: String) = { + val json = s"""{"id":"$documentId"}""" + val node = objectMapper.readValue(json, classOf[ObjectNode]) + SparkRowItem( + cosmosRowConverter.fromObjectNodeToRow( + ItemsTable.defaultSchemaForInferenceDisabled, + node, + SchemaConversionModes.Strict + ), + None) + } + + private def generateMockedCosmosPagedFlux( + continuationToken: String, + pageCount: Int, + errorRate: Double, + transientErrorCounter: AtomicLong, + injectEmptyPages: Boolean + ) = { + + require(pageCount > 20) + + val flux = generateFeedResponseFlux( + "Batch", pageCount, errorRate, + Option(continuationToken), transientErrorCounter, injectEmptyPages) + + UtilBridgeInternal.createCosmosPagedFlux(_ => flux) + } + + private def generateMockedCosmosPagedFluxWithDeterministicError( + continuationToken: String, + pageCount: Int, + errorAfterPage: Int, + errorInjected: AtomicLong + ) = { + + val flux = generateFeedResponseFluxWithDeterministicError( + "Batch", pageCount, errorAfterPage, + Option(continuationToken), errorInjected) + + UtilBridgeInternal.createCosmosPagedFlux(_ => flux) + } + + private def generateFeedResponseFlux( + prefix: String, + pageCount: Int, + errorRate: Double, + requestContinuationToken: Option[String], + transientErrorCounter: AtomicLong, + injectEmptyPages: Boolean + ): Flux[FeedResponse[SparkRowItem]] = { + + val responses = Array.range(1, pageCount + 1) + .map(i => generateFeedResponse( + prefix, i, + if (injectEmptyPages && i > 1 && i <= 20 && i % 2 == 0) -1 else 1)) + .filter(response => requestContinuationToken.isEmpty || + requestContinuationToken.get == null || + requestContinuationToken.get < response.getContinuationToken) + + Flux + .fromArray(responses) + .map(response => if (errorRate > 0 && rnd.nextDouble() < errorRate) { + transientErrorCounter.incrementAndGet() + throw new DummyTransientCosmosException + } else { + response + }) + } + + private def generateFeedResponseFluxWithDeterministicError( + prefix: String, + pageCount: Int, + errorAfterPage: Int, + requestContinuationToken: Option[String], + errorInjected: AtomicLong + ): Flux[FeedResponse[SparkRowItem]] = { + + val responses = Array.range(1, pageCount + 1) + .map(i => generateFeedResponse(prefix, i, 1)) + .filter(response => requestContinuationToken.isEmpty || + requestContinuationToken.get == null || + requestContinuationToken.get < response.getContinuationToken) + + Flux + .fromArray(responses) + .map(response => { + // Extract page sequence number from continuation token to inject error deterministically + val token = response.getContinuationToken + val pageNum = token.split("_Page")(1).split("_")(0).toInt + if (pageNum == errorAfterPage + 1 && errorInjected.compareAndSet(0, 1)) { + throw new DummyTransientCosmosException + } else { + response + } + }) + } + + private def generateFeedResponse( + prefix: String, + pageSequenceNumber: Int, + documentStartIndex: Int + ): FeedResponse[SparkRowItem] = { + + val continuationToken = f"$prefix%s_Page$pageSequenceNumber%05d_ContinuationToken" + val items = if (documentStartIndex < 0) { + Array.empty[SparkRowItem] + } else { + val id1 = f"$prefix%s_Page$pageSequenceNumber%05d_$documentStartIndex%05d" + val id2 = f"$prefix%s_Page$pageSequenceNumber%05d_${documentStartIndex + 1}%05d" + Array[SparkRowItem](getDocumentDefinition(id1), getDocumentDefinition(id2)) + } + + val r = ModelBridgeInternal.createFeedResponse( + items.toList.asJava, + new ConcurrentHashMap[String, String] + ) + ModelBridgeInternal.setFeedResponseContinuationToken(continuationToken, r) + r + } + + private class DummyTransientCosmosException + extends CosmosException(500, "Dummy Internal Server Error") + + private class DummyNonTransientCosmosException + extends CosmosException(404, "Dummy Not Found Error") +} +//scalastyle:on magic.number +//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 33fa94dda6c0..28c503b4dfec 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 @@ -3,6 +3,8 @@ ### 4.48.0-beta.1 (Unreleased) #### Features Added +* Added new `CosmosItemsDataSource.readManyByPartitionKeys` Spark function to execute bulk queries by a list of pk-values with better efficiency. Configure null handling via `spark.cosmos.read.readManyByPk.nullHandling` - default `Null` treats a null PK column as JSON null (`addNullValue`), `None` treats it as `PartitionKey.NONE` (`addNoneValue` / `NOT IS_DEFINED`). These route to different physical partitions - picking the wrong mode silently returns zero rows. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added Spark config `spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch` (default `1`) to bound the per-task prefetch parallelism the SDK uses inside `readManyByPartitionKeys`. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java index 5275af383f2a..5ece83be3156 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosMultiHashTest.java @@ -33,6 +33,7 @@ import java.time.Duration; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Objects; @@ -51,6 +52,7 @@ public class CosmosMultiHashTest extends TestSuiteBase { private CosmosClient client; private CosmosDatabase createdDatabase; private CosmosContainer createdMultiHashContainer; + private CosmosContainer createdNestedPathContainer; @Factory(dataProvider = "clientBuilders") public CosmosMultiHashTest(CosmosClientBuilder clientBuilder) { @@ -79,6 +81,19 @@ public void before_CosmosMultiHashTest() { //MultiHash collection read createdMultiHashContainer = createdDatabase.getContainer(collectionName); + + String nestedPathCollectionName = UUID.randomUUID().toString(); + PartitionKeyDefinition nestedPartitionKeyDefinition = new PartitionKeyDefinition(); + nestedPartitionKeyDefinition.setKind(PartitionKind.HASH); + nestedPartitionKeyDefinition.setVersion(PartitionKeyDefinitionVersion.V2); + nestedPartitionKeyDefinition.setPaths(Collections.singletonList("/address/city")); + + CosmosContainerProperties nestedPathContainerProperties = getCollectionDefinition( + nestedPathCollectionName, + nestedPartitionKeyDefinition); + + createdDatabase.createContainer(nestedPathContainerProperties); + createdNestedPathContainer = createdDatabase.getContainer(nestedPathCollectionName); } @AfterClass(groups = {"emulator"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) @@ -155,6 +170,42 @@ private void validateResponse(FeedResponse response, .collect(Collectors.toList()) ); } + @Test(groups = { "emulator" }, timeOut = TIMEOUT) + public void readManySupportsNestedPartitionKeyPaths() { + String city = "nested-readmany-" + UUID.randomUUID(); + + ObjectNode firstItem = createNestedPartitionKeyItem(UUID.randomUUID().toString(), city, "downtown"); + ObjectNode secondItem = createNestedPartitionKeyItem(UUID.randomUUID().toString(), city, "overlake"); + + createdNestedPathContainer.createItem(firstItem); + createdNestedPathContainer.createItem(secondItem); + + List itemList = new ArrayList<>(); + itemList.add(new CosmosItemIdentity(new PartitionKey(city), firstItem.get("id").asText())); + itemList.add(new CosmosItemIdentity(new PartitionKey(city), secondItem.get("id").asText())); + + FeedResponse documentFeedResponse = createdNestedPathContainer.readMany(itemList, ObjectNode.class); + validateResponse(documentFeedResponse, itemList); + } + + @Test(groups = { "emulator" }, timeOut = TIMEOUT) + public void readAllItemsSupportsNestedPartitionKeyPaths() { + String city = "nested-readall-" + UUID.randomUUID(); + + ObjectNode firstItem = createNestedPartitionKeyItem(UUID.randomUUID().toString(), city, "north"); + ObjectNode secondItem = createNestedPartitionKeyItem(UUID.randomUUID().toString(), city, "south"); + ObjectNode otherItem = createNestedPartitionKeyItem(UUID.randomUUID().toString(), city + "-other", "east"); + + createdNestedPathContainer.createItem(firstItem); + createdNestedPathContainer.createItem(secondItem); + createdNestedPathContainer.createItem(otherItem); + + CosmosPagedIterable readAllResults = + createdNestedPathContainer.readAllItems(new PartitionKey(city), ObjectNode.class); + + assertThat(readAllResults.stream().map(item -> item.get("id").asText()).collect(Collectors.toList())) + .containsExactlyInAnyOrder(firstItem.get("id").asText(), secondItem.get("id").asText()); + } @Test(groups = { "emulator" }, timeOut = TIMEOUT) private void validateDocCRUDAndQuery() throws InterruptedException { @@ -550,6 +601,19 @@ private void extractPartitionKeyFromDocumentTests() { } } + private ObjectNode createNestedPartitionKeyItem(String id, String city, String neighborhood) { + ObjectNode item = new ObjectNode(JSON_NODE_FACTORY_INSTANCE); + item.put("id", id); + item.put("type", "nestedPartitionKeyRegression"); + + ObjectNode address = new ObjectNode(JSON_NODE_FACTORY_INSTANCE); + address.put("city", city); + address.put("neighborhood", neighborhood); + item.set("address", address); + + return item; + } + private ArrayList createItems() { ArrayList docs = new ArrayList<>(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java index 06de8524bcbd..f8e2979f544d 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -920,6 +920,29 @@ public void cosmosAsyncContainer( "readMany", samplingRate); mockTracer.reset(); + List partitionKeys = createdDocs + .stream() + .map(CosmosItemIdentity::getPartitionKey) + .collect(Collectors.toList()); + feedItemResponse = cosmosAsyncContainer + .readManyByPartitionKeys(partitionKeys, ObjectNode.class) + .byPage(1) + .blockFirst(); + assertThat(feedItemResponse).isNotNull(); + assertThat(feedItemResponse.getResults()).isNotEmpty(); + verifyTracerAttributes( + mockTracer, + "readManyByPartitionKeys." + cosmosAsyncContainer.getId(), + cosmosAsyncDatabase.getId(), + cosmosAsyncContainer.getId(), + feedItemResponse.getCosmosDiagnostics(), + null, + useLegacyTracing, + enableRequestLevelTracing, + forceThresholdViolations, + "readManyByPartitionKeys", + samplingRate); + mockTracer.reset(); } @Test(groups = { "fast", "simple" }, timeOut = 10 * TIMEOUT) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FITests_readManyByPartitionKeysAfterCreation.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FITests_readManyByPartitionKeysAfterCreation.java new file mode 100644 index 000000000000..f618951522b2 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FITests_readManyByPartitionKeysAfterCreation.java @@ -0,0 +1,307 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos; + +import com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions; +import com.azure.cosmos.models.FeedRange; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.test.faultinjection.CosmosFaultInjectionHelper; +import com.azure.cosmos.test.faultinjection.FaultInjectionConditionBuilder; +import com.azure.cosmos.test.faultinjection.FaultInjectionEndpointBuilder; +import com.azure.cosmos.test.faultinjection.FaultInjectionOperationType; +import com.azure.cosmos.test.faultinjection.FaultInjectionResultBuilders; +import com.azure.cosmos.test.faultinjection.FaultInjectionRule; +import com.azure.cosmos.test.faultinjection.FaultInjectionRuleBuilder; +import com.azure.cosmos.test.faultinjection.FaultInjectionServerErrorType; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.TestConfigurations; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.apache.commons.lang3.ArrayUtils; +import org.testng.annotations.Test; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import java.util.function.BiConsumer; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +public class FITests_readManyByPartitionKeysAfterCreation + extends FaultInjectionWithAvailabilityStrategyTestsBase { + + @Test(groups = {"fi-multi-master"}, dataProvider = "testConfigs_readManyByPartitionKeysAfterCreation", retryAnalyzer = FlakyTestRetryAnalyzer.class) + public void readManyByPartitionKeysAfterCreation( + String testCaseId, + Duration endToEndTimeout, + ThresholdBasedAvailabilityStrategy availabilityStrategy, + CosmosRegionSwitchHint regionSwitchHint, + ConnectionMode connectionMode, + Function readManyByPkOperation, + BiConsumer faultInjectionCallback, + BiConsumer validateStatusCode, + int expectedDiagnosticsContextCount, + Consumer[] firstDiagnosticsContextValidations, + Consumer[] otherDiagnosticsContextValidations, + Consumer responseValidator, + int numberOfOtherDocumentsWithSameId, + int numberOfOtherDocumentsWithSamePk, + boolean shouldInjectPreferredRegionsInClient) { + + execute( + testCaseId, + endToEndTimeout, + availabilityStrategy, + regionSwitchHint, + null, + notSpecifiedWhetherIdempotentWriteRetriesAreEnabled, + ArrayUtils.toArray(FaultInjectionOperationType.QUERY_ITEM), + readManyByPkOperation, + faultInjectionCallback, + validateStatusCode, + expectedDiagnosticsContextCount, + firstDiagnosticsContextValidations, + otherDiagnosticsContextValidations, + responseValidator, + numberOfOtherDocumentsWithSameId, + numberOfOtherDocumentsWithSamePk, + false, + connectionMode, + shouldInjectPreferredRegionsInClient); + } + + /** + * Validates continuation-token resume after fault injection causes a deterministic error. + * + * Strategy: use FeedRange-scoped fault injection so that queries against one physical + * partition fail while queries against other partitions succeed. With batch size 1, + * readManyByPartitionKeys processes one PK-batch at a time sequentially. The first + * batch(es) targeting the non-faulted partition succeed and emit pages with continuation + * tokens. When iteration reaches the faulted partition, the error surfaces to the caller. + * + * 1. Create documents across multiple PKs (spread across partitions) + * 2. Collect baseline (no faults) — all ids + * 3. Get feed ranges; pick the second one to fault + * 4. Inject sustained SERVICE_UNAVAILABLE scoped to that feed range + * 5. Iterate page-by-page; collect items + continuation tokens from successful pages + * 6. When error occurs: validate it's expected, capture last good continuation + * 7. Disable the fault injection rule + * 8. Resume from the last good continuation token + * 9. Assert: union of items before error + items from resume = all baseline items, no duplicates + */ + @Test(groups = {"fi-multi-master"}, timeOut = 180000, retryAnalyzer = FlakyTestRetryAnalyzer.class) + public void readManyByPartitionKeys_continuationResumeAfterFaultInjection() { + + String originalBatchSize = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + try { + // batch size 1 = one PK per batch = sequential processing across partitions + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "1"); + + CosmosAsyncClient client = new CosmosClientBuilder() + .endpoint(TestConfigurations.HOST) + .key(TestConfigurations.MASTER_KEY) + .contentResponseOnWriteEnabled(true) + .directMode() + .buildAsyncClient(); + + try { + CosmosAsyncContainer container = client + .getDatabase(this.getTestDatabaseId()) + .getContainer(this.getTestContainerId()); + + String uniqueTag = UUID.randomUUID().toString().substring(0, 8); + + // Create items across 3 PKs, 3 items each = 9 items total + List pkValues = Arrays.asList( + "ctResumePk1_" + uniqueTag, + "ctResumePk2_" + uniqueTag, + "ctResumePk3_" + uniqueTag); + + List allCreatedItems = new ArrayList<>(); + for (String pk : pkValues) { + for (int i = 0; i < 3; i++) { + ObjectNode item = com.azure.cosmos.implementation.Utils + .getSimpleObjectMapper().createObjectNode(); + item.put("id", UUID.randomUUID().toString()); + item.put("mypk", pk); + container.createItem(item).block(); + allCreatedItems.add(item); + } + } + + List partitionKeys = pkValues.stream() + .map(PartitionKey::new) + .collect(Collectors.toList()); + + // Step 1: Baseline — drain all pages without faults to know the complete set of ids + List> baselinePages = container + .readManyByPartitionKeys(partitionKeys, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(baselinePages).isNotNull(); + assertThat(baselinePages.size()).isGreaterThan(1); // with batch size 1 there must be multiple pages + + List baselineIds = baselinePages.stream() + .flatMap(p -> p.getResults().stream()) + .map(n -> n.get("id").asText()) + .sorted() + .collect(Collectors.toList()); + assertThat(baselineIds).hasSize(9); + assertThat(baselineIds).doesNotHaveDuplicates(); + + // Step 2: Get feed ranges and pick the LAST one to fault. + // With batch size 1, readManyByPartitionKeys processes batches sorted by EPK. + // Faulting the last feed range ensures the first batches succeed (giving us + // pages with continuation tokens) before the faulted partition is reached. + List feedRanges = container.getFeedRanges().block(); + assertThat(feedRanges).isNotNull(); + assertThat(feedRanges.size()).isGreaterThanOrEqualTo(1); + FeedRange faultedFeedRange = feedRanges.get(feedRanges.size() - 1); + + // Step 3: Inject sustained SERVICE_UNAVAILABLE scoped to the last feed range + FaultInjectionRule partitionScopedRule = new FaultInjectionRuleBuilder( + "readManyByPk-ct-resume-partition-scoped") + .condition(new FaultInjectionConditionBuilder() + .operationType(FaultInjectionOperationType.QUERY_ITEM) + .endpoints(new FaultInjectionEndpointBuilder(faultedFeedRange) + .replicaCount(4) + .includePrimary(true) + .build()) + .build()) + .result(FaultInjectionResultBuilders + .getResultBuilder(FaultInjectionServerErrorType.SERVICE_UNAVAILABLE) + .build()) + .duration(Duration.ofSeconds(120)) + .build(); + + CosmosFaultInjectionHelper + .configureFaultInjectionRules(container, Collections.singletonList(partitionScopedRule)) + .block(); + + // Step 4: Drain page-by-page. Pages from non-faulted partitions succeed; + // when the faulted partition is reached, the error surfaces. + List itemsBeforeError = new ArrayList<>(); + String lastGoodContinuation = null; + boolean errorOccurred = false; + + CosmosEndToEndOperationLatencyPolicyConfig e2ePolicy = + new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofSeconds(10)) + .enable(true) + .build(); + + CosmosReadManyByPartitionKeysRequestOptions faultOptions = + new CosmosReadManyByPartitionKeysRequestOptions(); + faultOptions.setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); + + try { + // Use Flux iteration (toIterable) so we can capture per-page state + for (FeedResponse page : container + .readManyByPartitionKeys(partitionKeys, faultOptions, ObjectNode.class) + .byPage() + .toIterable()) { + + for (ObjectNode item : page.getResults()) { + itemsBeforeError.add(item.get("id").asText()); + } + if (page.getContinuationToken() != null) { + lastGoodContinuation = page.getContinuationToken(); + } + } + } catch (Exception e) { + errorOccurred = true; + } + + // Step 5: The fault injection MUST have caused an error — pages from the + // faulted partition cannot succeed with SERVICE_UNAVAILABLE on all replicas. + assertThat(errorOccurred) + .as("Fault injection on the last feed range must cause an error") + .isTrue(); + + // We must have captured at least one continuation token from successful pages + assertThat(lastGoodContinuation) + .as("At least one page must have succeeded before the faulted partition") + .isNotNull(); + + // Items collected so far must be a strict subset of the baseline + assertThat(itemsBeforeError).doesNotHaveDuplicates(); + assertThat(itemsBeforeError.size()).isGreaterThan(0); + assertThat(itemsBeforeError.size()).isLessThan(baselineIds.size()); + + // Step 6: Disable fault injection rule + partitionScopedRule.disable(); + + // Step 7: Resume from the last good continuation token + CosmosReadManyByPartitionKeysRequestOptions resumeOptions = + new CosmosReadManyByPartitionKeysRequestOptions(); + resumeOptions.setContinuationToken(lastGoodContinuation); + + List> resumedPages = container + .readManyByPartitionKeys(partitionKeys, resumeOptions, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(resumedPages).isNotNull(); + + List resumedIds = resumedPages.stream() + .flatMap(p -> p.getResults().stream()) + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + + // Step 8: Assert completeness — union of before + after = all baseline items + List combined = new ArrayList<>(itemsBeforeError); + combined.addAll(resumedIds); + + assertThat(combined).doesNotHaveDuplicates(); + assertThat(combined).hasSameElementsAs(baselineIds); + + // Cleanup + for (ObjectNode item : allCreatedItems) { + try { + container.deleteItem( + item.get("id").asText(), + new PartitionKey(item.get("mypk").asText())).block(); + } catch (Exception ignore) { } + } + + } finally { + safeClose(client); + } + } finally { + if (originalBatchSize != null) { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalBatchSize); + } else { + System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + } + } + } + + // Helper to access testDatabaseId from base class + private String getTestDatabaseId() { + try { + java.lang.reflect.Field f = FaultInjectionWithAvailabilityStrategyTestsBase.class.getDeclaredField("testDatabaseId"); + f.setAccessible(true); + return (String) f.get(this); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private String getTestContainerId() { + try { + java.lang.reflect.Field f = FaultInjectionWithAvailabilityStrategyTestsBase.class.getDeclaredField("testContainerId"); + f.setAccessible(true); + return (String) f.get(this); + } catch (Exception e) { + throw new RuntimeException(e); + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FaultInjectionWithAvailabilityStrategyTestsBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FaultInjectionWithAvailabilityStrategyTestsBase.java index dbb158438a9e..296a1c06b7ac 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FaultInjectionWithAvailabilityStrategyTestsBase.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FaultInjectionWithAvailabilityStrategyTestsBase.java @@ -31,6 +31,7 @@ import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.PartitionKey; import com.azure.cosmos.models.PartitionKeyDefinition; +import com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions; import com.azure.cosmos.models.ThroughputProperties; import com.azure.cosmos.rx.TestSuiteBase; import com.azure.cosmos.test.faultinjection.CosmosFaultInjectionHelper; @@ -3887,6 +3888,315 @@ public Object[][] testConfigs_readManyAfterCreation() { return addBooleanFlagsToAllTestConfigs(testConfigs_readManyAfterCreation); } + private CosmosResponseWrapper readManyByPartitionKeysCore( + ItemOperationInvocationParameters params, + int numberOfOtherDocumentsWithSamePk + ) { + + List pkValues = new ArrayList<>(); + pkValues.add(new PartitionKey(params.idAndPkValuePair.getRight())); + + CosmosReadManyByPartitionKeysRequestOptions options = new CosmosReadManyByPartitionKeysRequestOptions(); + + CosmosEndToEndOperationLatencyPolicyConfig e2ePolicy = ImplementationBridgeHelpers + .CosmosItemRequestOptionsHelper + .getCosmosItemRequestOptionsAccessor() + .getEndToEndOperationLatencyPolicyConfig(params.options); + options.setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); + + List> returnedPages; + // Let CosmosException propagate to execute()'s catch block — it handles + // status code + sub-status code validation and diagnostics context extraction + // correctly for error cases (the response-level validation path passes null + // for sub-status which breaks validators like validateStatusCodeIsServiceUnavailable). + returnedPages = params.container + .readManyByPartitionKeys(pkValues, options, ObjectNode.class) + .byPage() + .collectList() + .block(); + + ArrayList foundCtxs = new ArrayList<>(); + + if (returnedPages == null || returnedPages.isEmpty()) { + return new CosmosResponseWrapper( + null, + HttpConstants.StatusCodes.OK, + HttpConstants.SubStatusCodes.UNKNOWN, + 0L); + } + + long totalRecordCount = 0L; + for (FeedResponse page : returnedPages) { + if (page.getCosmosDiagnostics() != null) { + foundCtxs.add(page.getCosmosDiagnostics().getDiagnosticsContext()); + } else { + foundCtxs.add(null); + } + + if (page.getResults() != null && page.getResults().size() > 0) { + totalRecordCount += page.getResults().size(); + } + } + + return new CosmosResponseWrapper( + foundCtxs.toArray(new CosmosDiagnosticsContext[0]), + HttpConstants.StatusCodes.OK, + HttpConstants.SubStatusCodes.UNKNOWN, + totalRecordCount); + } + + @DataProvider(name = "testConfigs_readManyByPartitionKeysAfterCreation") + public Object[][] testConfigs_readManyByPartitionKeysAfterCreation() { + + final int ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE = 10; + final int NO_OTHER_DOCS_WITH_SAME_PK = 0; + final int NO_OTHER_DOCS_WITH_SAME_ID = 0; + final int ENOUGH_DOCS_OTHER_PK_TO_HIT_EVERY_PARTITION = PHYSICAL_PARTITION_COUNT * 10; + final int SINGLE_REGION = 1; + final int TWO_REGIONS = 2; + + BiConsumer validateExpectedRecordCount = (response, expectedRecordCount) -> { + if (expectedRecordCount != null) { + assertThat(response).isNotNull(); + assertThat(response.getTotalRecordCount()).isNotNull(); + assertThat(response.getTotalRecordCount()).isEqualTo(expectedRecordCount); + } + }; + + Consumer validateExactlyOneRecordReturned = + (response) -> validateExpectedRecordCount.accept(response, 1L); + + Consumer validateAllRecordsSamePartitionReturned = + (response) -> validateExpectedRecordCount.accept( + response, + 1L + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE); + + BiConsumer validateCtxRegions = + (ctx, expectedNumberOfRegionsContacted) -> { + assertThat(ctx).isNotNull(); + if (ctx != null) { + assertThat(ctx.getContactedRegionNames().size()).isEqualTo(expectedNumberOfRegionsContacted); + } + }; + + Consumer validateCtxSingleRegion = + (ctx) -> validateCtxRegions.accept(ctx, SINGLE_REGION); + + Consumer validateCtxTwoRegions = + (ctx) -> validateCtxRegions.accept(ctx, TWO_REGIONS); + + Consumer validateCtxOnlyFeedResponses = + (ctx) -> { + assertThat(ctx).isNotNull(); + if (ctx != null) { + assertThat(ctx.getDiagnostics()).isNotNull(); + assertThat(ctx.getDiagnostics().size()).isGreaterThanOrEqualTo(1); + } + }; + + Function + readManyByPkSinglePartition = (inputParams) -> + readManyByPartitionKeysCore(inputParams, ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE); + + Function + readManyByPkSingleDoc = (inputParams) -> + readManyByPartitionKeysCore(inputParams, NO_OTHER_DOCS_WITH_SAME_PK); + + Object[][] testConfigs = new Object[][] { + // CONFIG description + // new Object[] { + // TestId - name identifying the test case + // End-to-end timeout + // Availability Strategy used + // Region switch hint + // ConnectionMode + // readManyByPartitionKeys operation callback + // Failure injection callback + // Status code/sub status code validation callback + // Expected number of DiagnosticsContext instances + // Diagnostics context validation callback applied to the first DiagnosticsContext + // Diagnostics context validation callback applied to all other DiagnosticsContext + // Consumer - callback to validate the response + // numberOfOtherDocumentsWithSameId + // numberOfOtherDocumentsWithSamePk + // }, + + // readManyByPartitionKeys - single partition, no failures, no availability strategy + new Object[] { + "ReadManyByPk_SinglePartition_AllGood_NoAvailabilityStrategy", + ONE_SECOND_DURATION, + noAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + noFailureInjection, + validateStatusCodeIs200Ok, + 1, + ArrayUtils.toArray( + validateCtxSingleRegion, + validateCtxOnlyFeedResponses + ), + null, + validateAllRecordsSamePartitionReturned, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + + // readManyByPartitionKeys - single doc, no failures, no availability strategy + new Object[] { + "ReadManyByPk_SingleDoc_AllGood_NoAvailabilityStrategy", + ONE_SECOND_DURATION, + noAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSingleDoc, + noFailureInjection, + validateStatusCodeIs200Ok, + 1, + ArrayUtils.toArray( + validateCtxSingleRegion, + validateCtxOnlyFeedResponses + ), + null, + validateExactlyOneRecordReturned, + NO_OTHER_DOCS_WITH_SAME_ID, + NO_OTHER_DOCS_WITH_SAME_PK + }, + + // readManyByPartitionKeys - 408 timeout in first region, eager availability strategy + // Should succeed via hedging to second region + new Object[] { + "ReadManyByPk_SinglePartition_408_FirstRegionOnly_EagerAvailabilityStrategy", + THREE_SECOND_DURATION, + eagerThresholdAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + injectTransitTimeoutIntoFirstRegionOnly, + validateStatusCodeIs200Ok, + 1, + ArrayUtils.toArray( + validateCtxTwoRegions, + validateCtxOnlyFeedResponses + ), + null, + validateAllRecordsSamePartitionReturned, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + + // readManyByPartitionKeys - 404/1002 in first region, remote preferred, reluctant strategy + // Client retry policy should failover to second region before hedging kicks in + new Object[] { + "ReadManyByPk_SinglePartition_404-1002_RemotePreferred_FirstRegionOnly_ReluctantAvailabilityStrategy", + Duration.ofSeconds(10), + reluctantThresholdAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + injectReadSessionNotAvailableIntoFirstRegionOnly, + validateStatusCodeIs200Ok, + 1, + ArrayUtils.toArray( + validateCtxTwoRegions, + validateCtxOnlyFeedResponses + ), + null, + validateAllRecordsSamePartitionReturned, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + + // readManyByPartitionKeys - 429/3200 in first region, default availability strategy + // Should succeed via hedging to second region + new Object[] { + "ReadManyByPk_SinglePartition_429-3200_FirstRegionOnly_DefaultAvailabilityStrategy", + THREE_SECOND_DURATION, + defaultAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + injectRequestRateTooLargeIntoFirstRegionOnly, + validateStatusCodeIs200Ok, + 1, + ArrayUtils.toArray( + validateCtxTwoRegions, + validateCtxOnlyFeedResponses + ), + null, + validateAllRecordsSamePartitionReturned, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + + // readManyByPartitionKeys - 429/3200 in first region, no availability strategy + // Should time out since no hedging and 429 retries locally until timeout + new Object[] { + "ReadManyByPk_SinglePartition_429-3200_FirstRegionOnly_NoAvailabilityStrategy", + THREE_SECOND_DURATION, + noAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + injectRequestRateTooLargeIntoFirstRegionOnly, + validateStatusCodeIsOperationCancelled, + 1, + ArrayUtils.toArray( + validateCtxSingleRegion, + validateCtxOnlyFeedResponses + ), + null, + null, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + + // readManyByPartitionKeys - 503 in first region, no availability strategy + // ClientRetryPolicy will failover to second region + new Object[] { + "ReadManyByPk_SinglePartition_503_FirstRegionOnly_NoAvailabilityStrategy", + Duration.ofSeconds(90), + noAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + injectServiceUnavailableIntoFirstRegionOnly, + validateStatusCodeIs200Ok, + 1, + ArrayUtils.toArray( + validateCtxTwoRegions, + validateCtxOnlyFeedResponses + ), + null, + validateAllRecordsSamePartitionReturned, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + + // readManyByPartitionKeys - 503 in all regions, eager availability strategy + // Both regions fail - should get 503 + new Object[] { + "ReadManyByPk_SinglePartition_503_AllRegions_EagerAvailabilityStrategy", + Duration.ofSeconds(10), + eagerThresholdAvailabilityStrategy, + noRegionSwitchHint, + ConnectionMode.DIRECT, + readManyByPkSinglePartition, + injectServiceUnavailableIntoAllRegions, + validateStatusCodeIsServiceUnavailable, + 1, + ArrayUtils.toArray( + validateCtxTwoRegions + ), + null, + null, + NO_OTHER_DOCS_WITH_SAME_ID, + ENOUGH_DOCS_SAME_PK_TO_EXCEED_PAGE_SIZE + }, + }; + + return addBooleanFlagsToAllTestConfigs(testConfigs); + } private CosmosResponseWrapper readAllReturnsTotalRecordCountCore( diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java new file mode 100644 index 000000000000..4b287457ee4b --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java @@ -0,0 +1,892 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + */ + +package com.azure.cosmos; + +import com.azure.cosmos.models.CosmosContainerProperties; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.PartitionKeyBuilder; +import com.azure.cosmos.models.PartitionKeyDefinition; +import com.azure.cosmos.models.PartitionKeyDefinitionVersion; +import com.azure.cosmos.models.PartitionKind; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; +import com.azure.cosmos.rx.TestSuiteBase; +import com.azure.cosmos.util.CosmosPagedIterable; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.fail; + +public class ReadManyByPartitionKeyTest extends TestSuiteBase { + + private String preExistingDatabaseId = CosmosDatabaseForTest.generateId(); + private CosmosClient client; + private CosmosDatabase createdDatabase; + + // Single PK container (/mypk) + private CosmosContainer singlePkContainer; + + // HPK container (/city, /zipcode, /areaCode) + private CosmosContainer multiHashContainer; + + @Factory(dataProvider = "clientBuilders") + public ReadManyByPartitionKeyTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) + public void before_ReadManyByPartitionKeyTest() { + client = getClientBuilder().buildClient(); + createdDatabase = createSyncDatabase(client, preExistingDatabaseId); + + // Single PK container + String singlePkContainerName = UUID.randomUUID().toString(); + CosmosContainerProperties singlePkProps = new CosmosContainerProperties(singlePkContainerName, "/mypk"); + createdDatabase.createContainer(singlePkProps); + singlePkContainer = createdDatabase.getContainer(singlePkContainerName); + + // HPK container + String multiHashContainerName = UUID.randomUUID().toString(); + PartitionKeyDefinition hpkDef = new PartitionKeyDefinition(); + hpkDef.setKind(PartitionKind.MULTI_HASH); + hpkDef.setVersion(PartitionKeyDefinitionVersion.V2); + ArrayList paths = new ArrayList<>(); + paths.add("/city"); + paths.add("/zipcode"); + paths.add("/areaCode"); + hpkDef.setPaths(paths); + + CosmosContainerProperties hpkProps = new CosmosContainerProperties(multiHashContainerName, hpkDef); + createdDatabase.createContainer(hpkProps); + multiHashContainer = createdDatabase.getContainer(multiHashContainerName); + } + + @AfterClass(groups = {"emulator"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + safeDeleteSyncDatabase(createdDatabase); + safeCloseSyncClient(client); + } + + //region Single PK tests + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_basic() { + // Create items with different PKs + List items = createSinglePkItems("pk1", 3); + items.addAll(createSinglePkItems("pk2", 2)); + items.addAll(createSinglePkItems("pk3", 4)); + + // Read by 2 partition keys + List pkValues = Arrays.asList( + new PartitionKey("pk1"), + new PartitionKey("pk2")); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys(pkValues, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).hasSize(5); // 3 + 2 + resultList.forEach(item -> { + String pk = item.get("mypk").asText(); + assertThat(pk).isIn("pk1", "pk2"); + }); + + // Cleanup + cleanupContainer(singlePkContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_withProjection() { + List items = createSinglePkItems("pkProj", 2); + + List pkValues = Collections.singletonList(new PartitionKey("pkProj")); + SqlQuerySpec customQuery = new SqlQuerySpec("SELECT c.id, c.mypk FROM c"); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys( + pkValues, customQuery, null, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).hasSize(2); + // Should only have id and mypk fields (plus system properties) + resultList.forEach(item -> { + assertThat(item.has("id")).isTrue(); + assertThat(item.has("mypk")).isTrue(); + }); + + cleanupContainer(singlePkContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_withAdditionalFilter() { + // Create items with different "status" values + createSinglePkItemsWithStatus("pkFilter", "active", 3); + createSinglePkItemsWithStatus("pkFilter", "inactive", 2); + + List pkValues = Collections.singletonList(new PartitionKey("pkFilter")); + SqlQuerySpec customQuery = new SqlQuerySpec( + "SELECT * FROM c WHERE c.status = @status", + Arrays.asList(new SqlParameter("@status", "active"))); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys( + pkValues, customQuery, null, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).hasSize(3); + resultList.forEach(item -> { + assertThat(item.get("status").asText()).isEqualTo("active"); + }); + + cleanupContainer(singlePkContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_emptyResults() { + List pkValues = Collections.singletonList(new PartitionKey("nonExistent")); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys(pkValues, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).isEmpty(); + } + + //endregion + + //region HPK tests + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void hpk_readManyByPartitionKey_fullPk() { + createHpkItems(); + + // Read by full PKs + List pkValues = Arrays.asList( + new PartitionKeyBuilder().add("Redmond").add("98053").add(1).build(), + new PartitionKeyBuilder().add("Pittsburgh").add("15232").add(2).build()); + + CosmosPagedIterable results = multiHashContainer.readManyByPartitionKeys(pkValues, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + // Redmond/98053/1 has 2 items, Pittsburgh/15232/2 has 1 item + assertThat(resultList).hasSize(3); + + cleanupContainer(multiHashContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void hpk_readManyByPartitionKey_partialPk_singleLevel() { + createHpkItems(); + + // Read by partial PK (only city) + List pkValues = Collections.singletonList( + new PartitionKeyBuilder().add("Redmond").build()); + + CosmosPagedIterable results = multiHashContainer.readManyByPartitionKeys(pkValues, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + // Redmond has 3 items total (2 with 98053/1 and 1 with 12345/1) + assertThat(resultList).hasSize(3); + resultList.forEach(item -> { + assertThat(item.get("city").asText()).isEqualTo("Redmond"); + }); + + cleanupContainer(multiHashContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void hpk_readManyByPartitionKey_partialPk_twoLevels() { + createHpkItems(); + + // Read by partial PK (city + zipcode) + List pkValues = Collections.singletonList( + new PartitionKeyBuilder().add("Redmond").add("98053").build()); + + CosmosPagedIterable results = multiHashContainer.readManyByPartitionKeys(pkValues, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + // Redmond/98053 has 2 items + assertThat(resultList).hasSize(2); + resultList.forEach(item -> { + assertThat(item.get("city").asText()).isEqualTo("Redmond"); + assertThat(item.get("zipcode").asText()).isEqualTo("98053"); + }); + + cleanupContainer(multiHashContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + @SuppressWarnings("deprecation") + public void hpk_readManyByPartitionKey_withNoneComponent() { + try { + createHpkItems(); + + ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); + item.put("id", UUID.randomUUID().toString()); + item.put("city", "Redmond"); + item.put("zipcode", "98053"); + + try { + multiHashContainer.createItem(item); + fail("Should have thrown CosmosException for HPK item with missing trailing partition key component"); + } catch (CosmosException e) { + assertThat(e.getMessage()).contains("wrong-pk-value"); + } + + try { + new PartitionKeyBuilder().add("Redmond").add("98053").addNoneValue().build(); + fail("Should have thrown IllegalStateException for HPK addNoneValue"); + } catch (IllegalStateException e) { + assertThat(e.getMessage()).contains("PartitionKey.None can't be used with multiple paths"); + } + } finally { + cleanupContainer(multiHashContainer); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void hpk_readManyByPartitionKey_withProjection() { + createHpkItems(); + + List pkValues = Collections.singletonList( + new PartitionKeyBuilder().add("Redmond").add("98053").add(1).build()); + + SqlQuerySpec customQuery = new SqlQuerySpec("SELECT c.id, c.city FROM c"); + + CosmosPagedIterable results = multiHashContainer.readManyByPartitionKeys( + pkValues, customQuery, null, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).hasSize(2); + + cleanupContainer(multiHashContainer); + } + + //endregion + + //region Negative/validation tests + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsAggregateQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec aggregateQuery = new SqlQuerySpec("SELECT COUNT(1) FROM c"); + + try { + singlePkContainer.readManyByPartitionKeys(pkValues, aggregateQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for aggregate query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("aggregates"); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsOrderByQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec orderByQuery = new SqlQuerySpec("SELECT * FROM c ORDER BY c.id"); + + try { + singlePkContainer.readManyByPartitionKeys(pkValues, orderByQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for ORDER BY query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("ORDER BY"); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsDistinctQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec distinctQuery = new SqlQuerySpec("SELECT DISTINCT c.mypk FROM c"); + + try { + singlePkContainer.readManyByPartitionKeys(pkValues, distinctQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for DISTINCT query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("DISTINCT"); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsGroupByQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec groupByQuery = new SqlQuerySpec("SELECT c.mypk FROM c GROUP BY c.mypk"); + + try { + singlePkContainer.readManyByPartitionKeys(pkValues, groupByQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for GROUP BY query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("GROUP BY"); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsGroupByWithAggregateQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec groupByWithAggregateQuery = new SqlQuerySpec("SELECT c.mypk, COUNT(1) as cnt FROM c GROUP BY c.mypk"); + + try { + singlePkContainer.readManyByPartitionKeys(pkValues, groupByWithAggregateQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for GROUP BY with aggregate query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("GROUP BY"); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT, expectedExceptions = NullPointerException.class) + public void rejectsNullPartitionKeyList() { + singlePkContainer.readManyByPartitionKeys((List) null, ObjectNode.class); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT, expectedExceptions = IllegalArgumentException.class) + public void rejectsEmptyPartitionKeyList() { + singlePkContainer.readManyByPartitionKeys(new ArrayList<>(), ObjectNode.class) + .stream().collect(Collectors.toList()); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsOffsetQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec offsetQuery = new SqlQuerySpec("SELECT * FROM c OFFSET 0 LIMIT 10"); + + try { + singlePkContainer.readManyByPartitionKeys(pkValues, offsetQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for OFFSET query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("OFFSET"); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void rejectsTopQuery() { + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + SqlQuerySpec topQuery = new SqlQuerySpec("SELECT TOP 5 * FROM c"); + + try { + singlePkContainer.readManyByPartitionKeys(pkValues, topQuery, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + fail("Should have thrown IllegalArgumentException for TOP query"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()).contains("TOP"); + } + } + + // DCOUNT, standalone LIMIT, and hybrid/vector/full-text search cannot be tested against the + // emulator: DCOUNT is not recognized as a built-in function, standalone LIMIT is not valid + // Cosmos SQL syntax (only valid with OFFSET, already covered by rejectsOffsetQuery), and + // hybrid search requires vector indexes. All three are covered by unit tests in + // ReadManyByPartitionKeyQueryPlanValidationTest (rejectsDCountQueryPlan, rejectsLimitQueryPlan, + // rejectsHybridSearchQueryPlanWithoutDereferencingNullQueryInfo). + + + + //endregion + + + //region Batch size tests (#10) + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_withSmallBatchSize() { + // Temporarily set batch size to 2 to exercise the batching/interleaving logic + String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + try { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "2"); + + // Create items across 4 PKs (more than the batch size of 2) + List items = createSinglePkItems("batchPk1", 2); + items.addAll(createSinglePkItems("batchPk2", 2)); + items.addAll(createSinglePkItems("batchPk3", 2)); + items.addAll(createSinglePkItems("batchPk4", 2)); + + // Read all 4 PKs — should be split into batches of 2 + List pkValues = Arrays.asList( + new PartitionKey("batchPk1"), + new PartitionKey("batchPk2"), + new PartitionKey("batchPk3"), + new PartitionKey("batchPk4")); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys(pkValues, ObjectNode.class); + List> pages = new ArrayList<>(); + results.iterableByPage().forEach(pages::add); + List resultList = pages.stream() + .flatMap(page -> page.getResults().stream()) + .collect(Collectors.toList()); + + assertThat(resultList).hasSize(8); // 2 items per PK * 4 PKs + assertThat(pages.size()).isGreaterThan(1); + resultList.forEach(item -> { + String pk = item.get("mypk").asText(); + assertThat(pk).isIn("batchPk1", "batchPk2", "batchPk3", "batchPk4"); + }); + + cleanupContainer(singlePkContainer); + } finally { + if (originalValue != null) { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); + } else { + System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + } + } + } + + //endregion + + //region Custom serializer regression tests (#5) + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_withRequestOptions() { + List items = createSinglePkItems("pkOpts", 3); + + List pkValues = Collections.singletonList(new PartitionKey("pkOpts")); + com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions options = new com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions(); + AtomicInteger deserializeCount = new AtomicInteger(); + options.setCustomItemSerializer(new CosmosItemSerializerNoExceptionWrapping() { + @Override + public Map serialize(T item) { + return CosmosItemSerializer.DEFAULT_SERIALIZER.serialize(item); + } + + @Override + public T deserialize(Map jsonNodeMap, Class classType) { + deserializeCount.incrementAndGet(); + return CosmosItemSerializer.DEFAULT_SERIALIZER.deserialize(jsonNodeMap, classType); + } + }); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys( + pkValues, options, ReadManyByPartitionKeyPojo.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).hasSize(3); + assertThat(deserializeCount.get()).isEqualTo(3); + assertThat(resultList.stream().map(item -> item.mypk)).containsOnly("pkOpts"); + + cleanupContainer(singlePkContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_readManyByPartitionKey_withRequestOptionsAndMaxConcurrentBatchPrefetch() { + // Regression test: passing non-null CosmosReadManyByPartitionKeysRequestOptions + // with maxConcurrentBatchPrefetch set should not throw NullPointerException + // from auto-unboxing a null MaxDegreeOfParallelism during options cloning. + List items = createSinglePkItems("pkMdop", 3); + + List pkValues = Collections.singletonList(new PartitionKey("pkMdop")); + com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions options = + new com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions(); + options.setMaxConcurrentBatchPrefetch(2); + + CosmosPagedIterable results = singlePkContainer.readManyByPartitionKeys( + pkValues, options, ObjectNode.class); + List resultList = results.stream().collect(Collectors.toList()); + + assertThat(resultList).hasSize(3); + resultList.forEach(item -> { + assertThat(item.get("mypk").asText()).isEqualTo("pkMdop"); + }); + + cleanupContainer(singlePkContainer); + } + //endregion + + //region Continuation token tests + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_sequentialExecution_emitsContinuationTokens() { + // Use small batch size so we get multiple batches (and thus multiple continuation tokens) + String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + try { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "2"); + + // Create items across 3 PKs + createSinglePkItems("seqPk1", 3); + createSinglePkItems("seqPk2", 3); + createSinglePkItems("seqPk3", 3); + + List pkValues = Arrays.asList( + new PartitionKey("seqPk1"), + new PartitionKey("seqPk2"), + new PartitionKey("seqPk3")); + + // Use the async container to collect FeedResponse pages with continuation tokens + CosmosAsyncContainer asyncContainer = client.asyncClient() + .getDatabase(preExistingDatabaseId) + .getContainer(singlePkContainer.getId()); + + List> pages = asyncContainer + .readManyByPartitionKeys(pkValues, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(pages).isNotNull(); + assertThat(pages).isNotEmpty(); + + // All non-final pages should have a non-null continuation token + for (int i = 0; i < pages.size() - 1; i++) { + assertThat(pages.get(i).getContinuationToken()) + .as("Page %d should have a continuation token", i) + .isNotNull() + .isNotEmpty(); + } + + // The final page should have null continuation + assertThat(pages.get(pages.size() - 1).getContinuationToken()) + .as("Last page should have null continuation token") + .isNull(); + + // Total items should be 9 + long totalItems = pages.stream() + .mapToLong(p -> p.getResults().size()) + .sum(); + assertThat(totalItems).isEqualTo(9); + + cleanupContainer(singlePkContainer); + } finally { + if (originalValue != null) { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); + } else { + System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + } + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_continuationToken_resumesCorrectly() { + // Use small batch size to force multiple batches + String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + try { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "1"); + + // Create items across 3 PKs + createSinglePkItems("resumePk1", 2); + createSinglePkItems("resumePk2", 2); + createSinglePkItems("resumePk3", 2); + + List pkValues = Arrays.asList( + new PartitionKey("resumePk1"), + new PartitionKey("resumePk2"), + new PartitionKey("resumePk3")); + + CosmosAsyncContainer asyncContainer = client.asyncClient() + .getDatabase(preExistingDatabaseId) + .getContainer(singlePkContainer.getId()); + + // First pass: collect all items + List> allPages = asyncContainer + .readManyByPartitionKeys(pkValues, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(allPages).isNotNull(); + assertThat(allPages.size()).isGreaterThan(1); + + // Pick a continuation token from the first page + String continuationAfterFirstPage = allPages.get(0).getContinuationToken(); + assertThat(continuationAfterFirstPage).isNotNull(); + List itemsFromFirstPage = allPages.get(0).getResults(); + + // Second pass: resume from the continuation token + com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions options2 = + new com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions(); + options2.setContinuationToken(continuationAfterFirstPage); + + List> remainingPages = asyncContainer + .readManyByPartitionKeys(pkValues, options2, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(remainingPages).isNotNull(); + + // Collect all item ids + List firstPageIds = itemsFromFirstPage.stream() + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + List remainingIds = remainingPages.stream() + .flatMap(p -> p.getResults().stream()) + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + List allIds = allPages.stream() + .flatMap(p -> p.getResults().stream()) + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + + // Items from first page + remaining should equal all items (no overlap, no gap) + List combined = new ArrayList<>(firstPageIds); + combined.addAll(remainingIds); + + assertThat(combined).hasSameElementsAs(allIds); + + // No duplicates + assertThat(combined).doesNotHaveDuplicates(); + + cleanupContainer(singlePkContainer); + } finally { + if (originalValue != null) { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); + } else { + System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + } + } + } + + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void singlePk_continuationToken_resumesCorrectly_whenInputContainsDuplicatePartitionKeys() { + String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + try { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "1"); + + createSinglePkItems("dupResumePk1", 2); + createSinglePkItems("dupResumePk2", 2); + createSinglePkItems("dupResumePk3", 2); + + List pkValues = Arrays.asList( + new PartitionKey("dupResumePk2"), + new PartitionKey("dupResumePk1"), + new PartitionKey("dupResumePk2"), + new PartitionKey("dupResumePk3"), + new PartitionKey("dupResumePk1")); + + CosmosAsyncContainer asyncContainer = client.asyncClient() + .getDatabase(preExistingDatabaseId) + .getContainer(singlePkContainer.getId()); + + List> allPages = asyncContainer + .readManyByPartitionKeys(pkValues, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(allPages).isNotNull(); + assertThat(allPages.size()).isGreaterThan(1); + + String continuationAfterFirstPage = allPages.get(0).getContinuationToken(); + assertThat(continuationAfterFirstPage).isNotNull(); + List itemsFromFirstPage = allPages.get(0).getResults(); + + com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions options2 = + new com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions(); + options2.setContinuationToken(continuationAfterFirstPage); + + List> remainingPages = asyncContainer + .readManyByPartitionKeys(pkValues, options2, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(remainingPages).isNotNull(); + + List firstPageIds = itemsFromFirstPage.stream() + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + List remainingIds = remainingPages.stream() + .flatMap(p -> p.getResults().stream()) + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + List allIds = allPages.stream() + .flatMap(p -> p.getResults().stream()) + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + + List combined = new ArrayList<>(firstPageIds); + combined.addAll(remainingIds); + + assertThat(combined).hasSameElementsAs(allIds); + assertThat(combined).doesNotHaveDuplicates(); + assertThat(combined).hasSize(6); + + cleanupContainer(singlePkContainer); + } finally { + if (originalValue != null) { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); + } else { + System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + } + } + } + + //endregion + + //region Continuation-token resume correctness tests + + @Test(groups = {"emulator"}, timeOut = TIMEOUT * 3) + public void singlePk_continuationToken_resumeAtEveryPageBoundary_noLossNoDuplicates() { + // Validates that resuming from ANY page boundary produces a complete, duplicate-free + // result set when combined with the items from earlier pages. + String originalValue = System.getProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + try { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", "2"); + + createSinglePkItems("resumePk1", 4); + createSinglePkItems("resumePk2", 4); + createSinglePkItems("resumePk3", 4); + createSinglePkItems("resumePk4", 4); + + List pkValues = Arrays.asList( + new PartitionKey("resumePk1"), + new PartitionKey("resumePk2"), + new PartitionKey("resumePk3"), + new PartitionKey("resumePk4")); + + CosmosAsyncContainer asyncContainer = client.asyncClient() + .getDatabase(preExistingDatabaseId) + .getContainer(singlePkContainer.getId()); + + // Collect all pages in a single pass + List> allPages = asyncContainer + .readManyByPartitionKeys(pkValues, ObjectNode.class) + .byPage() + .collectList() + .block(); + + assertThat(allPages).isNotNull(); + assertThat(allPages.size()).isGreaterThan(1); + + List allIds = allPages.stream() + .flatMap(p -> p.getResults().stream()) + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + assertThat(allIds).hasSize(16); + assertThat(allIds).doesNotHaveDuplicates(); + + // For each page boundary, resume from that boundary's continuation token + // and verify: items before + items after = all items, no duplicates + for (int splitAt = 0; splitAt < allPages.size() - 1; splitAt++) { + String continuation = allPages.get(splitAt).getContinuationToken(); + if (continuation == null) { + continue; // last page has null continuation + } + + // Items from pages 0..splitAt + List beforeIds = new ArrayList<>(); + for (int p = 0; p <= splitAt; p++) { + allPages.get(p).getResults().forEach(n -> beforeIds.add(n.get("id").asText())); + } + + // Resume from continuation + com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions resumeOptions = + new com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions(); + resumeOptions.setContinuationToken(continuation); + + List resumedResults = asyncContainer + .readManyByPartitionKeys(pkValues, resumeOptions, ObjectNode.class) + .byPage() + .flatMapIterable(FeedResponse::getResults) + .collectList() + .block(); + + assertThat(resumedResults).isNotNull(); + List afterIds = resumedResults.stream() + .map(n -> n.get("id").asText()) + .collect(Collectors.toList()); + + List combined = new ArrayList<>(beforeIds); + combined.addAll(afterIds); + + assertThat(combined) + .as("Resume at page boundary %d should produce all items", splitAt) + .doesNotHaveDuplicates(); + assertThat(combined) + .as("Resume at page boundary %d should cover all 16 items", splitAt) + .hasSameElementsAs(allIds); + } + + cleanupContainer(singlePkContainer); + } finally { + if (originalValue != null) { + System.setProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE", originalValue); + } else { + System.clearProperty("COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"); + } + } + } + + //endregion + + //region helper methods + + private List createSinglePkItems(String pkValue, int count) { + List items = new ArrayList<>(); + for (int i = 0; i < count; i++) { + ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); + item.put("id", UUID.randomUUID().toString()); + item.put("mypk", pkValue); + singlePkContainer.createItem(item); + items.add(item); + } + return items; + } + + private List createSinglePkItemsWithStatus(String pkValue, String status, int count) { + List items = new ArrayList<>(); + for (int i = 0; i < count; i++) { + ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); + item.put("id", UUID.randomUUID().toString()); + item.put("mypk", pkValue); + item.put("status", status); + singlePkContainer.createItem(item); + items.add(item); + } + return items; + } + + private void createHpkItems() { + // Same data as CosmosMultiHashTest.createItems() + createHpkItem("Redmond", "98053", 1); + createHpkItem("Redmond", "98053", 1); + createHpkItem("Pittsburgh", "15232", 2); + createHpkItem("Stonybrook", "11790", 3); + createHpkItem("Stonybrook", "11794", 3); + createHpkItem("Stonybrook", "11791", 3); + createHpkItem("Redmond", "12345", 1); + } + + private void createHpkItem(String city, String zipcode, int areaCode) { + ObjectNode item = com.azure.cosmos.implementation.Utils.getSimpleObjectMapper().createObjectNode(); + item.put("id", UUID.randomUUID().toString()); + item.put("city", city); + item.put("zipcode", zipcode); + item.put("areaCode", areaCode); + multiHashContainer.createItem(item); + } + + private void cleanupContainer(CosmosContainer container) { + CosmosPagedIterable allItems = container.queryItems( + "SELECT * FROM c", new com.azure.cosmos.models.CosmosQueryRequestOptions(), ObjectNode.class); + allItems.forEach(item -> { + try { + container.deleteItem(item, new CosmosItemRequestOptions()); + } catch (CosmosException e) { + // ignore cleanup failures + } + }); + } + + private static class ReadManyByPartitionKeyPojo { + public String id; + public String mypk; + } + + //endregion +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java index 2ed7e7e091f4..2cf84c721107 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionServerErrorRuleOnDirectTests.java @@ -24,6 +24,7 @@ import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedRange; +import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.PartitionKey; import com.azure.cosmos.test.faultinjection.CosmosFaultInjectionHelper; import com.azure.cosmos.test.faultinjection.FaultInjectionConditionBuilder; @@ -1078,6 +1079,103 @@ public void faultInjectionServerErrorRuleTests_HitLimit() throws JsonProcessingE } } + @Test(groups = {"multi-region"}, timeOut = TIMEOUT * 3) + public void readManyByPartitionKeys_goneError_retriedInternally_noLossNoDuplicates() { + // Injects GONE (410/0) errors on query operations targeting readManyByPartitionKeys. + // 410/0 triggers a partition-key-range refresh + retry inside the SDK's + // DocumentProducer retry loop (direct mode): + // + // DocumentProducer.executeFeedOperationCore [DocumentProducer.java:144] + // → ObservableHelper.inlineIfPossibleAsObs(func, finalRetryPolicy) [DocumentProducer.java:147] + // → BackoffRetryUtility.executeRetry() [ObservableHelper.java:43] + // → client.executeQueryAsync(request) [IDocumentQueryClient impl in RxDocumentClientImpl.java:5371] + // → RxDocumentClientImpl.query(request) [RxDocumentClientImpl.java:1822] + // → getStoreProxy(request).processMessage(request) + // → [Direct: RntbdTransportClient] ← GONE 410/0 injected by fault injection + // + // The finalRetryPolicy is a ClientRetryPolicy created via: + // resetSessionTokenRetryPolicy.getRequestPolicy() [ParallelDocumentQueryExecutionContextBase.java:115] + // ClientRetryPolicy handles 410 (GONE) by refreshing the partition key range cache + // and retrying in the same region — this works regardless of region count. + + CosmosAsyncContainer container = getSharedMultiPartitionCosmosContainer(clientWithoutPreferredRegions); + + // Create test items across 3 partition keys + String pk1 = "readManyGone_" + UUID.randomUUID(); + String pk2 = "readManyGone_" + UUID.randomUUID(); + String pk3 = "readManyGone_" + UUID.randomUUID(); + + List createdItems = new ArrayList<>(); + for (String pk : Arrays.asList(pk1, pk2, pk3)) { + for (int i = 0; i < 3; i++) { + TestObject item = TestObject.create(pk); + container.createItem(item).block(); + createdItems.add(item); + } + } + + List pkValues = Arrays.asList( + new PartitionKey(pk1), + new PartitionKey(pk2), + new PartitionKey(pk3)); + + // Baseline — no faults + List baselineResults = container + .readManyByPartitionKeys(pkValues, TestObject.class) + .byPage() + .flatMapIterable(FeedResponse::getResults) + .collectList() + .block(); + + assertThat(baselineResults).isNotNull(); + assertThat(baselineResults.size()).isEqualTo(9); + List baselineIds = baselineResults.stream() + .map(TestObject::getId) + .sorted() + .collect(Collectors.toList()); + + // Inject GONE (410/0) on QUERY_ITEM — direct mode, no connectionType needed + FaultInjectionRule goneRule = new FaultInjectionRuleBuilder("readManyByPk-gone-direct") + .condition(new FaultInjectionConditionBuilder() + .operationType(FaultInjectionOperationType.QUERY_ITEM) + .build()) + .result(FaultInjectionResultBuilders + .getResultBuilder(FaultInjectionServerErrorType.GONE) + .build()) + .hitLimit(3) + .build(); + + CosmosFaultInjectionHelper + .configureFaultInjectionRules(container, Arrays.asList(goneRule)) + .block(); + + try { + List faultInjectedResults = container + .readManyByPartitionKeys(pkValues, TestObject.class) + .byPage() + .flatMapIterable(FeedResponse::getResults) + .collectList() + .block(); + + assertThat(faultInjectedResults).isNotNull(); + List faultInjectedIds = faultInjectedResults.stream() + .map(TestObject::getId) + .sorted() + .collect(Collectors.toList()); + + org.assertj.core.api.Assertions.assertThat(faultInjectedIds).doesNotHaveDuplicates(); + org.assertj.core.api.Assertions.assertThat(faultInjectedIds).hasSameElementsAs(baselineIds); + org.assertj.core.api.Assertions.assertThat(faultInjectedIds).hasSize(9); + } finally { + goneRule.disable(); + } + + // Cleanup + for (TestObject item : createdItems) { + container.deleteItem(item.getId(), new PartitionKey(item.getMypk())).block(); + } + } + @AfterClass(groups = {"multi-region", "long", "fast", "fi-multi-master", "multi-region-strong"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) public void afterClass() { safeClose(clientWithoutPreferredRegions); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java new file mode 100644 index 000000000000..86057a353021 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java @@ -0,0 +1,401 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + */ + +package com.azure.cosmos.implementation; + +import com.azure.cosmos.implementation.routing.Range; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.models.ModelBridgeInternal; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; +import org.testng.annotations.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public class ReadManyByPartitionKeyContinuationTokenTest { + + private static final String TEST_COLLECTION_RID = "dbs/testDb/colls/testColl"; + private static final String TEST_QUERY_HASH = "12345"; + private static final String TEST_PARTITION_KEY_SET_HASH = "67890"; + + /** + * Builds a BatchDefinition whose batchFilter is the half-open EPK range [min, max). + */ + private static ReadManyByPartitionKeyContinuationToken.BatchDefinition bd(String min, String max) { + return new ReadManyByPartitionKeyContinuationToken.BatchDefinition( + new Range<>(min, max, true, false)); + } + + private static List bds( + ReadManyByPartitionKeyContinuationToken.BatchDefinition... defs) { + return new ArrayList<>(Arrays.asList(defs)); + } + + @Test(groups = { "unit" }) + public void roundtrip_withBackendContinuation() { + List remaining = bds( + bd("05C1E0", "0BF333"), + bd("0BF333", "FF")); + ReadManyByPartitionKeyContinuationToken.BatchDefinition current = bd("", "05C1E0"); + String backendCont = "eyJDb21wb3NpdGVUb2tlbg=="; + + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + remaining, current, backendCont, TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + + String serialized = token.serialize(); + assertThat(serialized).isNotNull().isNotEmpty(); + + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + assertThat(deserialized.getBackendContinuation()).isEqualTo(backendCont); + assertThat(deserialized.getCollectionRid()).isEqualTo(TEST_COLLECTION_RID); + assertThat(deserialized.getQueryHash()).isEqualTo(TEST_QUERY_HASH); + assertThat(deserialized.getPartitionKeySetHash()).isEqualTo(TEST_PARTITION_KEY_SET_HASH); + + ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBatch = deserialized.getCurrentBatch(); + assertThat(currentBatch.getBatchFilter().getMin()).isEqualTo(""); + assertThat(currentBatch.getBatchFilter().getMax()).isEqualTo("05C1E0"); + + List remainingBatches = + deserialized.getRemainingBatches(); + assertThat(remainingBatches).hasSize(2); + assertThat(remainingBatches.get(0).getBatchFilter().getMin()).isEqualTo("05C1E0"); + assertThat(remainingBatches.get(0).getBatchFilter().getMax()).isEqualTo("0BF333"); + assertThat(remainingBatches.get(1).getBatchFilter().getMin()).isEqualTo("0BF333"); + assertThat(remainingBatches.get(1).getBatchFilter().getMax()).isEqualTo("FF"); + } + + @Test(groups = { "unit" }) + public void roundtrip_withNullBackendContinuation() { + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + bds(bd("0BF333", "FF")), bd("05C1E0", "0BF333"), + null, TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + assertThat(deserialized.getBackendContinuation()).isNull(); + assertThat(deserialized.getCurrentBatch().getBatchFilter().getMin()).isEqualTo("05C1E0"); + assertThat(deserialized.getCurrentBatch().getBatchFilter().getMax()).isEqualTo("0BF333"); + assertThat(deserialized.getRemainingBatches()).hasSize(1); + } + + @Test(groups = { "unit" }) + public void roundtrip_emptyRemainingBatches() { + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + Collections.emptyList(), bd("0BF333", "FF"), + "someCont", TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + assertThat(deserialized.getRemainingBatches()).isEmpty(); + assertThat(deserialized.getCurrentBatch().getBatchFilter().getMin()).isEqualTo("0BF333"); + assertThat(deserialized.getBackendContinuation()).isEqualTo("someCont"); + } + + @Test(groups = { "unit" }) + public void roundtrip_lastBatchNoContinuation() { + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + Collections.emptyList(), bd("0BF333", "FF"), + null, TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + assertThat(deserialized.getRemainingBatches()).isEmpty(); + assertThat(deserialized.getBackendContinuation()).isNull(); + } + + @Test(groups = { "unit" }) + public void setFeedResponseContinuationToken_handlesEmptyHeadersWithoutCopyingNormalCase() { + Map immutableEmptyHeaders = Collections.emptyMap(); + FeedResponse emptyResponse = ModelBridgeInternal.createFeedResponse( + Collections.emptyList(), + immutableEmptyHeaders); + + ModelBridgeInternal.setFeedResponseContinuationToken(null, emptyResponse); + + assertThat(emptyResponse.getContinuationToken()).isNull(); + assertThat(emptyResponse.getResponseHeaders()).isSameAs(immutableEmptyHeaders); + assertThat(emptyResponse.getResponseHeaders()).isEmpty(); + + Map normalHeaders = new HashMap<>(); + normalHeaders.put(HttpConstants.HttpHeaders.ACTIVITY_ID, "test-activity-id"); + FeedResponse normalResponse = ModelBridgeInternal.createFeedResponse( + Collections.emptyList(), + normalHeaders); + + ModelBridgeInternal.setFeedResponseContinuationToken("token", normalResponse); + + assertThat(normalResponse.getContinuationToken()).isEqualTo("token"); + assertThat(normalResponse.getResponseHeaders()).isSameAs(normalHeaders); + } + + @Test(groups = { "unit" }) + public void deserialize_malformedInput_throws() { + // Either the base64 decoder or the JSON parsing layer rejects garbage; both raise + // IllegalArgumentException, which is the contract callers depend on. + assertThatThrownBy(() -> + ReadManyByPartitionKeyContinuationToken.deserialize("not-valid-base64!!!") + ).isInstanceOf(IllegalArgumentException.class); + } + + @Test(groups = { "unit" }) + public void deserialize_emptyString_throws() { + assertThatThrownBy(() -> + ReadManyByPartitionKeyContinuationToken.deserialize("") + ).isInstanceOf(IllegalArgumentException.class); + } + + @Test(groups = { "unit" }) + public void deserialize_null_throws() { + assertThatThrownBy(() -> + ReadManyByPartitionKeyContinuationToken.deserialize(null) + ).isInstanceOf(NullPointerException.class); + } + + @Test(groups = { "unit" }) + public void deserialize_nullRemainingBatches_throws() { + // Hand-craft a token with "rb":null - a malformed or tampered token. + String json = "{\"v\":1,\"rb\":null,\"cb\":{\"bf\":{\"min\":\"\",\"max\":\"FF\"}}," + + "\"bc\":null,\"cr\":\"" + TEST_COLLECTION_RID + "\",\"qh\":\"" + TEST_QUERY_HASH + "\",\"ph\":\"" + TEST_PARTITION_KEY_SET_HASH + "\"}"; + String serialized = java.util.Base64.getEncoder().encodeToString(json.getBytes(java.nio.charset.StandardCharsets.UTF_8)); + + assertThatThrownBy(() -> ReadManyByPartitionKeyContinuationToken.deserialize(serialized)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("remainingBatches"); + } + + @Test(groups = { "unit" }) + public void deserialize_nullCurrentBatch_throws() { + // Hand-craft a token with "cb":null - a malformed or tampered token. + String json = "{\"v\":1,\"rb\":[],\"cb\":null," + + "\"bc\":null,\"cr\":\"" + TEST_COLLECTION_RID + "\",\"qh\":\"" + TEST_QUERY_HASH + "\",\"ph\":\"" + TEST_PARTITION_KEY_SET_HASH + "\"}"; + String serialized = java.util.Base64.getEncoder().encodeToString(json.getBytes(java.nio.charset.StandardCharsets.UTF_8)); + + assertThatThrownBy(() -> ReadManyByPartitionKeyContinuationToken.deserialize(serialized)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("currentBatch"); + } + + @Test(groups = { "unit" }) + public void deserialize_unsupportedVersion_throws() { + // Hand-craft a token JSON with version=999 (a future format) and ensure it is rejected. + String json = "{\"v\":999,\"rb\":[],\"cb\":{\"bf\":{\"min\":\"\",\"max\":\"FF\"}}," + + "\"bc\":null,\"cr\":\"" + TEST_COLLECTION_RID + "\",\"qh\":\"" + TEST_QUERY_HASH + "\",\"ph\":\"" + TEST_PARTITION_KEY_SET_HASH + "\"}"; + String serialized = java.util.Base64.getEncoder().encodeToString(json.getBytes(java.nio.charset.StandardCharsets.UTF_8)); + + // The root cause carries the precise "Unsupported version" message; the outer wrapper + // gives a generic "Failed to deserialize" hint. Match against the full chain. + assertThatThrownBy(() -> ReadManyByPartitionKeyContinuationToken.deserialize(serialized)) + .isInstanceOf(IllegalArgumentException.class) + .hasStackTraceContaining("Unsupported readManyByPartitionKeys continuation token version"); + } + + @Test(groups = { "unit" }) + public void serialized_includesVersionField_andIsBase64() { + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + Collections.emptyList(), bd("", "FF"), null, + TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + + String serialized = token.serialize(); + assertThat(serialized).matches("[A-Za-z0-9+/=]+"); + + String json = new String( + java.util.Base64.getDecoder().decode(serialized), + java.nio.charset.StandardCharsets.UTF_8); + assertThat(json).startsWith("{"); + assertThat(json).endsWith("}"); + // The wire format must include the version field so future SDKs can detect/reject + // tokens from incompatible versions. + assertThat(json).contains("\"v\":1"); + } + + @Test(groups = { "unit" }) + public void serialized_doesNotIncludePartitionScope() { + // The partition routing scope is intentionally NOT persisted in the continuation token. + // It is rederived at execution time from the live PartitionKeyRange cache so partition + // splits never cause stale routing information to be embedded in a token. + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + bds(bd("AA", "BB"), bd("BB", "CC")), + bd("", "AA"), null, + TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + + String json = new String( + java.util.Base64.getDecoder().decode(token.serialize()), + java.nio.charset.StandardCharsets.UTF_8); + + assertThat(json).doesNotContain("\"ps\""); + assertThat(json).contains("\"bf\""); + } + + @Test(groups = { "unit" }) + public void version_roundtripsAsCurrent() { + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + Collections.emptyList(), bd("", "FF"), null, + TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(token.serialize()); + + assertThat(deserialized.getVersion()).isEqualTo(ReadManyByPartitionKeyContinuationToken.CURRENT_VERSION); + assertThat(ReadManyByPartitionKeyContinuationToken.CURRENT_VERSION).isEqualTo(1); + } + + @Test(groups = { "unit" }) + public void rangesPreserveMinMaxInclusive() { + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + bds(bd("CD", "EF")), bd("AB", "CD"), null, + TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + Range currentFilter = deserialized.getCurrentBatch().getBatchFilter(); + assertThat(currentFilter.isMinInclusive()).isTrue(); + assertThat(currentFilter.isMaxInclusive()).isFalse(); + Range remainingFilter = deserialized.getRemainingBatches().get(0).getBatchFilter(); + assertThat(remainingFilter.isMinInclusive()).isTrue(); + assertThat(remainingFilter.isMaxInclusive()).isFalse(); + } + + @Test(groups = { "unit" }) + public void collectionRidAndQueryHash_roundtrip() { + String rid = "dbs/myDb/colls/myColl"; + String hash = "98765"; + + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + Collections.emptyList(), bd("", "FF"), null, rid, hash, TEST_PARTITION_KEY_SET_HASH); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + assertThat(deserialized.getCollectionRid()).isEqualTo(rid); + assertThat(deserialized.getQueryHash()).isEqualTo(hash); + } + + @Test(groups = { "unit" }) + public void partitionKeySetHash_roundtrip() { + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + Collections.emptyList(), bd("", "FF"), null, + TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + assertThat(deserialized.getPartitionKeySetHash()).isEqualTo(TEST_PARTITION_KEY_SET_HASH); + } + + @Test(groups = { "unit" }) + public void computePartitionKeySetHash_isStableAcrossDuplicateAndReorderedEpks() { + List epks1 = Arrays.asList("BB", "AA", "BB", "CC"); + List epks2 = Arrays.asList("CC", "AA", "BB"); + + assertThat(ReadManyByPartitionKeyContinuationToken.computePartitionKeySetHash(epks1)) + .isEqualTo(ReadManyByPartitionKeyContinuationToken.computePartitionKeySetHash(epks2)); + } + + @Test(groups = { "unit" }) + public void computePartitionKeySetHash_returnsStableHexDigest() { + String hash = ReadManyByPartitionKeyContinuationToken.computePartitionKeySetHash( + Arrays.asList("BB", "AA", "BB", "CC")); + + assertThat(hash).matches("[0-9a-f]{32}"); + } + + @Test(groups = { "unit" }) + public void computeQueryHash_nullSpec_returnsZero() { + assertThat(ReadManyByPartitionKeyContinuationToken.computeQueryHash(null)).isEqualTo("0"); + } + + @Test(groups = { "unit" }) + public void computeQueryHash_sameQueryText_sameHash() { + SqlQuerySpec spec1 = new SqlQuerySpec("SELECT * FROM c WHERE c.pk = @pk", + Collections.singletonList(new SqlParameter("@pk", "value1"))); + SqlQuerySpec spec2 = new SqlQuerySpec("SELECT * FROM c WHERE c.pk = @pk", + Collections.singletonList(new SqlParameter("@pk", "value1"))); + + assertThat(ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec1)) + .isEqualTo(ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec2)); + } + + @Test(groups = { "unit" }) + public void computeQueryHash_differentParams_differentHash() { + SqlQuerySpec spec1 = new SqlQuerySpec("SELECT * FROM c WHERE c.pk = @pk", + Collections.singletonList(new SqlParameter("@pk", "value1"))); + SqlQuerySpec spec2 = new SqlQuerySpec("SELECT * FROM c WHERE c.pk = @pk", + Collections.singletonList(new SqlParameter("@pk", "value2"))); + + assertThat(ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec1)) + .isNotEqualTo(ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec2)); + } + + @Test(groups = { "unit" }) + public void computeQueryHash_differentQueryText_differentHash() { + SqlQuerySpec spec1 = new SqlQuerySpec("SELECT * FROM c"); + SqlQuerySpec spec2 = new SqlQuerySpec("SELECT c.id FROM c"); + + assertThat(ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec1)) + .isNotEqualTo(ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec2)); + } + + @Test(groups = { "unit" }) + public void computeQueryHash_noParams_stableHash() { + SqlQuerySpec spec = new SqlQuerySpec("SELECT * FROM c"); + String hash1 = ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec); + String hash2 = ReadManyByPartitionKeyContinuationToken.computeQueryHash(spec); + assertThat(hash1).isEqualTo(hash2); + } + + @Test(groups = { "unit" }) + public void batchDefinition_roundtrip() { + ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBd = bd("01", "03"); + List remaining = + Collections.singletonList(bd("05C1E0", "0A")); + + ReadManyByPartitionKeyContinuationToken token = + new ReadManyByPartitionKeyContinuationToken( + remaining, currentBd, "cont", TEST_COLLECTION_RID, TEST_QUERY_HASH, TEST_PARTITION_KEY_SET_HASH); + + String serialized = token.serialize(); + ReadManyByPartitionKeyContinuationToken deserialized = + ReadManyByPartitionKeyContinuationToken.deserialize(serialized); + + ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBatch = deserialized.getCurrentBatch(); + assertThat(currentBatch.getBatchFilter().getMin()).isEqualTo("01"); + assertThat(currentBatch.getBatchFilter().getMax()).isEqualTo("03"); + + ReadManyByPartitionKeyContinuationToken.BatchDefinition remainingBatch = + deserialized.getRemainingBatches().get(0); + assertThat(remainingBatch.getBatchFilter().getMin()).isEqualTo("05C1E0"); + assertThat(remainingBatch.getBatchFilter().getMax()).isEqualTo("0A"); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java new file mode 100644 index 000000000000..d9cab76b16ac --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java @@ -0,0 +1,873 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + */ + +package com.azure.cosmos.implementation; + +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.PartitionKeyBuilder; +import com.azure.cosmos.models.PartitionKeyDefinition; +import com.azure.cosmos.models.PartitionKeyDefinitionVersion; +import com.azure.cosmos.models.PartitionKind; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; +import com.azure.cosmos.implementation.routing.Range; +import org.testng.annotations.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public class ReadManyByPartitionKeyQueryHelperTest { + + //region Single PK (HASH) tests + + @Test(groups = { "unit" }) + public void singlePk_defaultQuery_singleValue() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("SELECT * FROM c\n WHERE"); + assertThat(result.getQueryText()).contains("IN ("); + assertThat(result.getQueryText()).contains("@__rmPk_0"); + assertThat(result.getParameters()).hasSize(1); + assertThat(result.getParameters().get(0).getValue(Object.class)).isEqualTo("pk1"); + } + + @Test(groups = { "unit" }) + public void singlePk_defaultQuery_multipleValues() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Arrays.asList( + new PartitionKey("pk1"), + new PartitionKey("pk2"), + new PartitionKey("pk3")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("IN ("); + assertThat(result.getQueryText()).contains("@__rmPk_0"); + assertThat(result.getQueryText()).contains("@__rmPk_1"); + assertThat(result.getQueryText()).contains("@__rmPk_2"); + assertThat(result.getParameters()).hasSize(3); + } + + @Test(groups = { "unit" }) + public void singlePk_customQuery_noWhere() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Arrays.asList(new PartitionKey("pk1"), new PartitionKey("pk2")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT c.name, c.age FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).startsWith("SELECT c.name, c.age FROM c\n WHERE"); + assertThat(result.getQueryText()).contains("IN ("); + } + + @Test(groups = { "unit" }) + public void singlePk_customQuery_withExistingWhere() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@minAge", 18)); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.age > @minAge", baseParams, pkValues, selectors, pkDef); + + // Should AND the PK filter to the existing WHERE clause + assertThat(result.getQueryText()).contains("WHERE (c.age > @minAge\n) AND ("); + assertThat(result.getQueryText()).contains("IN ("); + assertThat(result.getParameters()).hasSize(2); // @minAge + @__rmPk_0 + assertThat(result.getParameters().get(0).getName()).isEqualTo("@minAge"); + } + + //endregion + + //region HPK (MULTI_HASH) tests + + @Test(groups = { "unit" }) + public void hpk_fullPk_defaultQuery() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + + PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); + List pkValues = Collections.singletonList(pk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("SELECT * FROM c\n WHERE"); + // Should use OR/AND pattern, not IN + assertThat(result.getQueryText()).doesNotContain("IN ("); + assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); + assertThat(result.getQueryText()).contains("AND"); + assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_1"); + assertThat(result.getParameters()).hasSize(2); + assertThat(result.getParameters().get(0).getValue(Object.class)).isEqualTo("Redmond"); + assertThat(result.getParameters().get(1).getValue(Object.class)).isEqualTo("98052"); + } + + @Test(groups = { "unit" }) + public void hpk_fullPk_multipleValues() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + + List pkValues = Arrays.asList( + new PartitionKeyBuilder().add("Redmond").add("98052").build(), + new PartitionKeyBuilder().add("Seattle").add("98101").build()); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("OR"); + assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); + assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_1"); + assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_2"); + assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_3"); + assertThat(result.getParameters()).hasSize(4); + } + + @Test(groups = { "unit" }) + public void hpk_partialPk_singleLevel() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); + List selectors = createSelectors(pkDef); + + // Partial PK — only first level + PartitionKey partialPk = new PartitionKeyBuilder().add("Redmond").build(); + List pkValues = Collections.singletonList(partialPk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); + // Should NOT include zipcode or areaCode since it's partial + assertThat(result.getQueryText()).doesNotContain("zipcode"); + assertThat(result.getQueryText()).doesNotContain("areaCode"); + assertThat(result.getParameters()).hasSize(1); + } + + @Test(groups = { "unit" }) + public void hpk_partialPk_twoLevels() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); + List selectors = createSelectors(pkDef); + + // Partial PK — first two levels + PartitionKey partialPk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); + List pkValues = Collections.singletonList(partialPk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); + assertThat(result.getQueryText()).contains("c[\"zipcode\"] = @__rmPk_1"); + assertThat(result.getQueryText()).doesNotContain("areaCode"); + assertThat(result.getParameters()).hasSize(2); + } + + @Test(groups = { "unit" }) + public void hpk_customQuery_withWhere() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@status", "active")); + + PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); + List pkValues = Collections.singletonList(pk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT c.name FROM c WHERE c.status = @status", baseParams, pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("WHERE (c.status = @status\n) AND ("); + assertThat(result.getQueryText()).contains("c[\"city\"] = @__rmPk_0"); + assertThat(result.getParameters()).hasSize(3); // @status + 2 pk params + } + + @Test(groups = { "unit" }) + public void normalizePartitionKeys_removesSubsumedFullHpkValues() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); + + List pkValues = Arrays.asList( + new PartitionKeyBuilder().add("Redmond").build(), + new PartitionKeyBuilder().add("Redmond").add("98052").add(1).build(), + new PartitionKeyBuilder().add("Seattle").build()); + + List normalizedPartitionKeys = + RxDocumentClientImpl.normalizePartitionKeys(pkValues, pkDef); + + assertThat(normalizedPartitionKeys) + .extracting(normalizedPk -> normalizedPk.effectivePkInternal.toJson()) + .containsExactly("[\"Redmond\"]", "[\"Seattle\"]"); + } + + @Test(groups = { "unit" }) + public void normalizePartitionKeys_rejectsNoneForMultiHashPartitionKeys() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + + assertThatThrownBy(() -> RxDocumentClientImpl.normalizePartitionKeys( + Collections.singletonList(PartitionKey.NONE), + pkDef)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("PartitionKey.NONE is not supported for multi-path partition keys"); + } + + @Test(groups = { "unit" }) + public void normalizePartitionKeys_epkAtBatchBoundary_isCorrectlyContained() { + // Verifies that a PK whose EPK sits exactly at a batchFilter boundary is correctly + // included in its owning range and excluded from the next range. The batchFilter is + // [minInclusive, maxExclusive) — a PK at the minInclusive boundary must be included, + // and a PK at the maxExclusive boundary must be excluded (it belongs to the next batch). + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List pkValues = Arrays.asList( + new PartitionKey("alpha"), + new PartitionKey("bravo"), + new PartitionKey("charlie")); + + List normalized = + RxDocumentClientImpl.normalizePartitionKeys(pkValues, pkDef); + + // After normalization the EPKs are sorted; take the middle one's EPK as a boundary + assertThat(normalized).hasSize(3); + String middleEpk = normalized.get(1).effectivePartitionKeyString; + + // Range [middleEpk, ...) must include the middle PK (minInclusive) + Range rangeIncluding = new Range<>(middleEpk, "FF", true, false); + assertThat(rangeIncluding.contains(middleEpk)).isTrue(); + + // Range [..., middleEpk) must exclude the middle PK (maxExclusive) + Range rangeExcluding = new Range<>("", middleEpk, true, false); + assertThat(rangeExcluding.contains(middleEpk)).isFalse(); + } + + @Test(groups = { "unit" }) + public void normalizePartitionKeys_noneInHpk_alwaysRejected() { + // PartitionKey.NONE must be rejected for MULTI_HASH regardless of whether this is the + // first call or a continuation-resume call (normalizePartitionKeys is invoked on both + // code paths). This test verifies the invariant holds for mixed inputs too. + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + + // NONE alone + assertThatThrownBy(() -> RxDocumentClientImpl.normalizePartitionKeys( + Collections.singletonList(PartitionKey.NONE), pkDef)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("PartitionKey.NONE is not supported for multi-path partition keys"); + + // NONE mixed with valid HPK values + assertThatThrownBy(() -> RxDocumentClientImpl.normalizePartitionKeys( + Arrays.asList( + new PartitionKeyBuilder().add("Redmond").add("98052").build(), + PartitionKey.NONE), + pkDef)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("PartitionKey.NONE is not supported for multi-path partition keys"); + } + + + @Test(groups = { "unit" }) + public void normalizePartitionKeys_twoPartialPrefixesAtDifferentDepths() { + // Two partial PKs at different prefix depths: (Redmond) and (Redmond, 98052). + // (Redmond) is a prefix of (Redmond, 98052), so (Redmond, 98052) should be collapsed. + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); + + List pkValues = Arrays.asList( + new PartitionKeyBuilder().add("Redmond").build(), + new PartitionKeyBuilder().add("Redmond").add("98052").build()); + + List normalized = + RxDocumentClientImpl.normalizePartitionKeys(pkValues, pkDef); + + assertThat(normalized) + .extracting(normalizedPk -> normalizedPk.effectivePkInternal.toJson()) + .containsExactly("[\"Redmond\"]"); + } + + @Test(groups = { "unit" }) + public void normalizePartitionKeys_partialPkDoesNotSubsumeUnrelatedPk() { + // (Redmond) and (Seattle, 98052) are unrelated - neither subsumes the other. + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); + + List pkValues = Arrays.asList( + new PartitionKeyBuilder().add("Redmond").build(), + new PartitionKeyBuilder().add("Seattle").add("98052").build()); + + List normalized = + RxDocumentClientImpl.normalizePartitionKeys(pkValues, pkDef); + + assertThat(normalized).hasSize(2); + assertThat(normalized) + .extracting(normalizedPk -> normalizedPk.effectivePkInternal.toJson()) + .containsExactly("[\"Redmond\"]", "[\"Seattle\",\"98052\"]"); + } + + @Test(groups = { "unit" }) + public void normalizePartitionKeys_allPartialPksPreserved() { + // All PKs are partial (single-level) and unrelated - all should survive. + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode", "/areaCode"); + + List pkValues = Arrays.asList( + new PartitionKeyBuilder().add("Redmond").build(), + new PartitionKeyBuilder().add("Seattle").build(), + new PartitionKeyBuilder().add("Pittsburgh").build()); + + List normalized = + RxDocumentClientImpl.normalizePartitionKeys(pkValues, pkDef); + + assertThat(normalized).hasSize(3); + } + //endregion + + //region findTopLevelWhereIndex tests + + @Test(groups = { "unit" }) + public void findWhere_simpleQuery() { + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM C WHERE C.ID = 1"); + assertThat(idx).isEqualTo(16); + } + + @Test(groups = { "unit" }) + public void findWhere_noWhere() { + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM C"); + assertThat(idx).isEqualTo(-1); + } + + @Test(groups = { "unit" }) + public void findWhere_whereInSubquery() { + // WHERE inside parentheses (subquery) should be ignored + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT * FROM C WHERE EXISTS(SELECT VALUE T FROM T IN C.TAGS WHERE T = 'FOO')"); + // Should find the outer WHERE, not the inner one + assertThat(idx).isEqualTo(16); + } + + @Test(groups = { "unit" }) + public void findWhere_caseInsensitive() { + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM C WHERE C.X = 1"); + assertThat(idx).isGreaterThan(0); + } + + @Test(groups = { "unit" }) + public void findWhere_whereNotKeyword() { + // "ELSEWHERE" should not match + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex("SELECT * FROM ELSEWHERE"); + assertThat(idx).isEqualTo(-1); + } + + //endregion + + //region Custom alias tests + + @Test(groups = { "unit" }) + public void singlePk_customAlias() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Arrays.asList(new PartitionKey("pk1"), new PartitionKey("pk2")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT x.id, x.mypk FROM x", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).startsWith("SELECT x.id, x.mypk FROM x\n WHERE"); + assertThat(result.getQueryText()).contains("x[\"mypk\"] IN ("); + assertThat(result.getQueryText()).doesNotContain("c[\"mypk\"]"); + } + + @Test(groups = { "unit" }) + public void singlePk_customAlias_withWhere() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@cat", "HelloWorld")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT x.id, x.mypk FROM x WHERE x.category = @cat", baseParams, pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("WHERE (x.category = @cat\n) AND (x[\"mypk\"] IN ("); + } + + @Test(groups = { "unit" }) + public void hpk_customAlias() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + + PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); + List pkValues = Collections.singletonList(pk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT r.name FROM root r", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("r[\"city\"] = @__rmPk_0"); + assertThat(result.getQueryText()).contains("r[\"zipcode\"] = @__rmPk_1"); + assertThat(result.getQueryText()).doesNotContain("c[\""); + } + + //endregion + + //region extractTableAlias tests + + @Test(groups = { "unit" }) + public void extractAlias_defaultC() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT * FROM c")).isEqualTo("c"); + } + + @Test(groups = { "unit" }) + public void extractAlias_customX() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT x.id FROM x WHERE x.age > 5")).isEqualTo("x"); + } + + @Test(groups = { "unit" }) + public void extractAlias_rootWithAlias() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT r.name FROM root r")).isEqualTo("r"); + } + + @Test(groups = { "unit" }) + public void extractAlias_rootNoAlias() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT * FROM root")).isEqualTo("root"); + } + + @Test(groups = { "unit" }) + public void extractAlias_containerWithWhere() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("SELECT * FROM items WHERE items.status = 'active'")).isEqualTo("items"); + } + + @Test(groups = { "unit" }) + public void extractAlias_caseInsensitive() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias("select * from MyContainer where MyContainer.id = '1'")).isEqualTo("MyContainer"); + } + + //endregion + + + //region String literal handling tests (#1) + + @Test(groups = { "unit" }) + public void findWhere_ignoresWhereInsideStringLiteral() { + // WHERE inside a string literal should be ignored + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT * FROM c WHERE c.msg = 'use WHERE clause here'"); + // Should find the outer WHERE at position 16, not the one inside the string + assertThat(idx).isEqualTo(16); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresParenthesesInsideStringLiteral() { + // Parentheses inside string literal should not affect depth tracking + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT * FROM c WHERE c.name = 'foo(bar)' AND c.x = 1"); + assertThat(idx).isEqualTo(16); + } + + @Test(groups = { "unit" }) + public void findWhere_handlesUnbalancedParenInStringLiteral() { + // Unbalanced paren inside string literal must not corrupt depth + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT * FROM c WHERE c.val = 'open(' AND c.active = true"); + assertThat(idx).isEqualTo(16); + } + + @Test(groups = { "unit" }) + public void findWhere_handlesStringLiteralBeforeWhere() { + // String literal in SELECT before WHERE + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT 'WHERE' as label FROM c WHERE c.id = '1'"); + // The WHERE inside quotes should be ignored; the real WHERE is further along + assertThat(idx).isGreaterThan(30); + } + + @Test(groups = { "unit" }) + public void singlePk_customQuery_withStringLiteralContainingParens() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@msg", "hello")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.msg = 'test(value)WHERE'", baseParams, pkValues, selectors, pkDef); + + // Should correctly AND the PK filter to the real WHERE clause + assertThat(result.getQueryText()).contains("WHERE (c.msg = 'test(value)WHERE'\n) AND ("); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresKeywordInsideDoubleQuotedBracketNotation() { + // c["WHERE"] uses double-quoted bracket notation — the WHERE inside quotes + // must not be matched as the real WHERE keyword. + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT c[\"WHERE\"] FROM c WHERE c.status = 'active'"); + // The real WHERE is after "FROM c ", not inside the brackets + assertThat(idx).isEqualTo(25); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresFromInsideDoubleQuotedBracketNotation() { + // Property named "FROM" in bracket notation should not confuse the FROM-clause parser + String query = "SELECT c[\"FROM\"] FROM c WHERE c.x = 1"; + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias(query)).isEqualTo("c"); + assertThat(ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex(query)).isEqualTo(24); + } + + //endregion + + //region Dotted/underscore/bracket property access boundary tests + + @Test(groups = { "unit" }) + public void findWhere_ignoresWhereInDottedPropertyAccess() { + // c.where is a property named "where" — not the WHERE keyword + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT c.where, c.id FROM c WHERE c.status = 'active'"); + // "SELECT c.where, c.id FROM c " = 28 chars, so WHERE at 28 + assertThat(idx).isEqualTo(28); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresFromInDottedPropertyAccess() { + // c.from is a property — the real FROM is later + String query = "SELECT c.from FROM c WHERE c.x = 1"; + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias(query)).isEqualTo("c"); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresOrderInDottedPropertyAccess() { + // c.order shouldn't match ORDER keyword + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelKeywordIndex( + "SELECT c.order FROM c WHERE c.x = 1", "ORDER"); + assertThat(idx).isEqualTo(-1); // no ORDER BY in this query + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresKeywordInUnderscoredIdentifier() { + // where_clause is an identifier, not WHERE + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT c.where_clause FROM c WHERE c.x = 1"); + // "SELECT c.where_clause FROM c " = 29 chars, WHERE at 29 + assertThat(idx).isEqualTo(29); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresKeywordPrecededByUnderscore() { + // _where is an identifier + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT c._where FROM c WHERE c.x = 1"); + assertThat(idx).isEqualTo(23); + } + + @Test(groups = { "unit" }) + public void extractAlias_ignoresKeywordPrefixInPropertyName() { + // c.offset is a property — should not confuse isFollowedByReservedKeyword + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( + "SELECT c.offset_val FROM c WHERE c.x = 1")).isEqualTo("c"); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresKeywordFollowedByBracket() { + // WHERE[ should not match as a standalone WHERE keyword + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT c.WHERE[0] FROM c WHERE c.x = 1"); + // The real WHERE is at position 25 + assertThat(idx).isEqualTo(25); + } + + @Test(groups = { "unit" }) + public void findWhere_ignoresKeywordFollowedByDollar() { + // WHERE$1 is an identifier, not WHERE + int idx = ReadManyByPartitionKeyQueryHelper.findTopLevelWhereIndex( + "SELECT WHERE$1 FROM c WHERE c.x = 1"); + assertThat(idx).isEqualTo(22); + } + + //endregion + + //region OFFSET/LIMIT/HAVING alias detection tests (#9) + + @Test(groups = { "unit" }) + public void extractAlias_containerWithOffset() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( + "SELECT * FROM c OFFSET 10 LIMIT 5")).isEqualTo("c"); + } + + @Test(groups = { "unit" }) + public void extractAlias_containerWithLimit() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( + "SELECT * FROM c LIMIT 10")).isEqualTo("c"); + } + + @Test(groups = { "unit" }) + public void extractAlias_containerWithHaving() { + assertThat(ReadManyByPartitionKeyQueryHelper.extractTableAlias( + "SELECT c.status, COUNT(1) FROM c GROUP BY c.status HAVING COUNT(1) > 1")).isEqualTo("c"); + } + + @Test(groups = { "unit" }) + public void createSelectors_nestedPath() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/address/city"); + + assertThat(PartitionKeyQueryHelper.createPkSelectors(pkDef)) + .containsExactly("[\"address\"][\"city\"]"); + } + + @Test(groups = { "unit" }) + public void createSelectors_escapesQuotesInPathParts() { + // Verify the escaping logic directly: a path part containing a double quote + // must produce \" in the selector, not just a bare backslash. + // Use an unquoted simple path /my so PathParser returns "my" cleanly, + // then verify a path whose PathParser output contains a quote character. + // PathParser for /"my\"field" yields token: my\"field (literal \, ") + PartitionKeyDefinition pkDef = createSinglePkDefinition("/\"my\\\"field\""); + + List selectors = PartitionKeyQueryHelper.createPkSelectors(pkDef); + assertThat(selectors).hasSize(1); + String selector = selectors.get(0); + // Must contain the escaped quote sequence \" + assertThat(selector).contains("\\\""); + // Must NOT be just ["my\field"] (old bug: quote replaced with bare backslash) + assertThat(selector).isNotEqualTo("[\"my\\field\"]"); + } + + //endregion + + //region PartitionKey.NONE query generation tests + + @Test(groups = { "unit" }) + public void singlePk_nonePartitionKey_generatesNotIsDefined() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(PartitionKey.NONE); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("NOT IS_DEFINED(c[\"mypk\"])"); + assertThat(result.getParameters()).isEmpty(); + } + + @Test(groups = { "unit" }) + public void singlePk_mixedNoneAndNormal_generatesCombinedFilter() { + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Arrays.asList(new PartitionKey("pk1"), PartitionKey.NONE); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("IN ("); + assertThat(result.getQueryText()).contains("NOT IS_DEFINED(c[\"mypk\"])"); + assertThat(result.getQueryText()).contains("OR"); + } + + @Test(groups = { "unit" }) + public void hpk_nonePartitionKey_throwsForMultiHash() { + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(PartitionKey.NONE); + + // PartitionKey.NONE is not supported for multi-path partition keys - + // the SDK rejects it in normalizePartitionKeys before reaching the query helper. + // The query helper itself also rejects it defensively. + assertThatThrownBy(() -> ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c", new ArrayList<>(), pkValues, selectors, pkDef)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("PartitionKey.NONE is not supported for multi-path partition keys"); + } + + //endregion + + //region Trailing single-line comment tests + + @Test(groups = { "unit" }) + public void singlePk_customQuery_trailingSingleLineComment_noWhere() { + // A trailing -- comment must not swallow the appended WHERE clause + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c -- my note", new ArrayList<>(), pkValues, selectors, pkDef); + + // WHERE must appear on a new line so it is not inside the -- comment + assertThat(result.getQueryText()).contains("\n WHERE"); + assertThat(result.getQueryText()).contains("IN ("); + assertThat(result.getQueryText()).contains("@__rmPk_0"); + assertThat(result.getParameters()).hasSize(1); + } + + @Test(groups = { "unit" }) + public void singlePk_customQuery_trailingBlockComment_noWhere() { + // Block comments are not affected, but verify the query still works + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c /* note */", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("WHERE"); + assertThat(result.getQueryText()).contains("IN ("); + assertThat(result.getParameters()).hasSize(1); + } + + @Test(groups = { "unit" }) + public void hpk_customQuery_trailingSingleLineComment_noWhere() { + // HPK path with trailing -- comment. + // Note: extractTableAlias does not skip comments, so "-- comment" after the + // FROM alias is mis-parsed as an alias token. This test verifies the \n WHERE + // fix and parameter correctness; alias-aware comment handling is a separate concern. + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + + PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); + List pkValues = Collections.singletonList(pk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c -- WHERE comment", new ArrayList<>(), pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("\n WHERE"); + assertThat(result.getQueryText()).contains("@__rmPk_0"); + assertThat(result.getQueryText()).contains("@__rmPk_1"); + assertThat(result.getParameters()).hasSize(2); + } + + @Test(groups = { "unit" }) + public void singlePk_customQuery_trailingComment_withExistingWhere() { + // When the query already has a WHERE clause, trailing comment after the condition + // is handled by the AND-merge path, not the no-WHERE path + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@minAge", 18)); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.age > @minAge -- filter", baseParams, pkValues, selectors, pkDef); + + assertThat(result.getQueryText()).contains("WHERE (c.age > @minAge -- filter\n) AND ("); + assertThat(result.getQueryText()).contains("IN ("); + } + + @Test(groups = { "unit" }) + public void singlePk_existingWhere_trailingLineComment_andClauseNotSwallowed() { + // Without the \n before ), the -- comment would swallow ") AND (" and the entire + // PK filter, producing a query that returns unfiltered results. + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Arrays.asList(new PartitionKey("pk1"), new PartitionKey("pk2")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@status", "active")); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.status = @status -- only active", baseParams, pkValues, selectors, pkDef); + + // The \n must break the -- comment so that ) AND ( is on a new line + String queryText = result.getQueryText(); + assertThat(queryText).contains("-- only active\n) AND ("); + assertThat(queryText).contains("IN ("); + assertThat(queryText).contains("@__rmPk_0"); + assertThat(queryText).contains("@__rmPk_1"); + assertThat(result.getParameters()).hasSize(3); // @status + 2 pk params + } + + @Test(groups = { "unit" }) + public void hpk_existingWhere_trailingLineComment_andClauseNotSwallowed() { + // HPK variant: trailing -- in existing WHERE must not swallow the AND clause + PartitionKeyDefinition pkDef = createMultiHashPkDefinition("/city", "/zipcode"); + List selectors = createSelectors(pkDef); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@active", true)); + + PartitionKey pk = new PartitionKeyBuilder().add("Redmond").add("98052").build(); + List pkValues = Collections.singletonList(pk); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.active = @active -- check", baseParams, pkValues, selectors, pkDef); + + String queryText = result.getQueryText(); + assertThat(queryText).contains("-- check\n) AND ("); + assertThat(queryText).contains("c[\"city\"] = @__rmPk_0"); + assertThat(queryText).contains("c[\"zipcode\"] = @__rmPk_1"); + assertThat(result.getParameters()).hasSize(3); + } + + @Test(groups = { "unit" }) + public void singlePk_existingWhere_multipleLineComments_andClauseNotSwallowed() { + // WHERE clause with a -- comment mid-condition (on its own line) and at the end + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@minAge", 18)); + baseParams.add(new SqlParameter("@maxAge", 65)); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.age > @minAge -- lower bound\n AND c.age < @maxAge -- upper bound", + baseParams, pkValues, selectors, pkDef); + + String queryText = result.getQueryText(); + // The final -- upper bound must be broken by \n before ) + assertThat(queryText).contains("-- upper bound\n) AND ("); + assertThat(queryText).contains("IN ("); + assertThat(queryText).contains("@__rmPk_0"); + assertThat(result.getParameters()).hasSize(3); + } + + @Test(groups = { "unit" }) + public void singlePk_existingWhere_blockComment_noNewlineNeeded() { + // Block comments (/* */) are terminated by */ and don't need \n protection, + // but the \n is harmless. Verify the query is still correct. + PartitionKeyDefinition pkDef = createSinglePkDefinition("/mypk"); + List selectors = createSelectors(pkDef); + List pkValues = Collections.singletonList(new PartitionKey("pk1")); + + List baseParams = new ArrayList<>(); + baseParams.add(new SqlParameter("@val", 42)); + + SqlQuerySpec result = ReadManyByPartitionKeyQueryHelper.createReadManyByPkQuerySpec( + "SELECT * FROM c WHERE c.x = @val /* note */", baseParams, pkValues, selectors, pkDef); + + String queryText = result.getQueryText(); + assertThat(queryText).contains("WHERE (c.x = @val /* note */\n) AND ("); + assertThat(queryText).contains("IN ("); + assertThat(result.getParameters()).hasSize(2); + } + + //endregion + + //region helpers + private PartitionKeyDefinition createSinglePkDefinition(String path) { + PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); + pkDef.setKind(PartitionKind.HASH); + pkDef.setVersion(PartitionKeyDefinitionVersion.V2); + pkDef.setPaths(Collections.singletonList(path)); + return pkDef; + } + + private PartitionKeyDefinition createMultiHashPkDefinition(String... paths) { + PartitionKeyDefinition pkDef = new PartitionKeyDefinition(); + pkDef.setKind(PartitionKind.MULTI_HASH); + pkDef.setVersion(PartitionKeyDefinitionVersion.V2); + pkDef.setPaths(Arrays.asList(paths)); + return pkDef; + } + private List createSelectors(PartitionKeyDefinition pkDef) { + return PartitionKeyQueryHelper.createPkSelectors(pkDef); + } + + //endregion +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java new file mode 100644 index 000000000000..2bcaa8810d4f --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryPlanValidationTest.java @@ -0,0 +1,91 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + */ + +package com.azure.cosmos.implementation; + +import com.azure.cosmos.implementation.query.DCountInfo; +import com.azure.cosmos.implementation.query.PartitionedQueryExecutionInfo; +import com.azure.cosmos.implementation.query.QueryInfo; +import com.azure.cosmos.implementation.query.hybridsearch.HybridSearchQueryInfo; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.testng.annotations.Test; + +import static org.assertj.core.api.Assertions.assertThatCode; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public class ReadManyByPartitionKeyQueryPlanValidationTest { + + @Test(groups = { "unit" }) + public void rejectsDCountQueryPlan() { + QueryInfo queryInfo = new QueryInfo(); + DCountInfo dCountInfo = new DCountInfo(); + dCountInfo.setDCountAlias("countAlias"); + queryInfo.set("dCountInfo", dCountInfo); + + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKeys(createQueryPlan(queryInfo, null))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("DCOUNT"); + } + + @Test(groups = { "unit" }) + public void rejectsOffsetQueryPlan() { + QueryInfo queryInfo = new QueryInfo(); + queryInfo.set("offset", 10); + + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKeys(createQueryPlan(queryInfo, null))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("OFFSET"); + } + + @Test(groups = { "unit" }) + public void rejectsLimitQueryPlan() { + QueryInfo queryInfo = new QueryInfo(); + queryInfo.set("limit", 10); + + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKeys(createQueryPlan(queryInfo, null))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("LIMIT"); + } + + @Test(groups = { "unit" }) + public void rejectsTopQueryPlan() { + QueryInfo queryInfo = new QueryInfo(); + queryInfo.set("top", 5); + + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKeys(createQueryPlan(queryInfo, null))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("TOP"); + } + + @Test(groups = { "unit" }) + public void rejectsHybridSearchQueryPlanWithoutDereferencingNullQueryInfo() { + assertThatThrownBy(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKeys( + createQueryPlan(null, new HybridSearchQueryInfo()))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("hybrid/vector/full-text"); + } + + @Test(groups = { "unit" }) + public void acceptsSimpleQueryPlan() { + QueryInfo queryInfo = new QueryInfo(); + + assertThatCode(() -> RxDocumentClientImpl.validateQueryPlanForReadManyByPartitionKeys(createQueryPlan(queryInfo, null))) + .doesNotThrowAnyException(); + } + + private PartitionedQueryExecutionInfo createQueryPlan(QueryInfo queryInfo, HybridSearchQueryInfo hybridSearchQueryInfo) { + ObjectNode content = Utils.getSimpleObjectMapper().createObjectNode(); + content.put("partitionedQueryExecutionInfoVersion", Constants.PartitionedQueryExecutionInfo.VERSION_1); + + if (queryInfo != null) { + content.set("queryInfo", Utils.getSimpleObjectMapper().valueToTree(queryInfo.getMap())); + } + if (hybridSearchQueryInfo != null) { + content.set("hybridSearchQueryInfo", Utils.getSimpleObjectMapper().createObjectNode()); + } + + return new PartitionedQueryExecutionInfo(content, null); + } +} diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index c17648b95294..05949f653fbe 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -4,10 +4,13 @@ #### Features Added * Added support for change feed with `startFrom` point-in-time on merged partitions by enabling the `CHANGE_FEED_WITH_START_TIME_POST_MERGE` SDK capability. - See [PR 48752](https://github.com/Azure/azure-sdk-for-java/pull/48752) +* Added new `readManyByPartitionKeys` API on `CosmosAsyncContainer` / `CosmosContainer` to bulk-query all documents matching a list of partition key values with better efficiency than issuing individual queries. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added `CosmosReadManyByPartitionKeysRequestOptions` - a dedicated request-options type for `readManyByPartitionKeys` that exposes `setContinuationToken(String)` for resuming previous invocations and `setMaxConcurrentBatchPrefetch(int)` to bound per-call prefetch parallelism. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) #### Breaking Changes #### Bugs Fixed +* Fixed `readMany` and `readAllItems` returning incorrect results on containers whose partition key path is nested (e.g. `/address/city`) due to malformed selector generation. - See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) * Fixed an issue where the throughput control `throughputQueryMono` was always subscribed even when `targetThroughput` is used (not `targetThroughputThreshold`), causing unnecessary `throughputSettings/read` permission requirement for AAD principals. - See [PR 48800](https://github.com/Azure/azure-sdk-for-java/pull/48800) * Fixed JVM `` deadlock when multiple threads concurrently trigger Cosmos SDK class loading for the first time. - See [PR 48689](https://github.com/Azure/azure-sdk-for-java/pull/48689) * Fixed an issue where `CustomItemSerializer` was incorrectly applied to internal SDK query pipeline structures (e.g., `OrderByRowResult`, `Document`), causing deserialization failures in ORDER BY, GROUP BY, aggregate, DISTINCT, and hybrid search queries. - See [PR 48811](https://github.com/Azure/azure-sdk-for-java/pull/48811) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index ad871bb97c01..3bd56a58d385 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -60,6 +60,7 @@ import com.azure.cosmos.models.CosmosPatchItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions; import com.azure.cosmos.models.CosmosReadManyRequestOptions; import com.azure.cosmos.models.FeedRange; import com.azure.cosmos.models.FeedResponse; @@ -126,6 +127,10 @@ private static ImplementationBridgeHelpers.CosmosReadManyRequestOptionsHelper.Co return ImplementationBridgeHelpers.CosmosReadManyRequestOptionsHelper.getCosmosReadManyRequestOptionsAccessor(); } + private static ImplementationBridgeHelpers.CosmosReadManyByPartitionKeysRequestOptionsHelper.CosmosReadManyByPartitionKeysRequestOptionsAccessor readManyByPkOptionsAccessor() { + return ImplementationBridgeHelpers.CosmosReadManyByPartitionKeysRequestOptionsHelper.getCosmosReadManyByPartitionKeysRequestOptionsAccessor(); + } + private static ImplementationBridgeHelpers.CosmosDiagnosticsContextHelper.CosmosDiagnosticsContextAccessor ctxAccessor() { return ImplementationBridgeHelpers.CosmosDiagnosticsContextHelper.getCosmosDiagnosticsContextAccessor(); } @@ -148,6 +153,12 @@ private static ImplementationBridgeHelpers.CosmosBatchRequestOptionsHelper.Cosmo private final static Logger logger = LoggerFactory.getLogger(CosmosAsyncContainer.class); + // Sentinel values for CosmosQueryRequestOptions.maxDegreeOfParallelism used when no + // caller-supplied value is present: 0 == "uninitialized / SDK-chooses", + // -1 == "unbounded" (the value RxDocumentClientImpl interprets as no concurrency cap). + private static final int DEFAULT_MAX_DEGREE_OF_PARALLELISM = 0; + private static final int UNBOUNDED_MAX_DEGREE_OF_PARALLELISM = -1; + private final CosmosAsyncDatabase database; private final String id; private final String link; @@ -165,6 +176,7 @@ private static ImplementationBridgeHelpers.CosmosBatchRequestOptionsHelper.Cosmo private final String createItemSpanName; private final String readAllItemsSpanName; private final String readManyItemsSpanName; + private final String readManyByPartitionKeysSpanName; private final String readAllItemsOfLogicalPartitionSpanName; private final String queryItemsSpanName; private final String queryChangeFeedSpanName; @@ -198,6 +210,7 @@ protected CosmosAsyncContainer(CosmosAsyncContainer toBeWrappedContainer) { this.createItemSpanName = "createItem." + this.id; this.readAllItemsSpanName = "readAllItems." + this.id; this.readManyItemsSpanName = "readManyItems." + this.id; + this.readManyByPartitionKeysSpanName = "readManyByPartitionKeys." + this.id; this.readAllItemsOfLogicalPartitionSpanName = "readAllItemsOfLogicalPartition." + this.id; this.queryItemsSpanName = "queryItems." + this.id; this.queryChangeFeedSpanName = "queryChangeFeed." + this.id; @@ -1601,6 +1614,228 @@ private Mono> readManyInternal( context); } + /** + * Reads many documents matching the provided partition key values. + * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries + * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} + * as the base query. Duplicate partition key inputs are normalized with set-based semantics + * before batching, so repeated keys do not duplicate the results. + *

+ * {@link PartitionKey#NONE} is supported for single-path partition key containers only. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param classType class type + * @return a {@link CosmosPagedFlux} containing one or several feed response pages + */ + public CosmosPagedFlux readManyByPartitionKeys( + List partitionKeys, + Class classType) { + + return this.readManyByPartitionKeys(partitionKeys, null, null, classType); + } + + /** + * Reads many documents matching the provided partition key values. + * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries + * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} + * as the base query. Duplicate partition key inputs are normalized with set-based semantics + * before batching, so repeated keys do not duplicate the results. + *

+ * {@link PartitionKey#NONE} is supported for single-path partition key containers only. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param requestOptions the optional request options + * @param classType class type + * @return a {@link CosmosPagedFlux} containing one or several feed response pages + */ + public CosmosPagedFlux readManyByPartitionKeys( + List partitionKeys, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, + Class classType) { + + return this.readManyByPartitionKeys(partitionKeys, null, requestOptions, classType); + } + + /** + * Reads many documents matching the provided partition key values with a custom query. + * The custom query can be used to apply projections (e.g. {@code SELECT c.name, c.age FROM c}) + * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). + * The SDK will automatically append partition key filtering to the custom query. + *

+ * The custom query must be a simple streamable query - aggregates, ORDER BY, DISTINCT, + * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be + * rejected. + *

+ * Partial hierarchical partition keys are supported and will fan out to multiple + * physical partitions. Duplicate partition key inputs are normalized with set-based semantics + * before batching. + *

+ * {@link PartitionKey#NONE} is supported for single-path partition key containers only. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) + * - should not contain WHERE clause filters for PK + * @param classType class type + * @return a {@link CosmosPagedFlux} containing one or several feed response pages + */ + public CosmosPagedFlux readManyByPartitionKeys( + List partitionKeys, + SqlQuerySpec customQuery, + Class classType) { + + return this.readManyByPartitionKeys(partitionKeys, customQuery, null, classType); + } + + /** + * Reads many documents matching the provided partition key values with a custom query. + * The custom query can be used to apply projections (e.g. {@code SELECT c.name, c.age FROM c}) + * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). + * The SDK will automatically append partition key filtering to the custom query. + *

+ * The custom query must be a simple streamable query - aggregates, ORDER BY, DISTINCT, + * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be + * rejected. + *

+ * Partial hierarchical partition keys are supported and will fan out to multiple + * physical partitions. Duplicate partition key inputs are normalized with set-based semantics + * before batching. + *

+ * {@link PartitionKey#NONE} is allowed for single-path partition key containers — it + * queries documents where the partition key path is absent (generates + * {@code NOT IS_DEFINED(...)}). For hierarchical (multi-path) partition key containers, + * {@link PartitionKey#NONE} is rejected because {@code addNoneValue()} is not supported + * with multiple paths. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) + * - should not contain WHERE clause filters for PK + * @param requestOptions the optional request options + * @param classType class type + * @return a {@link CosmosPagedFlux} containing one or several feed response pages + */ + public CosmosPagedFlux readManyByPartitionKeys( + List partitionKeys, + SqlQuerySpec customQuery, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, + Class classType) { + + checkNotNull(partitionKeys, "Argument 'partitionKeys' must not be null."); + checkNotNull(classType, "Argument 'classType' must not be null."); + + if (partitionKeys.isEmpty()) { + throw new IllegalArgumentException("Argument 'partitionKeys' must not be empty."); + } + + for (PartitionKey pk : partitionKeys) { + if (pk == null) { + throw new IllegalArgumentException( + "Argument 'partitionKeys' must not contain null elements."); + } + } + + List partitionKeysSnapshot = new ArrayList<>(partitionKeys); + + return UtilBridgeInternal.createCosmosPagedFlux( + readManyByPartitionKeysInternalFunc(partitionKeysSnapshot, customQuery, requestOptions, classType)); + } + + private Function>> readManyByPartitionKeysInternalFunc( + List partitionKeys, + SqlQuerySpec customQuery, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, + Class classType) { + + CosmosAsyncClient client = this.getDatabase().getClient(); + + // Extract continuation token before entering the reactive chain. + // It will be set on the cloned CosmosQueryRequestOptions so that + // QueryFeedOperationState picks it up automatically. + String requestContinuation = requestOptions != null + ? readManyByPkOptionsAccessor().getContinuationToken(requestOptions) + : null; + + // Resolve the max-concurrent-batch-prefetch knob. Unlike queryItems (which fans out + // many small reads across partitions), readManyByPartitionKeys typically returns large + // result sets per partition, so a conservative default avoids overwhelming individual + // partitions with parallel requests and spiking RU consumption. + Integer prefetchOverride = requestOptions != null + ? readManyByPkOptionsAccessor().getMaxConcurrentBatchPrefetch(requestOptions) + : null; + int maxConcurrentBatchPrefetch = prefetchOverride != null + ? prefetchOverride + : Math.max(1, Math.min(Configs.getCPUCnt(), 8)); + + return (pagedFluxOptions -> { + CosmosQueryRequestOptions queryRequestOptions = requestOptions == null + ? new CosmosQueryRequestOptions() + : queryOptionsAccessor().clone(readManyByPkOptionsAccessor().getImpl(requestOptions)); + // CosmosQueryRequestOptionsBase initializes MaxDegreeOfParallelism to 0 (the + // "uninitialized / SDK-chooses" sentinel); -1 is the "unbounded" sentinel that + // RxDocumentClientImpl recognizes for query parallelism. Honor any caller-provided + // value but default the uninitialized case to unbounded for readManyByPartitionKeys. + // CosmosReadManyRequestOptions does not currently expose MDOP, so this only matters + // if it is plumbed through in the future. + // When cloning from CosmosReadManyByPartitionKeysRequestOptionsImpl (which does not + // define maxDegreeOfParallelism), the cloned CosmosQueryRequestOptions may have a + // null backing Integer. Guard against NPE from auto-unboxing by checking for null. + Integer mdop = queryOptionsAccessor().getImpl(queryRequestOptions).getMaxDegreeOfParallelism(); + if (mdop == null || mdop == DEFAULT_MAX_DEGREE_OF_PARALLELISM) { + queryRequestOptions.setMaxDegreeOfParallelism(UNBOUNDED_MAX_DEGREE_OF_PARALLELISM); + } + + // maxItemCount lives in CosmosQueryRequestOptionsImpl but is not copied by the + // clone path through CosmosQueryRequestOptionsBase, so propagate it explicitly. + Integer maxItems = queryOptionsAccessor().getImpl(queryRequestOptions).getMaxItemCount(); + if (maxItems == null && requestOptions != null) { + Integer userMaxItems = readManyByPkOptionsAccessor().getImpl(requestOptions).getMaxItemCount(); + if (userMaxItems != null) { + ModelBridgeInternal.setQueryRequestOptionsMaxItemCount(queryRequestOptions, userMaxItems); + } + } + + queryRequestOptions.setQueryName("readManyByPartitionKeys"); + + // Set the composite continuation token on the cloned query options so that + // QueryFeedOperationState carries it through to RxDocumentClientImpl. + if (requestContinuation != null) { + ModelBridgeInternal.setQueryRequestOptionsContinuationToken( + queryRequestOptions, requestContinuation); + } + + CosmosQueryRequestOptionsBase cosmosQueryRequestOptionsImpl = queryOptionsAccessor().getImpl(queryRequestOptions); + applyPolicies(OperationType.Query, ResourceType.Document, cosmosQueryRequestOptionsImpl, this.readManyByPartitionKeysSpanName); + + QueryFeedOperationState state = new QueryFeedOperationState( + client, + this.readManyByPartitionKeysSpanName, + database.getId(), + this.getId(), + ResourceType.Document, + OperationType.Query, + queryOptionsAccessor().getQueryNameOrDefault(queryRequestOptions, this.readManyByPartitionKeysSpanName), + queryRequestOptions, + pagedFluxOptions + ); + + pagedFluxOptions.setFeedOperationState(state); + + return CosmosBridgeInternal + .getAsyncDocumentClient(this.getDatabase()) + .readManyByPartitionKeys( + partitionKeys, + customQuery, + BridgeInternal.getLink(this), + state, + maxConcurrentBatchPrefetch, + classType) + .map(response -> prepareFeedResponse(response, false)); + }); + } + /** * Reads all the items of a logical partition * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 04a6060c1927..d0146a75157b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -21,6 +21,7 @@ import com.azure.cosmos.models.CosmosPatchItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions; import com.azure.cosmos.models.CosmosReadManyRequestOptions; import com.azure.cosmos.models.FeedRange; import com.azure.cosmos.models.FeedResponse; @@ -540,6 +541,106 @@ public FeedResponse readMany( classType)); } + /** + * Reads many documents matching the provided partition key values. + * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries + * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} + * as the base query. Duplicate partition key inputs are normalized with set-based semantics + * before batching, so repeated keys do not duplicate the results. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param classType class type + * @return a {@link CosmosPagedIterable} containing the results + */ + public CosmosPagedIterable readManyByPartitionKeys( + List partitionKeys, + Class classType) { + + return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKeys(partitionKeys, classType)); + } + + /** + * Reads many documents matching the provided partition key values. + * Unlike {@link #readMany(List, Class)} this method does not require item ids - it queries + * all documents matching the provided partition key values. Uses {@code SELECT * FROM c} + * as the base query. Duplicate partition key inputs are normalized with set-based semantics + * before batching, so repeated keys do not duplicate the results. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param requestOptions the optional request options + * @param classType class type + * @return a {@link CosmosPagedIterable} containing the results + */ + public CosmosPagedIterable readManyByPartitionKeys( + List partitionKeys, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, + Class classType) { + + return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKeys(partitionKeys, requestOptions, classType)); + } + + /** + * Reads many documents matching the provided partition key values with a custom query. + * The custom query can be used to apply projections (e.g. {@code SELECT c.name, c.age FROM c}) + * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). + * The SDK will automatically append partition key filtering to the custom query. + *

+ * The custom query must be a simple streamable query - aggregates, ORDER BY, DISTINCT, + * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be + * rejected. + *

+ * Partial hierarchical partition keys are supported and will fan out to multiple + * physical partitions. Duplicate partition key inputs are normalized with set-based semantics + * before batching. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) + * - should not contain WHERE clause filters for PK + * @param classType class type + * @return a {@link CosmosPagedIterable} containing the results + */ + public CosmosPagedIterable readManyByPartitionKeys( + List partitionKeys, + SqlQuerySpec customQuery, + Class classType) { + + return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKeys(partitionKeys, customQuery, classType)); + } + + /** + * Reads many documents matching the provided partition key values with a custom query. + * The custom query can be used to apply projections (e.g. {@code SELECT c.name, c.age FROM c}) + * and/or additional filters (e.g. {@code SELECT * FROM c WHERE c.status = 'active'}). + * The SDK will automatically append partition key filtering to the custom query. + *

+ * The custom query must be a simple streamable query - aggregates, ORDER BY, DISTINCT, + * GROUP BY, DCOUNT, vector search, and full-text search are not supported and will be + * rejected. + *

+ * Partial hierarchical partition keys are supported and will fan out to multiple + * physical partitions. Duplicate partition key inputs are normalized with set-based semantics + * before batching. + * + * @param the type parameter + * @param partitionKeys list of partition key values to read documents for + * @param customQuery optional custom query for projections/additional filters (null means SELECT * FROM c) + * - should not contain WHERE clause filters for PK + * @param requestOptions the optional request options + * @param classType class type + * @return a {@link CosmosPagedIterable} containing the results + */ + public CosmosPagedIterable readManyByPartitionKeys( + List partitionKeys, + SqlQuerySpec customQuery, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, + Class classType) { + + return getCosmosPagedIterable(this.asyncContainer.readManyByPartitionKeys(partitionKeys, customQuery, requestOptions, classType)); + } + /** * Reads all the items of a logical partition returning the results as {@link CosmosPagedIterable}. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java index 945e768a82ff..3018188f274c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java @@ -1584,6 +1584,30 @@ Mono> readMany( QueryFeedOperationState state, Class klass); + /** + * Reads many documents by partition key values. + * Unlike {@link #readMany(List, String, QueryFeedOperationState, Class)} this method does not require + * item ids - it queries all documents matching the provided partition key values. + * Partial hierarchical partition keys are supported and will fan out to multiple physical partitions. + * + * @param partitionKeys list of partition key values to read documents for + * @param customQuery optional custom query (for projections/additional filters) - null means SELECT * FROM c + * @param collectionLink link for the documentcollection/container to be queried + * @param state the query operation state (may carry a composite continuation token via requestContinuation) + * @param maxConcurrentBatchPrefetch the maximum number of per-physical-partition batches whose first + * page is prefetched concurrently. Must be >= 1. + * @param klass class type + * @param the type parameter + * @return a Flux with feed response pages of documents + */ + Flux> readManyByPartitionKeys( + List partitionKeys, + SqlQuerySpec customQuery, + String collectionLink, + QueryFeedOperationState state, + int maxConcurrentBatchPrefetch, + Class klass); + /** * Read all documents of a certain logical partition. *

@@ -1651,7 +1675,7 @@ Flux> readAllDocuments( */ void enableSDKThroughputControlGroup(SDKThroughputControlGroupInternal group, Mono throughputQueryMono); - + /*** * Enable server throughput control group. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index b52b4e443cf3..18eef0544e18 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -59,7 +59,7 @@ public class Configs { private static final String NETTY_HTTP_CLIENT_METRICS_ENABLED = "COSMOS.NETTY_HTTP_CLIENT_METRICS_ENABLED"; private static final String NETTY_HTTP_CLIENT_METRICS_ENABLED_VARIABLE = "COSMOS_NETTY_HTTP_CLIENT_METRICS_ENABLED"; - // Thin client connect/acquire timeout — controls CONNECT_TIMEOUT_MILLIS for Gateway V2 data plane endpoints. + // Thin client connect/acquire timeout - controls CONNECT_TIMEOUT_MILLIS for Gateway V2 data plane endpoints. // Data plane requests are routed to the thin client regional endpoint (from RegionalRoutingContext) // which uses a non-443 port. These get a shorter 5s connect/acquire timeout. // Metadata requests target Gateway V1 endpoint (port 443) and retain the full 45s/60s timeout (unchanged). @@ -250,6 +250,11 @@ public class Configs { public static final String MIN_TARGET_BULK_MICRO_BATCH_SIZE_VARIABLE = "COSMOS_MIN_TARGET_BULK_MICRO_BATCH_SIZE"; public static final int DEFAULT_MIN_TARGET_BULK_MICRO_BATCH_SIZE = 1; + // readManyByPartitionKeys: max number of PK values per query per physical partition + public static final String READ_MANY_BY_PK_MAX_BATCH_SIZE = "COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE"; + public static final String READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE = "COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE"; + public static final int DEFAULT_READ_MANY_BY_PK_MAX_BATCH_SIZE = 100; + public static final String MAX_BULK_MICRO_BATCH_CONCURRENCY = "COSMOS.MAX_BULK_MICRO_BATCH_CONCURRENCY"; public static final String MAX_BULK_MICRO_BATCH_CONCURRENCY_VARIABLE = "COSMOS_MAX_BULK_MICRO_BATCH_CONCURRENCY"; public static final int DEFAULT_MAX_BULK_MICRO_BATCH_CONCURRENCY = 1; @@ -684,7 +689,7 @@ public static int getThinClientConnectionTimeoutInMs() { } } - // Guard against invalid values — timeout must be at least 500ms + // Guard against invalid values - timeout must be at least 500ms if (value < 500) { logger.warn( "Invalid thin client connection timeout: {}ms. Must be >= 500. Falling back to default: {}ms.", @@ -822,6 +827,46 @@ public static int getMinTargetBulkMicroBatchSize() { return DEFAULT_MIN_TARGET_BULK_MICRO_BATCH_SIZE; } + public static int getReadManyByPkMaxBatchSize() { + Integer parsed = parsePositiveInt(System.getProperty(READ_MANY_BY_PK_MAX_BATCH_SIZE), READ_MANY_BY_PK_MAX_BATCH_SIZE); + if (parsed != null) { + return parsed; + } + + parsed = parsePositiveInt(System.getenv(READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE), READ_MANY_BY_PK_MAX_BATCH_SIZE_VARIABLE); + if (parsed != null) { + return parsed; + } + + return DEFAULT_READ_MANY_BY_PK_MAX_BATCH_SIZE; + } + + /** + * Parses a non-empty string as a positive integer (>= 1). On parse failure or + * non-positive result, logs a WARN and returns null so the caller can fall back + * to its default. A null/empty input is also treated as "no value". + */ + private static Integer parsePositiveInt(String value, String configName) { + if (value == null || value.isEmpty()) { + return null; + } + try { + int parsed = Integer.parseInt(value); + if (parsed < 1) { + logger.warn( + "Ignoring invalid value '{}' for config '{}'. Value must be >= 1. Falling back to default.", + value, configName); + return null; + } + return parsed; + } catch (NumberFormatException e) { + logger.warn( + "Ignoring non-numeric value '{}' for config '{}'. Falling back to default.", + value, configName); + return null; + } + } + public static int getMaxBulkMicroBatchConcurrency() { String valueFromSystemProperty = System.getProperty(MAX_BULK_MICRO_BATCH_CONCURRENCY); if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeysRequestOptionsImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeysRequestOptionsImpl.java new file mode 100644 index 000000000000..d92eb0d2f1ab --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosReadManyByPartitionKeysRequestOptionsImpl.java @@ -0,0 +1,113 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation; + +/** + * Internal implementation backing the public {@code CosmosReadManyByPartitionKeysRequestOptions} + * facade. Holds state specific to the {@code readManyByPartitionKeys} operation. + */ +public class CosmosReadManyByPartitionKeysRequestOptionsImpl + extends CosmosQueryRequestOptionsBase { + + private String continuationToken; + private Integer maxConcurrentBatchPrefetch; + private Integer maxItemCount; + + public CosmosReadManyByPartitionKeysRequestOptionsImpl() { + super(); + } + + public CosmosReadManyByPartitionKeysRequestOptionsImpl(CosmosReadManyByPartitionKeysRequestOptionsImpl options) { + super(options); + this.continuationToken = options.continuationToken; + this.maxConcurrentBatchPrefetch = options.maxConcurrentBatchPrefetch; + this.maxItemCount = options.maxItemCount; + } + + /** + * Gets the composite continuation token for readManyByPartitionKeys. + * + * @return the continuation token, or null if not set. + */ + public String getContinuationToken() { + return this.continuationToken; + } + + /** + * Sets the composite continuation token for readManyByPartitionKeys. + * + * @param continuationToken the continuation token from a previous invocation. + * @return this instance. + */ + public CosmosReadManyByPartitionKeysRequestOptionsImpl setContinuationToken(String continuationToken) { + this.continuationToken = continuationToken; + return this; + } + + /** + * Gets the maximum number of per-physical-partition batches whose first page is + * prefetched concurrently. {@code null} means the SDK default applies. + * + * @return the max concurrent batch prefetch, or null if not set. + */ + public Integer getMaxConcurrentBatchPrefetch() { + return this.maxConcurrentBatchPrefetch; + } + + /** + * Sets the maximum number of per-physical-partition batches whose first page is + * prefetched concurrently. + * + * @param maxConcurrentBatchPrefetch the max concurrent batch prefetch (must be >= 1). + * @return this instance. + */ + public CosmosReadManyByPartitionKeysRequestOptionsImpl setMaxConcurrentBatchPrefetch(int maxConcurrentBatchPrefetch) { + this.maxConcurrentBatchPrefetch = maxConcurrentBatchPrefetch; + return this; + } + + @Override + public Integer getMaxItemCount() { + return this.maxItemCount; + } + + public CosmosReadManyByPartitionKeysRequestOptionsImpl setMaxItemCount(Integer maxItemCount) { + this.maxItemCount = maxItemCount; + return this; + } + + @Override + public Boolean isContentResponseOnWriteEnabled() { + return null; + } + + @Override + public Boolean getNonIdempotentWriteRetriesEnabled() { + return null; + } + + @Override + public Boolean isScanInQueryEnabled() { + return null; + } + + @Override + public Integer getMaxDegreeOfParallelism() { + return null; + } + + @Override + public Integer getMaxBufferedItemCount() { + return null; + } + + @Override + public Integer getMaxPrefetchPageCount() { + return null; + } + + @Override + public String getQueryNameOrDefault(String defaultQueryName) { + return defaultQueryName; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java index ee98360d2af4..ab882d89cf64 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java @@ -69,6 +69,7 @@ import com.azure.cosmos.models.CosmosPatchItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.CosmosReadManyByPartitionKeysRequestOptions; import com.azure.cosmos.models.CosmosReadManyRequestOptions; import com.azure.cosmos.models.FeedRange; import com.azure.cosmos.models.FeedResponse; @@ -354,6 +355,43 @@ public interface CosmosReadManyRequestOptionsAccessor { } } + public static final class CosmosReadManyByPartitionKeysRequestOptionsHelper { + private final static AtomicBoolean cosmosReadManyByPkRequestOptionsClassLoaded = new AtomicBoolean(false); + private final static AtomicReference accessor = new AtomicReference<>(); + + private CosmosReadManyByPartitionKeysRequestOptionsHelper() {} + + public static void setCosmosReadManyByPartitionKeysRequestOptionsAccessor( + final CosmosReadManyByPartitionKeysRequestOptionsAccessor newAccessor) { + if (!accessor.compareAndSet(null, newAccessor)) { + logger.debug("CosmosReadManyByPartitionKeysRequestOptionsAccessor already initialized!"); + } else { + logger.debug("Setting CosmosReadManyByPartitionKeysRequestOptionsAccessor..."); + cosmosReadManyByPkRequestOptionsClassLoaded.set(true); + } + } + + public static CosmosReadManyByPartitionKeysRequestOptionsAccessor getCosmosReadManyByPartitionKeysRequestOptionsAccessor() { + if (!cosmosReadManyByPkRequestOptionsClassLoaded.get()) { + logger.debug("Initializing CosmosReadManyByPartitionKeysRequestOptionsAccessor..."); + initializeAllAccessors(); + } + + CosmosReadManyByPartitionKeysRequestOptionsAccessor snapshot = accessor.get(); + if (snapshot == null) { + logger.error("CosmosReadManyByPartitionKeysRequestOptionsAccessor is not initialized yet!"); + } + + return snapshot; + } + + public interface CosmosReadManyByPartitionKeysRequestOptionsAccessor { + CosmosQueryRequestOptionsBase getImpl(CosmosReadManyByPartitionKeysRequestOptions options); + String getContinuationToken(CosmosReadManyByPartitionKeysRequestOptions options); + Integer getMaxConcurrentBatchPrefetch(CosmosReadManyByPartitionKeysRequestOptions options); + } + } + public static final class CosmosChangeFeedRequestOptionsHelper { private final static AtomicBoolean cosmosChangeFeedRequestOptionsClassLoaded = new AtomicBoolean(false); private final static AtomicReference accessor = new AtomicReference<>(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/PartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/PartitionKeyQueryHelper.java new file mode 100644 index 000000000000..459e03c8fce3 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/PartitionKeyQueryHelper.java @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation; + +import com.azure.cosmos.models.PartitionKeyDefinition; + +import java.util.List; +import java.util.stream.Collectors; + +/** + * Helper for constructing SQL partition key selector fragments from partition key definitions. + */ +final class PartitionKeyQueryHelper { + + private PartitionKeyQueryHelper() { + } + + static List createPkSelectors(PartitionKeyDefinition partitionKeyDefinition) { + return partitionKeyDefinition.getPaths() + .stream() + .map(PathParser::getPathParts) + .map(pathParts -> pathParts.stream() + .map(pathPart -> "[\"" + pathPart.replace("\"", "\\\"") + "\"]") + .collect(Collectors.joining())) + .collect(Collectors.toList()); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java new file mode 100644 index 000000000000..5c972ed3f10f --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationToken.java @@ -0,0 +1,430 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation; + +import com.azure.cosmos.implementation.routing.MurmurHash3_128; +import com.azure.cosmos.implementation.routing.Range; +import com.azure.cosmos.implementation.routing.UInt128; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.io.ByteArrayOutputStream; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Base64; +import java.util.Collections; +import java.util.List; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +/** + * Composite continuation token for {@code readManyByPartitionKeys} operations. + *

+ * Captures the state needed to resume a readManyByPartitionKeys operation: + *

    + *
  • {@code remainingBatches} - batch definitions of batches not yet started
    • + *
  • {@code currentBatch} - batch definition of the batch currently being processed
    • + *
  • {@code backendContinuation} - backend query continuation within the current batch (nullable)
    • + *
+ * Each batch is identified solely by its {@link BatchDefinition#getBatchFilter() batchFilter} + * EPK range - the half-open range {@code [minInclusive, maxExclusive)} that contains the EPKs + * of all PKs assigned to the batch. The physical-partition routing range used at execution + * time (the FeedRange set on {@code CosmosQueryRequestOptions}) is not persisted; + * it is rederived at execution time from the current PartitionKeyRange cache by taking the + * union of all partition-key-range EPK ranges that overlap the batch filter + * ({@code [min(minEpk), max(maxEpk))}). + *

+ * This means the token survives partition splits without ever encoding stale partition + * boundaries: after a split, the rederived routing range exactly matches the new physical + * partition boundaries (one or more of them), keeping query-RU cost minimal. It also keeps + * the serialized token small (one EPK range per batch instead of two). + *

+ * Serialized as JSON -> Base64 to keep the token opaque. The serialized form embeds a + * {@code "v"} version field so future format evolutions can be detected and rejected + * (or migrated) cleanly without silently misinterpreting an older token. + */ +public final class ReadManyByPartitionKeyContinuationToken { + + /** Wire format version. Bump on any breaking change to the JSON shape. */ + public static final int CURRENT_VERSION = 1; + + private static final String VERSION_PROPERTY = "v"; + private static final String REMAINING_BATCHES_PROPERTY = "rb"; + private static final String CURRENT_BATCH_PROPERTY = "cb"; + private static final String BACKEND_CONTINUATION_PROPERTY = "bc"; + private static final String COLLECTION_RID_PROPERTY = "cr"; + private static final String QUERY_HASH_PROPERTY = "qh"; + private static final String PARTITION_KEY_SET_HASH_PROPERTY = "ph"; + + @JsonProperty(VERSION_PROPERTY) + private final int version; + + @JsonProperty(REMAINING_BATCHES_PROPERTY) + private final List remainingBatches; + + @JsonProperty(CURRENT_BATCH_PROPERTY) + private final BatchDefinitionDto currentBatch; + + @JsonProperty(BACKEND_CONTINUATION_PROPERTY) + private final String backendContinuation; + + @JsonProperty(COLLECTION_RID_PROPERTY) + private final String collectionRid; + + @JsonProperty(QUERY_HASH_PROPERTY) + private final String queryHash; + + @JsonProperty(PARTITION_KEY_SET_HASH_PROPERTY) + private final String partitionKeySetHash; + + /** + * Production constructor used by RxDocumentClientImpl when stamping each FeedResponse + * with a continuation token. Callers supply just the batch filter for both the current + * batch and every remaining batch, plus all three identity fingerprints + * (collectionRid + queryHash + partitionKeySetHash). Routing scopes are not persisted. + */ + public ReadManyByPartitionKeyContinuationToken( + List remainingBatches, + BatchDefinition currentBatch, + String backendContinuation, + String collectionRid, + String queryHash, + String partitionKeySetHash) { + + checkNotNull(currentBatch, "Argument 'currentBatch' must not be null."); + checkNotNull(remainingBatches, "Argument 'remainingBatches' must not be null."); + + this.version = CURRENT_VERSION; + this.remainingBatches = new ArrayList<>(remainingBatches.size()); + for (BatchDefinition bd : remainingBatches) { + this.remainingBatches.add(BatchDefinitionDto.fromBatchDefinition(bd)); + } + this.currentBatch = BatchDefinitionDto.fromBatchDefinition(currentBatch); + this.backendContinuation = backendContinuation; + this.collectionRid = collectionRid; + this.queryHash = queryHash; + this.partitionKeySetHash = partitionKeySetHash; + } + + /** + * Deserialization constructor invoked by Jackson. Validates the version field so a + * token from an incompatible future SDK is rejected with a clear error rather than + * being silently misinterpreted. + */ + @JsonCreator + ReadManyByPartitionKeyContinuationToken( + @JsonProperty(VERSION_PROPERTY) Integer version, + @JsonProperty(REMAINING_BATCHES_PROPERTY) List remainingBatches, + @JsonProperty(CURRENT_BATCH_PROPERTY) BatchDefinitionDto currentBatch, + @JsonProperty(BACKEND_CONTINUATION_PROPERTY) String backendContinuation, + @JsonProperty(COLLECTION_RID_PROPERTY) String collectionRid, + @JsonProperty(QUERY_HASH_PROPERTY) String queryHash, + @JsonProperty(PARTITION_KEY_SET_HASH_PROPERTY) String partitionKeySetHash) { + + // Tokens written before the version field existed will deserialize with version == null. + // Treat null as version 1 (the format that existed when this field was introduced) to + // remain forward-compatible with any tokens emitted by an in-flight pre-versioned beta. + int effectiveVersion = (version == null) ? 1 : version; + if (effectiveVersion != CURRENT_VERSION) { + throw new IllegalArgumentException( + "Unsupported readManyByPartitionKeys continuation token version: " + effectiveVersion + + ". This SDK supports version " + CURRENT_VERSION + "."); + } + + this.version = effectiveVersion; + + if (remainingBatches == null) { + throw new IllegalArgumentException( + "Malformed readManyByPartitionKeys continuation token: 'remainingBatches' is required."); + } + if (currentBatch == null) { + throw new IllegalArgumentException( + "Malformed readManyByPartitionKeys continuation token: 'currentBatch' is required."); + } + this.remainingBatches = remainingBatches; + this.currentBatch = currentBatch; + this.backendContinuation = backendContinuation; + this.collectionRid = collectionRid; + this.queryHash = queryHash; + this.partitionKeySetHash = partitionKeySetHash; + } + + @JsonIgnore + public int getVersion() { + return version; + } + + @JsonIgnore + public List getRemainingBatches() { + List result = new ArrayList<>(remainingBatches.size()); + for (BatchDefinitionDto dto : remainingBatches) { + result.add(dto.toBatchDefinition()); + } + return result; + } + + @JsonIgnore + public BatchDefinition getCurrentBatch() { + return currentBatch.toBatchDefinition(); + } + + public String getBackendContinuation() { + return backendContinuation; + } + + public String getCollectionRid() { + return collectionRid; + } + + public String getQueryHash() { + return queryHash; + } + + public String getPartitionKeySetHash() { + return partitionKeySetHash; + } + + /** + * Computes a stable hash for a SqlQuerySpec (or null for the default SELECT * FROM c query). + * Hashes over both the query text and all parameter names/values to detect when a continuation + * token is reused with a different query or different parameter values. + */ + public static String computeQueryHash(SqlQuerySpec querySpec) { + if (querySpec == null) { + return "0"; + } + + try { + ByteArrayOutputStream output = new ByteArrayOutputStream(); + updateHashInput(output, querySpec.getQueryText()); + + List params = querySpec.getParameters(); + if (params != null) { + for (SqlParameter param : params) { + updateHashInput(output, param.getName()); + + Object value = param.getValue(Object.class); + if (value == null) { + updateHashInput(output, null); + } else { + output.write(Utils.getSimpleObjectMapper().writeValueAsBytes(value)); + output.write(0); + } + } + } + + return murmurHash128Hex(output.toByteArray()); + } catch (Exception e) { + throw new IllegalStateException("Failed to compute stable query hash for continuation token.", e); + } + } + + /** + * Computes a stable hash for the normalized set of partition key EPK values. + * Duplicate and reordered inputs intentionally produce the same digest. + */ + public static String computePartitionKeySetHash(List partitionKeyEpks) { + return computePartitionKeySetHash(partitionKeyEpks, false); + } + + /** + * Computes a stable hash for the normalized set of partition key EPK values. + * When {@code alreadySorted} is true, the input is assumed to be + * already sorted (e.g. from {@code normalizePartitionKeys}), + * skipping the O(n log n) sort pass. + */ + public static String computePartitionKeySetHash(List partitionKeyEpks, boolean alreadySorted) { + if (partitionKeyEpks == null || partitionKeyEpks.isEmpty()) { + return "0"; + } + + List normalizedEpks = null; + if (alreadySorted) { + normalizedEpks = partitionKeyEpks; + } else { + + normalizedEpks = new ArrayList<>(partitionKeyEpks.size()); + for (String epk : partitionKeyEpks) { + if (epk != null) { + normalizedEpks.add(epk); + } + } + + if (normalizedEpks.isEmpty()) { + return "0"; + } + + Collections.sort(normalizedEpks); + } + + ByteArrayOutputStream output = new ByteArrayOutputStream(); + String previous = null; + for (String epk : normalizedEpks) { + if (epk.equals(previous)) { + continue; + } + + updateHashInput(output, epk); + previous = epk; + } + + return murmurHash128Hex(output.toByteArray()); + } + + private static void updateHashInput(ByteArrayOutputStream output, String value) { + if (value != null) { + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); + output.write(bytes, 0, bytes.length); + } + output.write(0); + } + + private static String murmurHash128Hex(byte[] bytes) { + if (bytes == null || bytes.length == 0) { + return "0"; + } + + UInt128 hash = MurmurHash3_128.hash128(bytes, bytes.length); + return toFixedHex(hash.getHigh()) + toFixedHex(hash.getLow()); + } + + private static String toFixedHex(long value) { + String hex = Long.toHexString(value); + if (hex.length() >= 16) { + return hex; + } + + StringBuilder builder = new StringBuilder(16); + for (int i = hex.length(); i < 16; i++) { + builder.append('0'); + } + builder.append(hex); + return builder.toString(); + } + + /** + * Serializes this token to a URL-safe Base64-encoded JSON string. + */ + public String serialize() { + try { + String json = Utils.getSimpleObjectMapper().writeValueAsString(this); + return Base64.getUrlEncoder().encodeToString(json.getBytes(StandardCharsets.UTF_8)); + } catch (Exception e) { + throw new IllegalStateException("Failed to serialize ReadManyByPartitionKeyContinuationToken.", e); + } + } + + /** + * Deserializes a URL-safe Base64-encoded JSON string into a continuation token. + * + * @param serialized the serialized token (URL-safe Base64 of JSON) + * @return the deserialized token + * @throws IllegalArgumentException if the token is malformed + */ + public static ReadManyByPartitionKeyContinuationToken deserialize(String serialized) { + checkNotNull(serialized, "Argument 'serialized' must not be null."); + checkArgument(!serialized.isEmpty(), "Argument 'serialized' must not be empty."); + + try { + byte[] decoded = Base64.getUrlDecoder().decode(serialized); + String json = new String(decoded, StandardCharsets.UTF_8); + return Utils.getSimpleObjectMapper().readValue(json, ReadManyByPartitionKeyContinuationToken.class); + } catch (IllegalArgumentException e) { + // Preserve our own IllegalArgumentException (version mismatch, null fields) without wrapping. + throw e; + } catch (Exception e) { + // Jackson wraps constructor-thrown IllegalArgumentException inside JsonMappingException. + // Unwrap to preserve the actionable error message from our null checks. + Throwable cause = e.getCause(); + if (cause instanceof IllegalArgumentException) { + throw (IllegalArgumentException) cause; + } + throw new IllegalArgumentException( + "Failed to deserialize ReadManyByPartitionKeyContinuationToken. " + + "The continuation token may be malformed or from an incompatible version.", e); + } + } + + /** + * Identifies a single batch in a readManyByPartitionKeys operation by its EPK filter range. + *

+ * The {@code batchFilter} is the half-open EPK range {@code [minInclusive, maxExclusive)} + * containing the EPKs of all PKs assigned to the batch. It is the only piece of routing + * data persisted in the continuation token; the physical-partition scope used as the + * query FeedRange is rederived at execution time from the current PartitionKeyRange + * cache (union of overlapping partition-key-range EPK ranges) so partition splits do not + * cause stale routing information to be embedded in the token. + */ + public static final class BatchDefinition { + private final Range batchFilter; + + public BatchDefinition(Range batchFilter) { + this.batchFilter = checkNotNull(batchFilter, "Argument 'batchFilter' must not be null."); + } + + public Range getBatchFilter() { + return batchFilter; + } + } + + /** + * Compact DTO for JSON serialization of a batch definition. + * Persists only the batch filter; routing scope is rederived at execution time. + */ + static final class BatchDefinitionDto { + private final EpkRangeDto bf; + + @JsonCreator + BatchDefinitionDto(@JsonProperty("bf") EpkRangeDto bf) { + this.bf = bf; + } + + @JsonProperty("bf") + EpkRangeDto getBf() { return bf; } + + static BatchDefinitionDto fromBatchDefinition(BatchDefinition bd) { + return new BatchDefinitionDto(EpkRangeDto.fromRange(bd.batchFilter)); + } + + BatchDefinition toBatchDefinition() { + return new BatchDefinition(bf.toRange()); + } + } + + /** + * Minimal DTO for EPK range serialization. Uses short field names to keep the + * serialized token compact. + */ + static final class EpkRangeDto { + @JsonProperty("min") + private final String min; + @JsonProperty("max") + private final String max; + + @JsonCreator + EpkRangeDto( + @JsonProperty("min") String min, + @JsonProperty("max") String max) { + this.min = min; + this.max = max; + } + + static EpkRangeDto fromRange(Range range) { + if (!range.isMinInclusive()) { + throw new IllegalArgumentException("EPK ranges must be minInclusive."); + } + if (range.isMaxInclusive()) { + throw new IllegalArgumentException("EPK ranges must be maxExclusive."); + } + return new EpkRangeDto(range.getMin(), range.getMax()); + } + + Range toRange() { + return new Range<>(min, max, true, false); + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java new file mode 100644 index 000000000000..f498a452fad3 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java @@ -0,0 +1,321 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation; + +import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.implementation.routing.PartitionKeyInternal; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.PartitionKeyDefinition; +import com.azure.cosmos.models.PartitionKind; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; + +import java.util.ArrayList; +import java.util.List; +import java.util.Locale; + +/** + * Helper for constructing SqlQuerySpec instances for readManyByPartitionKeys operations. + * This class is not intended to be used directly by end-users. + */ +public class ReadManyByPartitionKeyQueryHelper { + + private static final String DEFAULT_TABLE_ALIAS = "c"; + // Internal parameter prefix - uses double-underscore to avoid collisions with user-provided parameters + private static final String PK_PARAM_PREFIX = "@__rmPk_"; + + public static SqlQuerySpec createReadManyByPkQuerySpec( + String baseQueryText, + List baseParameters, + List pkValues, + List partitionKeySelectors, + PartitionKeyDefinition pkDefinition) { + + // Guard against collisions with our internal parameter names - callers cannot realistically + // use the @__rmPk_ prefix for their own parameters, but if they do we surface a clear error + // rather than letting the server reject a SqlQuerySpec with duplicate parameter names. + for (SqlParameter baseParam : baseParameters) { + String name = baseParam.getName(); + if (name != null && name.startsWith(PK_PARAM_PREFIX)) { + throw new IllegalArgumentException( + "Custom query parameter name '" + name + "' collides with the reserved " + + "readManyByPartitionKeys internal prefix '" + PK_PARAM_PREFIX + "'. Rename the parameter."); + } + } + + // Extract the table alias from the FROM clause (e.g. "FROM x" -> "x", "FROM c" -> "c") + String tableAlias = extractTableAlias(baseQueryText); + + StringBuilder pkFilter = new StringBuilder(); + List parameters = new ArrayList<>(baseParameters); + int paramCount = 0; + + boolean isSinglePathPk = partitionKeySelectors.size() == 1; + + if (isSinglePathPk && pkDefinition.getKind() != PartitionKind.MULTI_HASH) { + // Single PK path - use IN clause for normal values, OR NOT IS_DEFINED for NONE + // First, separate NONE PKs from normal PKs + boolean hasNone = false; + List normalPkValues = new ArrayList<>(); + for (PartitionKey pk : pkValues) { + PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); + if (pkInternal.getComponents() == null) { + hasNone = true; + } else { + normalPkValues.add(pk); + } + } + + pkFilter.append(" "); + boolean hasNormalValues = !normalPkValues.isEmpty(); + if (hasNormalValues && hasNone) { + pkFilter.append("("); + } + if (hasNormalValues) { + pkFilter.append(tableAlias); + pkFilter.append(partitionKeySelectors.get(0)); + pkFilter.append(" IN ( "); + for (int i = 0; i < normalPkValues.size(); i++) { + PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(normalPkValues.get(i)); + Object[] pkComponents = pkInternal.toObjectArray(); + String pkParamName = PK_PARAM_PREFIX + paramCount; + parameters.add(new SqlParameter(pkParamName, pkComponents[0])); + paramCount++; + + pkFilter.append(pkParamName); + if (i < normalPkValues.size() - 1) { + pkFilter.append(", "); + } + } + pkFilter.append(" )"); + } + if (hasNone) { + if (hasNormalValues) { + pkFilter.append(" OR "); + } + pkFilter.append("NOT IS_DEFINED("); + pkFilter.append(tableAlias); + pkFilter.append(partitionKeySelectors.get(0)); + pkFilter.append(")"); + } + if (hasNormalValues && hasNone) { + pkFilter.append(")"); + } + } else { + // Multiple PK paths (HPK) or MULTI_HASH - use OR of AND clauses + pkFilter.append(" "); + for (int i = 0; i < pkValues.size(); i++) { + PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pkValues.get(i)); + Object[] pkComponents = pkInternal.toObjectArray(); + + if (pkComponents == null) { + throw new IllegalArgumentException( + "PartitionKey.NONE is not supported for multi-path partition keys in readManyByPartitionKeys."); + } + + { + pkFilter.append("("); + for (int j = 0; j < pkComponents.length; j++) { + String pkParamName = PK_PARAM_PREFIX + paramCount; + parameters.add(new SqlParameter(pkParamName, pkComponents[j])); + paramCount++; + + if (j > 0) { + pkFilter.append(" AND "); + } + pkFilter.append(tableAlias); + pkFilter.append(partitionKeySelectors.get(j)); + pkFilter.append(" = "); + pkFilter.append(pkParamName); + } + pkFilter.append(")"); + } + + if (i < pkValues.size() - 1) { + pkFilter.append(" OR "); + } + } + } + + // Compose final query: handle existing WHERE clause in base query + String finalQuery; + int whereIndex = findTopLevelWhereIndex(baseQueryText); + if (whereIndex >= 0) { + // Base query has WHERE - AND our PK filter + String beforeWhere = baseQueryText.substring(0, whereIndex); + String afterWhere = baseQueryText.substring(whereIndex + 5); // skip "WHERE" + finalQuery = beforeWhere + "WHERE (" + afterWhere.trim() + "\n) AND (" + pkFilter.toString().trim() + ")"; + } else { + // No WHERE - add one. Use \n before WHERE so that a trailing single-line comment + // (-- ...) in the base query does not swallow the WHERE clause. + finalQuery = baseQueryText + "\n WHERE" + pkFilter.toString(); + } + + return new SqlQuerySpec(finalQuery, parameters); + } + + /** + * Extracts the table/collection alias from a SQL query's FROM clause. + * Handles: "SELECT * FROM c", "SELECT x.id FROM x WHERE ...", "SELECT * FROM root r", etc. + * Returns the alias used after FROM (last token before WHERE or end of FROM clause). + */ + static String extractTableAlias(String queryText) { + String upper = queryText.toUpperCase(Locale.ROOT); + int fromIndex = findTopLevelKeywordIndex(upper, "FROM"); + if (fromIndex < 0) { + return DEFAULT_TABLE_ALIAS; + } + + // Start scanning after "FROM" + int afterFrom = fromIndex + 4; + // Skip whitespace + while (afterFrom < queryText.length() && Character.isWhitespace(queryText.charAt(afterFrom))) { + afterFrom++; + } + + // Collect the container name token (could be "root", "c", etc.) + int tokenStart = afterFrom; + while (afterFrom < queryText.length() + && !Character.isWhitespace(queryText.charAt(afterFrom)) + && queryText.charAt(afterFrom) != '(' + && queryText.charAt(afterFrom) != ')') { + afterFrom++; + } + String containerName = queryText.substring(tokenStart, afterFrom); + + // Skip whitespace after container name + while (afterFrom < queryText.length() && Character.isWhitespace(queryText.charAt(afterFrom))) { + afterFrom++; + } + + // Check if there's an alias after the container name (before WHERE or end) + if (afterFrom < queryText.length()) { + String remaining = upper.substring(afterFrom); + // Reserved keywords that terminate the FROM clause - when the next token is one of these, + // containerName itself IS the alias used throughout the rest of the query. + if (isFollowedByReservedKeyword(remaining)) { + return containerName; + } + // Handle optional AS: "FROM root AS r" -> alias is "r" + if (remaining.startsWith("AS") + && (remaining.length() == 2 || !isIdentifierChar(remaining.charAt(2)))) { + afterFrom += 2; // skip AS + while (afterFrom < queryText.length() + && Character.isWhitespace(queryText.charAt(afterFrom))) { + afterFrom++; + } + } + // Otherwise the next token is the alias ("FROM root r" -> alias is "r") + int aliasStart = afterFrom; + while (afterFrom < queryText.length() + && !Character.isWhitespace(queryText.charAt(afterFrom)) + && queryText.charAt(afterFrom) != '(' + && queryText.charAt(afterFrom) != ')') { + afterFrom++; + } + if (afterFrom > aliasStart) { + return queryText.substring(aliasStart, afterFrom); + } + } + + return containerName; + } + + private static boolean isFollowedByReservedKeyword(String remainingUpper) { + String[] keywords = { "WHERE", "ORDER", "GROUP", "JOIN", "OFFSET", "LIMIT", "HAVING" }; + for (String kw : keywords) { + if (remainingUpper.startsWith(kw) + && (remainingUpper.length() == kw.length() + || !isIdentifierChar(remainingUpper.charAt(kw.length())))) { + return true; + } + } + return false; + } + + /** + * Finds the index of a top-level SQL keyword in the query text (case-insensitive), + * ignoring occurrences inside parentheses or string literals. + */ + static int findTopLevelKeywordIndex(String queryText, String keyword) { + String queryTextUpper = queryText.toUpperCase(Locale.ROOT); + String keywordUpper = keyword.toUpperCase(Locale.ROOT); + int depth = 0; + int keyLen = keywordUpper.length(); + int len = queryTextUpper.length(); + for (int i = 0; i <= len - keyLen; i++) { + char ch = queryText.charAt(i); + // Skip single-line comments: -- ... end-of-line + if (ch == '-' && i + 1 < len && queryText.charAt(i + 1) == '-') { + i += 2; + while (i < len && queryText.charAt(i) != '\n' && queryText.charAt(i) != '\r') { + i++; + } + continue; + } + // Skip block comments: /* ... */ + if (ch == '/' && i + 1 < len && queryText.charAt(i + 1) == '*') { + i += 2; + while (i + 1 < len + && !(queryText.charAt(i) == '*' && queryText.charAt(i + 1) == '/')) { + i++; + } + i++; // position on the '/'; loop post-increment moves past it + continue; + } + // Skip string literals enclosed in single quotes (handle '' escape) + if (ch == '\'') { + i++; + while (i < len) { + if (queryText.charAt(i) == '\'') { + if (i + 1 < len && queryText.charAt(i + 1) == '\'') { + i += 2; // escaped quote - skip both + continue; + } + break; // end of string literal + } + i++; + } + continue; + } + // Skip double-quoted identifiers (e.g. c["WHERE"]) + if (ch == '"') { + i++; + while (i < len && queryText.charAt(i) != '"') { + i++; + } + continue; + } + char upperCh = queryTextUpper.charAt(i); + if (upperCh == '(') { + depth++; + } else if (upperCh == ')') { + depth--; + } else if (depth == 0 && upperCh == keywordUpper.charAt(0) + && queryTextUpper.startsWith(keywordUpper, i) + && (i == 0 || !isIdentifierChar(queryTextUpper.charAt(i - 1))) + && (i + keyLen >= queryTextUpper.length() || !isIdentifierChar(queryTextUpper.charAt(i + keyLen)))) { + return i; + } + } + return -1; + } + + /** + * Returns true if the character can appear in a SQL identifier or property access, + * meaning it should NOT be treated as a word boundary for keyword matching. + * Covers letters, digits, underscore, dot (property access), bracket (bracket notation), + * and dollar sign (system properties). + */ + private static boolean isIdentifierChar(char ch) { + return Character.isLetterOrDigit(ch) || ch == '_' || ch == '.' || ch == '[' || ch == '$'; + } + + /** + * Finds the index of the top-level WHERE keyword in the query text, + * ignoring WHERE that appears inside parentheses (subqueries). + */ + public static int findTopLevelWhereIndex(String queryTextUpper) { + return findTopLevelKeywordIndex(queryTextUpper, "WHERE"); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 11121bca033e..2773ecdf6cf6 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -118,10 +118,12 @@ import java.util.Arrays; import java.util.Collection; import java.util.Collections; +import java.util.Comparator; import java.util.EnumSet; import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; +import java.util.LinkedHashMap; import java.util.List; import java.util.Locale; import java.util.Map; @@ -4365,13 +4367,717 @@ private Mono>> readMany( ); } + @Override + public Flux> readManyByPartitionKeys( + List partitionKeys, + SqlQuerySpec customQuery, + String collectionLink, + QueryFeedOperationState state, + int maxConcurrentBatchPrefetch, + Class klass) { + + checkNotNull(partitionKeys, "Argument 'partitionKeys' must not be null."); + checkArgument(!partitionKeys.isEmpty(), "Argument 'partitionKeys' must not be empty."); + checkArgument(maxConcurrentBatchPrefetch >= 1, + "Argument 'maxConcurrentBatchPrefetch' must be greater than or equal to 1."); + + final ScopedDiagnosticsFactory diagnosticsFactory = new ScopedDiagnosticsFactory(this, true); + state.registerDiagnosticsFactory( + () -> {}, // we never want to reset in readManyByPartitionKeys + (ctx) -> diagnosticsFactory.merge(ctx) + ); + + StaleResourceRetryPolicy staleResourceRetryPolicy = new StaleResourceRetryPolicy( + this.collectionCache, + null, + collectionLink, + queryOptionsAccessor().getProperties(state.getQueryOptions()), + queryOptionsAccessor().getHeaders(state.getQueryOptions()), + this.sessionContainer, + diagnosticsFactory, + ResourceType.Document + ); + + return ObservableHelper + .fluxInlineIfPossibleAsObs( + () -> readManyByPartitionKeys( + partitionKeys, customQuery, collectionLink, state, diagnosticsFactory, maxConcurrentBatchPrefetch, klass), + staleResourceRetryPolicy + ) + .onErrorMap(throwable -> { + if (throwable instanceof CosmosException) { + CosmosException cosmosException = (CosmosException) throwable; + CosmosDiagnostics diagnostics = cosmosException.getDiagnostics(); + if (diagnostics != null) { + state.mergeDiagnosticsContext(); + CosmosDiagnosticsContext ctx = state.getDiagnosticsContextSnapshot(); + if (ctx != null) { + ctxAccessor().recordOperation( + ctx, + cosmosException.getStatusCode(), + cosmosException.getSubStatusCode(), + 0, + cosmosException.getRequestCharge(), + diagnostics, + throwable + ); + diagAccessor() + .setDiagnosticsContext( + diagnostics, + state.getDiagnosticsContextSnapshot()); + } + } + + return cosmosException; + } + + return throwable; + }); + } + + private Flux> readManyByPartitionKeys( + List partitionKeys, + SqlQuerySpec customQuery, + String collectionLink, + QueryFeedOperationState state, + ScopedDiagnosticsFactory diagnosticsFactory, + int maxConcurrentBatchPrefetch, + Class klass) { + + String requestContinuation = state.getRequestContinuation(); + + String resourceLink = parentResourceLinkToQueryLink(collectionLink, ResourceType.Document); + RxDocumentServiceRequest request = RxDocumentServiceRequest.create(diagnosticsFactory, + OperationType.Query, + ResourceType.Document, + collectionLink, null + ); + + Mono> collectionObs = + collectionCache.resolveCollectionAsync(null, request); + + return collectionObs + .flatMapMany(documentCollectionResourceResponse -> { + final DocumentCollection collection = documentCollectionResourceResponse.v; + if (collection == null) { + return Flux.error(new IllegalStateException("Collection cannot be null")); + } + + final PartitionKeyDefinition pkDefinition = collection.getPartitionKey(); + + final String collectionRid = collection.getResourceId(); + final String queryHash = ReadManyByPartitionKeyContinuationToken.computeQueryHash(customQuery); + final List normalizedPartitionKeys = + normalizePartitionKeys(partitionKeys, pkDefinition); + final String partitionKeySetHash = ReadManyByPartitionKeyContinuationToken.computePartitionKeySetHash( + normalizedPartitionKeys.stream() + .map(normalizedPk -> normalizedPk.effectivePartitionKeyString) + .collect(Collectors.toList())); + + // When a continuation token is present, skip the routing map lookup and batch + // construction - the token already contains the batch definitions. Only resolve + // the routing map and build batches on the very first call (no continuation). + if (requestContinuation != null) { + // Continuation path: validate collection/query, then resume from token. + // The routing map is still required because the persisted token only carries + // batch FILTER ranges; the FeedRange used at execution time (for low-RU + // routing to a single physical partition) is rederived from the current + // routing map per batch via resolvePartitionScopeFromBatchFilter(). + ReadManyByPartitionKeyContinuationToken parsedContinuation = + ReadManyByPartitionKeyContinuationToken.deserialize(requestContinuation); + if (!collectionRid.equals(parsedContinuation.getCollectionRid())) { + return Flux.error(new IllegalArgumentException( + "Continuation token was created for a different collection (rid mismatch). " + + "Expected: " + collectionRid + ", token has: " + parsedContinuation.getCollectionRid())); + } + if (!queryHash.equals(parsedContinuation.getQueryHash())) { + return Flux.error(new IllegalArgumentException( + "Continuation token was created with a different query (hash mismatch). " + + "The same query must be used when resuming from a continuation token.")); + } + if (!partitionKeySetHash.equals(parsedContinuation.getPartitionKeySetHash())) { + return Flux.error(new IllegalArgumentException( + "Continuation token was created with a different partition-key set (hash mismatch). " + + "The same normalized set of partition key values must be used when resuming.")); + } + + Mono> resumeRoutingMapMono = partitionKeyRangeCache + .tryLookupAsync( + BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), + collection.getResourceId(), + null, + null); + + return resumeRoutingMapMono.flatMapMany(resumeRoutingMapHolder -> { + CollectionRoutingMap resumeRoutingMap = resumeRoutingMapHolder.v; + if (resumeRoutingMap == null) { + return Flux.error(new IllegalStateException( + "Failed to get routing map for readManyByPartitionKeys continuation.")); + } + return buildSequentialFluxFromContinuation( + parsedContinuation, normalizedPartitionKeys, customQuery, pkDefinition, + resumeRoutingMap, resourceLink, state, diagnosticsFactory, klass, + collectionRid, queryHash, partitionKeySetHash, maxConcurrentBatchPrefetch); + }); + } + + // First-call path: validate custom query, resolve routing map, build batches + Mono queryValidationMono; + if (customQuery != null) { + queryValidationMono = validateCustomQueryForReadManyByPartitionKeys( + customQuery, resourceLink, state.getQueryOptions()); + } else { + queryValidationMono = Mono.empty(); + } + + Mono> valueHolderMono = partitionKeyRangeCache + .tryLookupAsync( + BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), + collection.getResourceId(), + null, + null); + + return valueHolderMono + .delayUntil(ignored -> queryValidationMono) + .flatMapMany(routingMapHolder -> { + CollectionRoutingMap routingMap = routingMapHolder.v; + if (routingMap == null) { + return Flux.error(new IllegalStateException("Failed to get routing map.")); + } + + return buildSequentialFluxFromScratch( + normalizedPartitionKeys, customQuery, pkDefinition, routingMap, + resourceLink, state, diagnosticsFactory, klass, + collectionRid, queryHash, partitionKeySetHash, maxConcurrentBatchPrefetch); + }); + }); + } + + /** + * Normalizes the input partition key list into a deterministic, set-based representation. + * Duplicates are removed by effective partition key string and the remaining keys are sorted + * lexicographically by EPK so batching and continuation-token hashes stay stable. + */ + static List normalizePartitionKeys( + List partitionKeys, + PartitionKeyDefinition pkDefinition) { + + Map normalizedByEpk = new HashMap<>(); + + for (PartitionKey pk : partitionKeys) { + PartitionKeyInternal pkInternal = BridgeInternal.getPartitionKeyInternal(pk); + if (pkDefinition.getKind() == PartitionKind.MULTI_HASH && pkInternal.getComponents() == null) { + throw new IllegalArgumentException( + "PartitionKey.NONE is not supported for multi-path partition keys in readManyByPartitionKeys."); + } + + PartitionKeyInternal effectivePkInternal = pkInternal.getComponents() == null + ? PartitionKeyInternal.UndefinedPartitionKey + : pkInternal; + String effectivePartitionKeyString = PartitionKeyInternalHelper + .getEffectivePartitionKeyString(effectivePkInternal, pkDefinition); + + normalizedByEpk.putIfAbsent( + effectivePartitionKeyString, + new NormalizedPartitionKey(pk, effectivePkInternal, effectivePartitionKeyString)); + } + + List normalizedPartitionKeys = new ArrayList<>(normalizedByEpk.values()); + normalizedPartitionKeys.sort(Comparator.comparing(normalizedPk -> normalizedPk.effectivePartitionKeyString)); + + if (pkDefinition.getKind() != PartitionKind.MULTI_HASH || normalizedPartitionKeys.size() < 2) { + return normalizedPartitionKeys; + } + + List collapsedPartitionKeys = new ArrayList<>(normalizedPartitionKeys.size()); + for (NormalizedPartitionKey candidate : normalizedPartitionKeys) { + if (!collapsedPartitionKeys.isEmpty()) { + NormalizedPartitionKey previous = collapsedPartitionKeys.get(collapsedPartitionKeys.size() - 1); + if (previous.effectivePkInternal.contains(candidate.effectivePkInternal)) { + // The previous PK is a prefix of the current PK, so the current PK is fully + // subsumed by the previous read scope and would only cause duplicate results + // if it ended up in a separate batch. + continue; + } + } + + collapsedPartitionKeys.add(candidate); + } + + return collapsedPartitionKeys; + } + + /** + * Builds the sequential Flux of FeedResponse pages for readManyByPartitionKeys when starting + * from scratch (no continuation token). Groups PKs by physical partition, creates batches, + * sorts by EPK, and executes sequentially. + */ + private Flux> buildSequentialFluxFromScratch( + List normalizedPartitionKeys, + SqlQuerySpec customQuery, + PartitionKeyDefinition pkDefinition, + CollectionRoutingMap routingMap, + String resourceLink, + QueryFeedOperationState state, + ScopedDiagnosticsFactory diagnosticsFactory, + Class klass, + String collectionRid, + String queryHash, + String partitionKeySetHash, + int maxConcurrentBatchPrefetch) { + + Map> partitionRangePkMap = + groupPartitionKeysByPhysicalPartition(normalizedPartitionKeys, pkDefinition, routingMap); + + List partitionKeySelectors = PartitionKeyQueryHelper.createPkSelectors(pkDefinition); + + String baseQueryText; + List baseParameters; + if (customQuery != null) { + baseQueryText = customQuery.getQueryText(); + baseParameters = customQuery.getParameters() != null + ? new ArrayList<>(customQuery.getParameters()) + : new ArrayList<>(); + } else { + baseQueryText = "SELECT * FROM c"; + baseParameters = new ArrayList<>(); + } + + int maxPksPerPartitionQuery = Configs.getReadManyByPkMaxBatchSize(); + + List allBatches = new ArrayList<>(); + + for (Map.Entry> entry : partitionRangePkMap.entrySet()) { + PartitionKeyRange pkRange = entry.getKey(); + Range partitionScope = pkRange.toRange(); + List allPks = entry.getValue(); + + // The per-range list already preserves the globally normalized EPK order, + // so no additional in-partition sort is required here. + for (int i = 0; i < allPks.size(); i += maxPksPerPartitionQuery) { + int batchEnd = Math.min(i + maxPksPerPartitionQuery, allPks.size()); + List batch = allPks.subList(i, batchEnd) + .stream() + .map(normalizedPk -> normalizedPk.partitionKey) + .collect(Collectors.toList()); + + // batchFilter is [epk(first), maxExclusive) where maxExclusive is: + // - epk of the next PK after this batch (first PK of the next batch), or + // - partitionScope.getMax() if this is the last batch in the partition + String batchMinInclusive = allPks.get(i).effectivePartitionKeyString; + String batchMaxExclusive = batchEnd < allPks.size() + ? allPks.get(batchEnd).effectivePartitionKeyString + : partitionScope.getMax(); + Range batchFilter = new Range<>(batchMinInclusive, batchMaxExclusive, true, false); + + SqlQuerySpec querySpec = ReadManyByPartitionKeyQueryHelper + .createReadManyByPkQuerySpec( + baseQueryText, baseParameters, batch, + partitionKeySelectors, pkDefinition); + allBatches.add(new BatchDescriptor(partitionScope, batchFilter, querySpec)); + } + } + + if (allBatches.isEmpty()) { + return Flux.empty(); + } + + // Sort batches by batchFilter EPK range for deterministic sequential processing. + // Tie-break on maxExclusive so that prefix HPK batches sharing the same minInclusive + // (the prefix EPK) always have a stable order regardless of routing-map traversal order. + allBatches.sort(Comparator.comparing((BatchDescriptor bd) -> bd.batchFilter.getMin()) + .thenComparing(bd -> bd.batchFilter.getMax())); + + return buildSequentialBatchFlux( + allBatches, null, + resourceLink, state, diagnosticsFactory, klass, + collectionRid, queryHash, partitionKeySetHash, maxConcurrentBatchPrefetch); + } + + /** + * Builds the sequential Flux of FeedResponse pages for readManyByPartitionKeys when resuming + * from a continuation token. Reconstructs batches from the token's batch definitions + * without needing the routing map. + */ + private Flux> buildSequentialFluxFromContinuation( + ReadManyByPartitionKeyContinuationToken parsedContinuation, + List normalizedPartitionKeys, + SqlQuerySpec customQuery, + PartitionKeyDefinition pkDefinition, + CollectionRoutingMap routingMap, + String resourceLink, + QueryFeedOperationState state, + ScopedDiagnosticsFactory diagnosticsFactory, + Class klass, + String collectionRid, + String queryHash, + String partitionKeySetHash, + int maxConcurrentBatchPrefetch) { + + List partitionKeySelectors = PartitionKeyQueryHelper.createPkSelectors(pkDefinition); + + String baseQueryText; + List baseParameters; + if (customQuery != null) { + baseQueryText = customQuery.getQueryText(); + baseParameters = customQuery.getParameters() != null + ? new ArrayList<>(customQuery.getParameters()) + : new ArrayList<>(); + } else { + baseQueryText = "SELECT * FROM c"; + baseParameters = new ArrayList<>(); + } + + ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBatchDef = + parsedContinuation.getCurrentBatch(); + List remainingBatchDefs = + parsedContinuation.getRemainingBatches(); + String initialBackendContinuation = parsedContinuation.getBackendContinuation(); + + List allBatchDefs = new ArrayList<>(); + allBatchDefs.add(currentBatchDef); + allBatchDefs.addAll(remainingBatchDefs); + + List allBatches = new ArrayList<>(); + + for (ReadManyByPartitionKeyContinuationToken.BatchDefinition batchDef : allBatchDefs) { + Range batchFilter = batchDef.getBatchFilter(); + + List batchPks = filterPartitionKeysByEpkRange( + normalizedPartitionKeys, batchFilter); + + if (batchPks.isEmpty()) { + continue; + } + + // Rederive the routing scope from the CURRENT routing map. After a partition + // split this naturally yields the new (potentially smaller) physical-partition + // boundaries that exactly cover this batch filter, keeping query RUs minimal. + // If the cache is briefly stale right after a split, the SDK's stale-resource + // retry will refresh it and rerun, so any RU-cost elevation is bounded. + Range partitionScope = resolvePartitionScopeFromBatchFilter(batchFilter, routingMap); + + SqlQuerySpec querySpec = ReadManyByPartitionKeyQueryHelper + .createReadManyByPkQuerySpec( + baseQueryText, baseParameters, batchPks, + partitionKeySelectors, pkDefinition); + + allBatches.add(new BatchDescriptor(partitionScope, batchFilter, querySpec)); + } + + if (allBatches.isEmpty()) { + return Flux.empty(); + } + + return buildSequentialBatchFlux( + allBatches, initialBackendContinuation, + resourceLink, state, diagnosticsFactory, klass, + collectionRid, queryHash, partitionKeySetHash, maxConcurrentBatchPrefetch); + } + + /** + * Filters partition keys to those whose EPK falls within the given range. + */ + private List filterPartitionKeysByEpkRange( + List normalizedPartitionKeys, + Range epkRange) { + + List result = new ArrayList<>(); + for (NormalizedPartitionKey normalizedPk : normalizedPartitionKeys) { + if (epkRange.contains(normalizedPk.effectivePartitionKeyString)) { + result.add(normalizedPk.partitionKey); + } + } + return result; + } + + /** + * Builds the sequential Flux that executes batches one at a time, stamping each + * FeedResponse with the composite continuation token. + */ + private Flux> buildSequentialBatchFlux( + List allBatches, + String initialBackendContinuation, + String resourceLink, + QueryFeedOperationState state, + ScopedDiagnosticsFactory diagnosticsFactory, + Class klass, + String collectionRid, + String queryHash, + String partitionKeySetHash, + int maxConcurrentBatchPrefetch) { + + List>> sequentialFluxes = new ArrayList<>(); + for (int i = 0; i < allBatches.size(); i++) { + final int batchIndex = i; + final BatchDescriptor bd = allBatches.get(i); + final String backendContinuation = (i == 0) ? initialBackendContinuation : null; + + final List remainingAfterThis = new ArrayList<>(); + for (int j = batchIndex + 1; j < allBatches.size(); j++) { + BatchDescriptor remaining = allBatches.get(j); + // Only the batch FILTER is persisted; the routing scope is rederived per batch + // at resume time from the live routing-map cache. + remainingAfterThis.add( + new ReadManyByPartitionKeyContinuationToken.BatchDefinition(remaining.batchFilter)); + } + + CosmosQueryRequestOptions batchQueryOptions = queryOptionsAccessor() + .clone(state.getQueryOptions()); + queryOptionsAccessor().disallowQueryPlanRetrieval(batchQueryOptions); + + batchQueryOptions.setFeedRange(new FeedRangeEpkImpl(bd.partitionScope)); + + ModelBridgeInternal.setQueryRequestOptionsContinuationToken( + batchQueryOptions, backendContinuation); + + Flux> batchFlux = createQueryInternal( + diagnosticsFactory, + resourceLink, + bd.querySpec, + batchQueryOptions, + klass, + ResourceType.Document, + documentQueryClientImpl(RxDocumentClientImpl.this, getOperationContextAndListenerTuple(batchQueryOptions)), + UUIDs.nonBlockingRandomUUID(), + new AtomicBoolean(false)); + + final ReadManyByPartitionKeyContinuationToken.BatchDefinition currentBatchDef = + new ReadManyByPartitionKeyContinuationToken.BatchDefinition(bd.batchFilter); + + Flux> stampedFlux = batchFlux.map(feedResponse -> { + String backendCont = feedResponse.getContinuationToken(); + boolean isLastBatch = remainingAfterThis.isEmpty(); + boolean batchExhausted = (backendCont == null); + + if (isLastBatch && batchExhausted) { + ModelBridgeInternal.setFeedResponseContinuationToken(null, feedResponse); + } else if (batchExhausted) { + ReadManyByPartitionKeyContinuationToken.BatchDefinition nextBatchDef = + remainingAfterThis.get(0); + List remaining = + remainingAfterThis.size() > 1 + ? remainingAfterThis.subList(1, remainingAfterThis.size()) + : Collections.emptyList(); + ReadManyByPartitionKeyContinuationToken compositeContinuation = + new ReadManyByPartitionKeyContinuationToken( + remaining, nextBatchDef, null, collectionRid, queryHash, partitionKeySetHash); + ModelBridgeInternal.setFeedResponseContinuationToken( + compositeContinuation.serialize(), feedResponse); + } else { + ReadManyByPartitionKeyContinuationToken compositeContinuation = + new ReadManyByPartitionKeyContinuationToken( + remainingAfterThis, currentBatchDef, backendCont, + collectionRid, queryHash, partitionKeySetHash); + ModelBridgeInternal.setFeedResponseContinuationToken( + compositeContinuation.serialize(), feedResponse); + } + return feedResponse; + }); + + sequentialFluxes.add(stampedFlux); + } + + int fluxConcurrency = Math.max(1, Math.min(maxConcurrentBatchPrefetch, sequentialFluxes.size())); + return Flux.mergeSequential(sequentialFluxes, fluxConcurrency, 1); + } + + /** + * Computes the EPK routing scope for a batch from its filter range and the current + * collection routing map. Looks up all partition-key ranges that overlap the batch + * filter and returns {@code [min(minInclusive), max(maxExclusive))} across them so + * the FeedRange set on the per-batch query options exactly aligns with one or more + * physical-partition boundaries (yielding minimal query RU cost). + *

+ * If a partition split happened between the time the continuation token was emitted + * and resume time, the rederived scope reflects the post-split boundaries; the SDK's + * stale-resource retry refreshes the cache promptly if the local cache snapshot was + * still pre-split. Either way no stale boundary is encoded in the token itself. + * + * @throws IllegalStateException if the routing map has no overlap for the batch filter. + */ + private static Range resolvePartitionScopeFromBatchFilter( + Range batchFilter, + CollectionRoutingMap routingMap) { + + List overlapping = routingMap.getOverlappingRanges(batchFilter); + if (overlapping == null || overlapping.isEmpty()) { + throw new IllegalStateException( + "Routing map returned no overlapping partition key ranges for batch filter " + + batchFilter + "."); + } + + String minInclusive = overlapping.get(0).getMinInclusive(); + String maxExclusive = overlapping.get(0).getMaxExclusive(); + for (int i = 1; i < overlapping.size(); i++) { + PartitionKeyRange r = overlapping.get(i); + if (r.getMinInclusive().compareTo(minInclusive) < 0) { + minInclusive = r.getMinInclusive(); + } + if (r.getMaxExclusive().compareTo(maxExclusive) > 0) { + maxExclusive = r.getMaxExclusive(); + } + } + + return new Range<>(minInclusive, maxExclusive, true, false); + } + + static final class NormalizedPartitionKey { + final PartitionKey partitionKey; + final PartitionKeyInternal effectivePkInternal; + final String effectivePartitionKeyString; + + private NormalizedPartitionKey( + PartitionKey partitionKey, + PartitionKeyInternal effectivePkInternal, + String effectivePartitionKeyString) { + + this.partitionKey = partitionKey; + this.effectivePkInternal = effectivePkInternal; + this.effectivePartitionKeyString = effectivePartitionKeyString; + } + } + + /** + * Descriptor for a single batch during execution of readManyByPartitionKeys. + *

+ * Each batch carries two EPK ranges: + *

    + *
  • {@code partitionScope} — the physical partition's EPK range, used as the + * FeedRange on CosmosQueryRequestOptions so the backend query request is sent + * to the matching physical partition with minimal RU overhead. This value is + * not persisted in the continuation token; it is derived per + * batch from the current routing map (via + * {@link #resolvePartitionScopeFromBatchFilter}) so post-split refreshes are + * reflected automatically.
    • + *
  • {@code batchFilter} — the EPK sub-range covering only the PKs in this batch. + * This IS persisted in the continuation token. When a physical partition has + * more PKs than maxBatchSize, multiple batches share the same partitionScope + * at execution time but have distinct batchFilter ranges. Used to reconstruct + * the correct PK set per batch when resuming.
    • + *
+ */ + private static final class BatchDescriptor { + final Range partitionScope; + final Range batchFilter; + final SqlQuerySpec querySpec; + + BatchDescriptor(Range partitionScope, Range batchFilter, SqlQuerySpec querySpec) { + this.partitionScope = partitionScope; + this.batchFilter = batchFilter; + this.querySpec = querySpec; + } + } + + private Mono validateCustomQueryForReadManyByPartitionKeys( + SqlQuerySpec customQuery, + String resourceLink, + CosmosQueryRequestOptions queryRequestOptions) { + + IDocumentQueryClient queryClient = documentQueryClientImpl( + RxDocumentClientImpl.this, getOperationContextAndListenerTuple(queryRequestOptions)); + + return DocumentQueryExecutionContextFactory + .fetchQueryPlanForValidation( + this, + queryClient, + customQuery, + resourceLink, + queryRequestOptions, + Configs.isQueryPlanCachingEnabled(), + this.getQueryPlanCache()) + .doOnNext(RxDocumentClientImpl::validateQueryPlanForReadManyByPartitionKeys) + .then(); + } + + static void validateQueryPlanForReadManyByPartitionKeys(PartitionedQueryExecutionInfo queryPlan) { + if (queryPlan.hasHybridSearchQueryInfo()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain hybrid/vector/full-text search."); + } + + QueryInfo queryInfo = queryPlan.getQueryInfo(); + if (queryInfo == null) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key is not supported because query plan details are unavailable."); + } + + if (queryInfo.hasGroupBy()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain GROUP BY."); + } + if (queryInfo.hasAggregates()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain aggregates."); + } + if (queryInfo.hasOrderBy()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain ORDER BY."); + } + if (queryInfo.hasDistinct()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain DISTINCT."); + } + if (queryInfo.hasDCount()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain DCOUNT."); + } + if (queryInfo.hasOffset()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain OFFSET."); + } + if (queryInfo.hasLimit()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain LIMIT."); + } + if (queryInfo.hasTop()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain TOP."); + } + if (queryInfo.hasNonStreamingOrderBy()) { + throw new IllegalArgumentException( + "Custom query for readMany by partition key must not contain non-streaming ORDER BY."); + } + } + + private Map> groupPartitionKeysByPhysicalPartition( + List normalizedPartitionKeys, + PartitionKeyDefinition pkDefinition, + CollectionRoutingMap routingMap) { + + Map> partitionRangePkMap = new LinkedHashMap<>(); + + for (NormalizedPartitionKey normalizedPk : normalizedPartitionKeys) { + int componentCount = normalizedPk.effectivePkInternal.getComponents().size(); + int definedPathCount = pkDefinition.getPaths().size(); + + List targetRanges; + + if (pkDefinition.getKind() == PartitionKind.MULTI_HASH && componentCount < definedPathCount) { + Range epkRange = PartitionKeyInternalHelper.getEPKRangeForPrefixPartitionKey( + normalizedPk.effectivePkInternal, pkDefinition); + targetRanges = routingMap.getOverlappingRanges(epkRange); + } else { + PartitionKeyRange range = routingMap.getRangeByEffectivePartitionKey( + normalizedPk.effectivePartitionKeyString); + targetRanges = Collections.singletonList(range); + } + + for (PartitionKeyRange range : targetRanges) { + partitionRangePkMap.computeIfAbsent(range, k -> new ArrayList<>()).add(normalizedPk); + } + } + + return partitionRangePkMap; + } + private Map getRangeQueryMap( Map> partitionRangeItemKeyMap, PartitionKeyDefinition partitionKeyDefinition) { //TODO: Optimise this to include all types of partitionkeydefinitions. ex: c["prop1./ab"]["key1"] Map rangeQueryMap = new HashMap<>(); - List partitionKeySelectors = createPkSelectors(partitionKeyDefinition); + List partitionKeySelectors = PartitionKeyQueryHelper.createPkSelectors(partitionKeyDefinition); for(Map.Entry> entry: partitionRangeItemKeyMap.entrySet()) { SqlQuerySpec sqlQuerySpec; @@ -4465,15 +5171,6 @@ private SqlQuerySpec createReadManyQuerySpec( return new SqlQuerySpec(queryStringBuilder.toString(), parameters); } - private List createPkSelectors(PartitionKeyDefinition partitionKeyDefinition) { - return partitionKeyDefinition.getPaths() - .stream() - .map(pathPart -> StringUtils.substring(pathPart, 1)) // skip starting / - .map(pathPart -> StringUtils.replace(pathPart, "\"", "\\")) // escape quote - .map(part -> "[\"" + part + "\"]") - .collect(Collectors.toList()); - } - private Flux> queryForReadMany( ScopedDiagnosticsFactory diagnosticsFactory, String parentResourceLink, @@ -4995,7 +5692,7 @@ public Flux> readAllDocuments( } PartitionKeyDefinition pkDefinition = collection.getPartitionKey(); - List partitionKeySelectors = createPkSelectors(pkDefinition); + List partitionKeySelectors = PartitionKeyQueryHelper.createPkSelectors(pkDefinition); SqlQuerySpec querySpec = createLogicalPartitionScanQuerySpec(partitionKey, partitionKeySelectors); String resourceLink = parentResourceLinkToQueryLink(collectionLink, ResourceType.Document); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java index e62d8ed3d754..11528ed4f7ea 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java @@ -107,7 +107,6 @@ private static Mono getPartitionKeyRangesAn } Instant startTime = Instant.now(); - Mono queryExecutionInfoMono; if (queryRequestOptionsAccessor() .isQueryPlanRetrievalDisallowed(cosmosQueryRequestOptions)) { @@ -122,37 +121,53 @@ private static Mono getPartitionKeyRangesAn endTime); } + return fetchQueryPlan( + diagnosticsClientContext, + client, + query, + resourceLink, + cosmosQueryRequestOptions, + queryPlanCachingEnabled, + queryPlanCache) + .flatMap( + partitionedQueryExecutionInfo -> { + + Instant endTime = Instant.now(); + + return getTargetRangesFromQueryPlan(cosmosQueryRequestOptions, collection, queryExecutionContext, + partitionedQueryExecutionInfo, startTime, endTime); + }); + } + + private static Mono fetchQueryPlan( + DiagnosticsClientContext diagnosticsClientContext, + IDocumentQueryClient client, + SqlQuerySpec query, + String resourceLink, + CosmosQueryRequestOptions cosmosQueryRequestOptions, + boolean queryPlanCachingEnabled, + Map queryPlanCache) { + if (queryPlanCachingEnabled && - isScopedToSinglePartition(cosmosQueryRequestOptions) && - queryPlanCache.containsKey(query.getQueryText())) { - Instant endTime = Instant.now(); // endTime for query plan diagnostics + isScopedToSinglePartition(cosmosQueryRequestOptions) && + queryPlanCache.containsKey(query.getQueryText())) { PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = queryPlanCache.get(query.getQueryText()); if (partitionedQueryExecutionInfo != null) { logger.debug("Skipping query plan round trip by using the cached plan"); - return getTargetRangesFromQueryPlan(cosmosQueryRequestOptions, collection, queryExecutionContext, - partitionedQueryExecutionInfo, startTime, endTime); + return Mono.just(partitionedQueryExecutionInfo); } } - queryExecutionInfoMono = - QueryPlanRetriever.getQueryPlanThroughGatewayAsync( - diagnosticsClientContext, - client, - query, - resourceLink, - cosmosQueryRequestOptions); - - return queryExecutionInfoMono.flatMap( - partitionedQueryExecutionInfo -> { - - Instant endTime = Instant.now(); - + return QueryPlanRetriever.getQueryPlanThroughGatewayAsync( + diagnosticsClientContext, + client, + query, + resourceLink, + cosmosQueryRequestOptions) + .doOnNext(partitionedQueryExecutionInfo -> { if (queryPlanCachingEnabled && isScopedToSinglePartition(cosmosQueryRequestOptions)) { tryCacheQueryPlan(query, partitionedQueryExecutionInfo, queryPlanCache); } - - return getTargetRangesFromQueryPlan(cosmosQueryRequestOptions, collection, queryExecutionContext, - partitionedQueryExecutionInfo, startTime, endTime); }); } @@ -232,11 +247,11 @@ private static Mono getTargetRangesFromEmpt "Query plan retrieval must not be suppressed when not using FeedRanges"); } - QueryInfo queryInfo = QueryInfo.EMPTY; - queryInfo.setQueryPlanDiagnosticsContext( - new QueryInfo.QueryPlanDiagnosticsContext( - planFetchStartTime, - planFetchEndTime)); + // Do NOT use QueryInfo.EMPTY here — setQueryPlanDiagnosticsContext would mutate + // the shared static singleton. Instead, capture the diagnostics context and create + // a fresh per-request QueryInfo inside the reactive chain. + QueryInfo.QueryPlanDiagnosticsContext queryPlanDiagnosticsContext = + new QueryInfo.QueryPlanDiagnosticsContext(planFetchStartTime, planFetchEndTime); FeedRange userProvidedFeedRange = cosmosQueryRequestOptions.getFeedRange(); Mono> targetRange = queryExecutionContext @@ -249,7 +264,14 @@ private static Mono getTargetRangesFromEmpt ).map(pkRanges -> pkRanges.stream().map(PartitionKeyRange::toRange).collect(Collectors.toList())); return Mono.zip(targetRange, allRanges) - .map(tuple -> new PartitionKeyRangesAndQueryInfos(queryInfo, null, Collections.singletonList(tuple.getT1()), tuple.getT2())); + .map(tuple -> { + // Set diagnostics on a per-request basis inside the reactive chain, + // not on the shared EMPTY singleton. Create a fresh QueryInfo only here. + QueryInfo perRequestQueryInfo = new QueryInfo(); + perRequestQueryInfo.setQueryPlanDiagnosticsContext(queryPlanDiagnosticsContext); + return new PartitionKeyRangesAndQueryInfos( + perRequestQueryInfo, null, Collections.singletonList(tuple.getT1()), tuple.getT2()); + }); } synchronized private static void tryCacheQueryPlan( @@ -318,6 +340,25 @@ private static List getFeedRangeEpks(List> range return feedRanges; } + public static Mono fetchQueryPlanForValidation( + DiagnosticsClientContext diagnosticsClientContext, + IDocumentQueryClient queryClient, + SqlQuerySpec sqlQuerySpec, + String resourceLink, + CosmosQueryRequestOptions queryRequestOptions, + boolean queryPlanCachingEnabled, + Map queryPlanCache) { + + return fetchQueryPlan( + diagnosticsClientContext, + queryClient, + sqlQuerySpec, + resourceLink, + queryRequestOptions, + queryPlanCachingEnabled, + queryPlanCache); + } + public static Flux> createDocumentQueryExecutionContextAsync( DiagnosticsClientContext diagnosticsClientContext, IDocumentQueryClient client, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java new file mode 100644 index 000000000000..8b46dcc9c987 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosReadManyByPartitionKeysRequestOptions.java @@ -0,0 +1,402 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.models; + +import com.azure.cosmos.CosmosDiagnosticsThresholds; +import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfig; +import com.azure.cosmos.CosmosItemSerializer; +import com.azure.cosmos.ReadConsistencyStrategy; +import com.azure.cosmos.implementation.CosmosQueryRequestOptionsBase; +import com.azure.cosmos.implementation.CosmosReadManyByPartitionKeysRequestOptionsImpl; +import com.azure.cosmos.implementation.ImplementationBridgeHelpers; +import com.azure.cosmos.util.Beta; + +import java.time.Duration; +import java.util.List; +import java.util.Set; + +/** + * Specifies the options associated with the {@code readManyByPartitionKeys} operation + * in the Azure Cosmos DB database service. + *

+ * This is distinct from {@link CosmosReadManyRequestOptions} (used by the + * {@code readMany(List<CosmosItemIdentity>)} API). It exposes only the knobs that are + * applicable to {@code readManyByPartitionKeys} — for example, properties that influence + * query parallelism inside a single physical partition or a feed range filter are + * intentionally not exposed because the operation is fully managed by the SDK. + */ +public final class CosmosReadManyByPartitionKeysRequestOptions { + private final CosmosReadManyByPartitionKeysRequestOptionsImpl actualRequestOptions; + + /** + * Instantiates a new readManyByPartitionKeys request options. + */ + public CosmosReadManyByPartitionKeysRequestOptions() { + this.actualRequestOptions = new CosmosReadManyByPartitionKeysRequestOptionsImpl(); + } + + /** + * Copy constructor. + * + * @param options the options to copy. + */ + CosmosReadManyByPartitionKeysRequestOptions(CosmosReadManyByPartitionKeysRequestOptions options) { + this.actualRequestOptions = new CosmosReadManyByPartitionKeysRequestOptionsImpl(options.actualRequestOptions); + } + + /** + * Gets the composite continuation token used to resume a previous + * {@code readManyByPartitionKeys} invocation. + * + * @return the continuation token, or null if not set. + */ + public String getContinuationToken() { + return this.actualRequestOptions.getContinuationToken(); + } + + /** + * Sets the composite continuation token used to resume a previous + * {@code readManyByPartitionKeys} invocation. The token must have been returned by a prior + * invocation of {@code readManyByPartitionKeys} on the same container, with the same + * partition-key set and the same custom query. + * + * @param continuationToken the composite continuation token from a previous invocation. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeysRequestOptions setContinuationToken(String continuationToken) { + this.actualRequestOptions.setContinuationToken(continuationToken); + return this; + } + + /** + * Gets the maximum number of per-physical-partition batches whose first page is prefetched + * concurrently. This bounds the prefetch parallelism the SDK uses while sequentially + * draining batches. + * + * @return the max concurrent batch prefetch, or null if the SDK default is in effect. + */ + public Integer getMaxConcurrentBatchPrefetch() { + return this.actualRequestOptions.getMaxConcurrentBatchPrefetch(); + } + + /** + * Sets the maximum number of per-physical-partition batches whose first page is prefetched + * concurrently. The default is {@code Math.max(1, Math.min(Runtime.getRuntime().availableProcessors(), 8))}. + *

+ * Increase this to trade memory for lower end-to-end latency on wide containers; decrease it + * (e.g. to {@code 1}) when running in environments where a single task already saturates the + * network/CPU and additional prefetch only adds memory pressure. + * + * @param maxConcurrentBatchPrefetch the max concurrent batch prefetch (must be >= 1). + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + * @throws IllegalArgumentException if {@code maxConcurrentBatchPrefetch} is < 1. + */ + public CosmosReadManyByPartitionKeysRequestOptions setMaxConcurrentBatchPrefetch(int maxConcurrentBatchPrefetch) { + if (maxConcurrentBatchPrefetch < 1) { + throw new IllegalArgumentException( + "Argument 'maxConcurrentBatchPrefetch' must be greater than or equal to 1."); + } + this.actualRequestOptions.setMaxConcurrentBatchPrefetch(maxConcurrentBatchPrefetch); + return this; + } + + /** + * Gets the read consistency strategy for the request. + * + * @return the read consistency strategy. + */ + @Beta(value = Beta.SinceVersion.V4_69_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public ReadConsistencyStrategy getReadConsistencyStrategy() { + return this.actualRequestOptions.getReadConsistencyStrategy(); + } + + /** + * Sets the read consistency strategy required for the request. + * + * @param readConsistencyStrategy the read consistency strategy. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + */ + @Beta(value = Beta.SinceVersion.V4_69_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public CosmosReadManyByPartitionKeysRequestOptions setReadConsistencyStrategy( + ReadConsistencyStrategy readConsistencyStrategy) { + this.actualRequestOptions.setReadConsistencyStrategy(readConsistencyStrategy); + return this; + } + + /** + * Gets the session token for use with session consistency. + * + * @return the session token. + */ + public String getSessionToken() { + return this.actualRequestOptions.getSessionToken(); + } + + /** + * Sets the session token for use with session consistency. + * + * @param sessionToken the session token. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeysRequestOptions setSessionToken(String sessionToken) { + this.actualRequestOptions.setSessionToken(sessionToken); + return this; + } + + /** + * Sets the maximum size (in kilobytes) of the backend continuation token embedded inside the + * composite {@code readManyByPartitionKeys} continuation token. + *

+ * Note: this only constrains the per-batch backend continuation that the SDK wraps inside + * the public composite token; the public composite token itself is always larger because it + * also carries the remaining batch definitions, query hash, and partition-key-set hash. + * + * @param limitInKb backend continuation token size limit (must be >= 1). + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeysRequestOptions setMaxBackendContinuationTokenSizeInKb(int limitInKb) { + this.actualRequestOptions.setResponseContinuationTokenLimitInKb(limitInKb); + return this; + } + + /** + * Gets the maximum size (in kilobytes) of the backend continuation token embedded inside the + * composite {@code readManyByPartitionKeys} continuation token. Returns 0 if not set. + * + * @return the configured backend continuation token size limit, or 0 if not set. + */ + public int getMaxBackendContinuationTokenSizeInKb() { + return this.actualRequestOptions.getResponseContinuationTokenLimitInKb(); + } + + /** + * Gets the maximum number of items returned in a single page. + * + * @return the max item count, or null if not set (the SDK default applies). + */ + public Integer getMaxItemCount() { + return this.actualRequestOptions.getMaxItemCount(); + } + + /** + * Sets the maximum number of items returned in a single page. + * + * @param maxItemCount the maximum number of items per page. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeysRequestOptions setMaxItemCount(int maxItemCount) { + this.actualRequestOptions.setMaxItemCount(maxItemCount); + return this; + } + + /** + * Sets the {@link CosmosEndToEndOperationLatencyPolicyConfig} to be used for the request. + * + * @param cosmosEndToEndOperationLatencyPolicyConfig the {@link CosmosEndToEndOperationLatencyPolicyConfig} + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeysRequestOptions setCosmosEndToEndOperationLatencyPolicyConfig( + CosmosEndToEndOperationLatencyPolicyConfig cosmosEndToEndOperationLatencyPolicyConfig) { + this.actualRequestOptions + .setCosmosEndToEndOperationLatencyPolicyConfig(cosmosEndToEndOperationLatencyPolicyConfig); + return this; + } + + /** + * List of regions to be excluded for the request/retries. + * + * @param excludeRegions the regions to exclude + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} + */ + public CosmosReadManyByPartitionKeysRequestOptions setExcludedRegions(List excludeRegions) { + this.actualRequestOptions.setExcludedRegions(excludeRegions); + return this; + } + + /** + * Gets the list of regions to exclude for the request/retries. + * + * @return the list of excluded regions + */ + public List getExcludedRegions() { + return this.actualRequestOptions.getExcludedRegions(); + } + + /** + * Gets the option to enable populate query metrics. By default query metrics are enabled. + * + * @return whether query metrics are enabled + */ + public boolean isQueryMetricsEnabled() { + return this.actualRequestOptions.isQueryMetricsEnabled(); + } + + /** + * Sets the option to enable/disable query metrics. + * + * @param queryMetricsEnabled whether to enable or disable query metrics + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeysRequestOptions setQueryMetricsEnabled(boolean queryMetricsEnabled) { + this.actualRequestOptions.setQueryMetricsEnabled(queryMetricsEnabled); + return this; + } + + /** + * Gets the throughput control group name. + * + * @return the throughput control group name. + */ + public String getThroughputControlGroupName() { + return this.actualRequestOptions.getThroughputControlGroupName(); + } + + /** + * Sets the throughput control group name. + * + * @param throughputControlGroupName the throughput control group name. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeysRequestOptions setThroughputControlGroupName(String throughputControlGroupName) { + this.actualRequestOptions.setThroughputControlGroupName(throughputControlGroupName); + return this; + } + + /** + * Gets the dedicated gateway request options. + * + * @return the dedicated gateway request options. + */ + public DedicatedGatewayRequestOptions getDedicatedGatewayRequestOptions() { + return this.actualRequestOptions.getDedicatedGatewayRequestOptions(); + } + + /** + * Sets the dedicated gateway request options. + * + * @param dedicatedGatewayRequestOptions the dedicated gateway request options. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeysRequestOptions setDedicatedGatewayRequestOptions( + DedicatedGatewayRequestOptions dedicatedGatewayRequestOptions) { + this.actualRequestOptions.setDedicatedGatewayRequestOptions(dedicatedGatewayRequestOptions); + return this; + } + + /** + * Gets the latency threshold for diagnostics on tracer. + * + * @return the latency threshold for diagnostics on tracer. + */ + public Duration getThresholdForDiagnosticsOnTracer() { + return this.actualRequestOptions.getThresholdForDiagnosticsOnTracer(); + } + + /** + * Sets the latency threshold for diagnostics on tracer. + * + * @param thresholdForDiagnosticsOnTracer the latency threshold for diagnostics on tracer. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeysRequestOptions setThresholdForDiagnosticsOnTracer( + Duration thresholdForDiagnosticsOnTracer) { + this.actualRequestOptions.setThresholdForDiagnosticsOnTracer(thresholdForDiagnosticsOnTracer); + return this; + } + + /** + * Allows overriding the diagnostic thresholds for a specific operation. + * + * @param operationSpecificThresholds the diagnostic threshold override for this operation + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeysRequestOptions setDiagnosticsThresholds( + CosmosDiagnosticsThresholds operationSpecificThresholds) { + this.actualRequestOptions.setDiagnosticsThresholds(operationSpecificThresholds); + return this; + } + + /** + * Gets the diagnostic thresholds used as an override for a specific operation. + * + * @return the diagnostic thresholds for this operation. + */ + public CosmosDiagnosticsThresholds getDiagnosticsThresholds() { + return this.actualRequestOptions.getDiagnosticsThresholds(); + } + + /** + * Gets the custom item serializer defined for this instance of request options. + * + * @return the custom item serializer. + */ + public CosmosItemSerializer getCustomItemSerializer() { + return this.actualRequestOptions.getCustomItemSerializer(); + } + + /** + * Sets a custom item serializer to be used for this operation. + * + * @param customItemSerializer the custom item serializer for this operation. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeysRequestOptions setCustomItemSerializer( + CosmosItemSerializer customItemSerializer) { + this.actualRequestOptions.setCustomItemSerializer(customItemSerializer); + return this; + } + + /** + * Sets the custom keyword identifiers. + * + * @param keywordIdentifiers the custom keyword identifiers. + * @return the {@link CosmosReadManyByPartitionKeysRequestOptions} for fluent chaining. + */ + public CosmosReadManyByPartitionKeysRequestOptions setKeywordIdentifiers(Set keywordIdentifiers) { + this.actualRequestOptions.setKeywordIdentifiers(keywordIdentifiers); + return this; + } + + /** + * Gets the custom keyword identifiers. + * + * @return the custom keyword identifiers. + */ + public Set getKeywordIdentifiers() { + return this.actualRequestOptions.getKeywordIdentifiers(); + } + + CosmosQueryRequestOptionsBase getImpl() { + return this.actualRequestOptions; + } + + /////////////////////////////////////////////////////////////////////////////////////////// + // the following helper/accessor only helps to access this class outside of this package.// + /////////////////////////////////////////////////////////////////////////////////////////// + static void initialize() { + ImplementationBridgeHelpers.CosmosReadManyByPartitionKeysRequestOptionsHelper + .setCosmosReadManyByPartitionKeysRequestOptionsAccessor( + new ImplementationBridgeHelpers.CosmosReadManyByPartitionKeysRequestOptionsHelper + .CosmosReadManyByPartitionKeysRequestOptionsAccessor() { + @Override + public CosmosQueryRequestOptionsBase getImpl( + CosmosReadManyByPartitionKeysRequestOptions options) { + return options.actualRequestOptions; + } + + @Override + public String getContinuationToken(CosmosReadManyByPartitionKeysRequestOptions options) { + return options.actualRequestOptions.getContinuationToken(); + } + + @Override + public Integer getMaxConcurrentBatchPrefetch( + CosmosReadManyByPartitionKeysRequestOptions options) { + return options.actualRequestOptions.getMaxConcurrentBatchPrefetch(); + } + }); + } + + static { initialize(); } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java index fed27eb48945..da8764c95f85 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java @@ -55,6 +55,21 @@ private static ImplementationBridgeHelpers.CosmosDiagnosticsHelper.CosmosDiagnos private QueryInfo queryInfo; private QueryInfo.QueryPlanDiagnosticsContext queryPlanDiagnosticsContext; + // All header maps are produced by the SDK's own query pipeline. Non-null maps + // are always mutable (HashMap or ConcurrentHashMap) - the SDK intentionally + // allows callers to add/modify headers on FeedResponse. The only known + // exception is empty-page responses where the query pipeline may pass null. + // We do NOT clone non-null maps here to avoid unnecessary allocations on every + // FeedResponse construction - the wider blast radius of cloning (every query, + // change feed, readMany response) is not justified by the narrow null case. + // If a future code path introduces an immutable non-null header map, the + // setContinuationTokenInternal method will fail fast with + // UnsupportedOperationException, and the fix should be to make the upstream + // pipeline emit a mutable map rather than adding defensive cloning here. + private static Map ensureMutableHeadersMap(Map headers) { + return headers == null ? new HashMap<>() : headers; + } + FeedResponse(List results, Map headers) { this(results, headers, false, false, new ConcurrentHashMap<>()); } @@ -112,7 +127,7 @@ private FeedResponse( boolean nochanges, ConcurrentMap queryMetricsMap) { this.results = results; - this.header = header; + this.header = ensureMutableHeadersMap(header); this.usageHeaders = new HashMap<>(); this.quotaHeaders = new HashMap<>(); this.useEtagAsContinuation = useEtagAsContinuation; @@ -129,7 +144,7 @@ private FeedResponse( ConcurrentMap queryMetricsMap, CosmosDiagnostics diagnostics) { this.results = results; - this.header = header; + this.header = ensureMutableHeadersMap(header); this.usageHeaders = new HashMap<>(); this.quotaHeaders = new HashMap<>(); this.useEtagAsContinuation = useEtagAsContinuation; @@ -145,7 +160,7 @@ private FeedResponse( // NOTE - it is important to use HashMap over ConcurrentHashMap here because some keys/values might be null // and this is not allowed in ConcurrentHashMap - while it is ok in HashMap - this.header = toBeCloned.header != null ? new HashMap<>(toBeCloned.header) : null; + this.header = toBeCloned.header != null ? new HashMap<>(toBeCloned.header) : new HashMap<>(); this.usageHeaders = toBeCloned.usageHeaders != null ? new HashMap<>(toBeCloned.usageHeaders) : null; this.quotaHeaders = toBeCloned.quotaHeaders != null ? new HashMap<>(toBeCloned.quotaHeaders) : null; @@ -430,9 +445,23 @@ void setContinuationToken(String continuationToken) { ? HttpConstants.HttpHeaders.E_TAG : HttpConstants.HttpHeaders.CONTINUATION; + setContinuationTokenInternal(headerName, continuationToken); + } + + private void setContinuationTokenInternal(String headerName, String continuationToken) { if (!Strings.isNullOrWhiteSpace(continuationToken)) { this.header.put(headerName, continuationToken); - } else { + } else if (!this.header.isEmpty() && this.header.containsKey(headerName)) { + // The query API returns unmodifiable header collections for empty + // responses (no documents returned - when only header set is request charge) + // the protection here to check for existence of the header before attempting + // to remove it would not be robust enough against unknown headers + // but since we only ever call our own query pipeline + // avoiding cloning in all cases and gating on continuation header + // existence is a reasonable trade-off - test coverage exists that uncovered + // the problem - so, this acts as regression test as well + // --> the test coverage is in ItemsPartitionReaderWithReadManyByPartitionKeyITest + // it should "return empty results for non-existent partition keys" this.header.remove(headerName); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java index c34df7c4e925..7dde81e06e98 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java @@ -758,6 +758,7 @@ public static void initializeAllAccessors() { CosmosItemRequestOptions.initialize(); CosmosItemResponse.initialize(); CosmosPatchOperations.initialize(); + CosmosReadManyByPartitionKeysRequestOptions.initialize(); CosmosReadManyRequestOptions.initialize(); CosmosQueryRequestOptions.initialize(); CosmosOperationDetails.initialize(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java index 5ccd99db2fd0..53072b09326d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java @@ -107,6 +107,8 @@ public enum SinceVersion { /** v4.71.0 */ V4_71_0, /** v4.74.0 */ - V4_74_0 + V4_74_0, + /** v4.75.0 */ + V4_75_0 } } diff --git a/sdk/cosmos/cspell.yaml b/sdk/cosmos/cspell.yaml new file mode 100644 index 000000000000..2237242789e8 --- /dev/null +++ b/sdk/cosmos/cspell.yaml @@ -0,0 +1,8 @@ +import: + - ../../.vscode/cspell.json +overrides: + - filename: "**/sdk/cosmos/*" + words: + - DCOUNT + - dedupe + - colls diff --git a/sdk/cosmos/docs/readManyByPartitionKey-design.md b/sdk/cosmos/docs/readManyByPartitionKey-design.md new file mode 100644 index 000000000000..7edc808ba854 --- /dev/null +++ b/sdk/cosmos/docs/readManyByPartitionKey-design.md @@ -0,0 +1,474 @@ +# readManyByPartitionKeys - Design & Implementation + +## Overview + +New `readManyByPartitionKeys` methods on `CosmosAsyncContainer` / `CosmosContainer` that accept a +`List` (without item ids). The SDK splits the PK values by physical partition, +generates batched streaming queries per physical partition, and returns results as +`CosmosPagedFlux` / `CosmosPagedIterable`. The result stream is strictly batch-by-batch +ordered (see Phase 4) so each `FeedResponse` carries a usable composite continuation token. + +An optional `SqlQuerySpec` parameter lets callers supply a custom query for projections +and additional filters. The SDK appends the auto-generated PK WHERE clause to it and rejects +non-streamable shapes (aggregates / ORDER BY / DISTINCT / etc.) up front via a one-time +gateway query-plan validation. + +## High-level flow + +```mermaid +flowchart TD + A["Caller
readManyByPartitionKeys(pks, ?customQuery, ?options)"] --> B{"Continuation
token in options?"} + B -- no --> C[Resolve collection + PK definition] + C --> D[Normalize PKs
dedupe by EPK + sort] + D --> E[Optional: validate custom query
via gateway query plan] + E --> F[Look up routing map
PartitionKeyRangeCache] + F --> G[Group PKs per physical partition,
split into batches of N PKs,
compute batchFilter EPK range] + G --> H[Sort all batches by
batchFilter.minInclusive] + H --> I[mergeSequential
execution] + B -- yes --> J[Deserialize composite token] + J --> K[Validate version + collectionRid
+ queryHash + partitionKeySetHash] + K --> L[Look up routing map
PartitionKeyRangeCache] + L --> M[For each persisted batchFilter:
filter PKs into batch +
resolve partitionScope from cache] + M --> I + I --> N["FeedResponse stream
(each stamped with composite token)"] +``` + +## Decisions + +| Topic | Decision | +|-------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| API name | `readManyByPartitionKeys` - distinct name to avoid ambiguity with existing `readMany(List)` | +| Return type | `CosmosPagedFlux` (async) / `CosmosPagedIterable` (sync) | +| Custom query format | `SqlQuerySpec` - full query with parameters; SDK ANDs the PK filter | +| Partial HPK | Supported from the start; prefix PKs fan out via `getOverlappingRanges` | +| PK deduplication | SDK normalizes and deduplicates the partition-key set by EPK before batching; Spark callers should still dedupe the input DataFrame when practical for efficiency | +| Spark UDF | New `GetCosmosPartitionKeyValue` UDF | +| Custom query validation | Gateway query plan via the standard SDK query-plan retrieval path; reject aggregates/ORDER BY/DISTINCT/GROUP BY/DCount/OFFSET/LIMIT/non-streaming ORDER BY/vector/fulltext | +| PK list size | No hard upper-bound enforced; SDK batches internally per physical partition (default 100 PKs per batch, configurable via `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE`) | +| Eager validation | Null and empty PK list rejected eagerly (not lazily in reactive chain) | +| Telemetry | Separate span name `readManyByPartitionKeys.` (distinct from existing `readManyItems`) | +| Query construction | Table alias auto-detected from FROM clause; string literals and subqueries handled correctly | +| Continuation tokens | Composite token persists only the per-batch EPK FILTER; the routing scope (`FeedRange`) is rederived from the live `PartitionKeyRangeCache` per batch on resume so post-split refreshes are reflected automatically | + +## Phase 1 - SDK Core (`azure-cosmos`) + +### Step 1: New public overloads in CosmosAsyncContainer + +```java + CosmosPagedFlux readManyByPartitionKeys(List partitionKeys, Class classType) + CosmosPagedFlux readManyByPartitionKeys(List partitionKeys, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, + Class classType) + CosmosPagedFlux readManyByPartitionKeys(List partitionKeys, + SqlQuerySpec customQuery, + Class classType) + CosmosPagedFlux readManyByPartitionKeys(List partitionKeys, + SqlQuerySpec customQuery, + CosmosReadManyByPartitionKeysRequestOptions requestOptions, + Class classType) +``` + +All four overloads delegate to a private `readManyByPartitionKeyInternalFunc(...)`. + +> **Note:** the request-options type for `readManyByPartitionKeys` is the dedicated +> `CosmosReadManyByPartitionKeysRequestOptions` class (not the shared +> `CosmosReadManyRequestOptions` used by `readMany(List)`). It exposes +> `setContinuationToken(String)` / `getContinuationToken()` directly on the public surface +> and a `setMaxConcurrentBatchPrefetch(int)` knob that defaults to +> `Runtime.getRuntime().availableProcessors()` (the Spark connector overrides this default to +> `1` and exposes a config key `spark.cosmos.read.readManyByPk.maxConcurrentBatchPrefetch` for +> per-task tuning). + +**Eager validation:** The 4-arg method validates `partitionKeys` is non-null and non-empty before +constructing the reactive pipeline, throwing `IllegalArgumentException` synchronously. + +### Step 2: Sync wrappers in CosmosContainer + +Same four signatures returning `CosmosPagedIterable`, each delegating to the async container. + +### Step 3: Internal orchestration (RxDocumentClientImpl) + +#### Step 3a: First-call path (no continuation token) + +1. Resolve collection metadata + PK definition from cache. +2. Normalize the input PK list: compute the effective partition key (EPK) for each input, + deduplicate by EPK string, sort the surviving entries by EPK. +3. Optionally validate the custom query (Step 4) - one gateway query-plan call, executed + in parallel with the routing-map lookup. +4. Fetch the routing map from `partitionKeyRangeCache`. +5. For each `PartitionKeyRange` (physical partition): + - Bucket the normalized PKs whose EPK falls in the partition's range. + - Split into batches of `Configs.getReadManyByPkMaxBatchSize()` PKs (default 100). + - For each batch, compute the **batchFilter** EPK range: + `[epkOf(firstPkInBatch), epkOf(firstPkInNextBatch))` - or, for the last batch in a + partition, `[epkOf(firstPkInBatch), partitionScope.maxExclusive)`. + - The batch's **partitionScope** is the partition's EPK range itself. + - Build the per-batch `SqlQuerySpec` (Step 5). +6. Sort all batches across all partitions by `batchFilter.minInclusive` for deterministic + global EPK ordering. +7. Execute sequentially with `Flux.mergeSequential(sources, maxConcurrency, prefetch=1)` + so emission order is strictly batch-by-batch (items from batch *N+1* are never interleaved + with items from batch *N*) while still allowing eager subscription to the next batch + (small Reactor prefetch) for throughput. `maxConcurrency = Math.min(batchCount, cpuCount)`. +8. As each `FeedResponse` flows through, stamp it with a composite continuation token (Phase 4). + +#### Step 3b: Continuation path (token present) + +1. Resolve collection metadata + PK definition from cache (same as 3a step 1). +2. Normalize the input PK list (same as 3a step 2). +3. Deserialize the composite token; validate `version`, `collectionRid`, `queryHash`, + `partitionKeySetHash`. Any mismatch throws `IllegalArgumentException` eagerly inside + the reactive chain. +4. Fetch the routing map from `partitionKeyRangeCache` - required to rederive each batch's + `partitionScope` for FeedRange routing. +5. For each persisted `batchFilter` (current + remaining): + - Filter the normalized PK list down to PKs whose EPK falls inside `batchFilter`. + Skip the batch if the filter happens to match no PKs (e.g. caller dropped some inputs). + - Compute the routing scope by calling `resolvePartitionScopeFromBatchFilter(batchFilter, routingMap)`: + `[min(minEpk), max(maxEpk))` across all `PartitionKeyRange`s overlapping the batch + filter. This naturally returns the post-split boundaries when a split has happened. + - Build the per-batch `SqlQuerySpec`. +6. Stitch the sequence: the first batch starts from the persisted `backendContinuation`; + subsequent batches start fresh. +7. Execute the same sequential merge as the first-call path. + +### Step 4: Custom query validation + +One-time call per invocation using the same query-plan retrieval path and cacheability rules as +regular SDK queries. + +- `QueryPlanRetriever.getQueryPlanThroughGatewayAsync()` for the user query. +- Reject (`IllegalArgumentException`) if: + - `queryInfo.hasGroupBy()` - checked first (takes precedence over aggregates since + `hasAggregates()` also returns true for GROUP BY queries) + - `queryInfo.hasAggregates()` + - `queryInfo.hasOrderBy()` + - `queryInfo.hasDistinct()` + - `queryInfo.hasDCount()` + - `queryInfo.hasOffset()` + - `queryInfo.hasLimit()` + - `queryInfo.hasTop()` + - `queryInfo.hasNonStreamingOrderBy()` + - `partitionedQueryExecutionInfo.hasHybridSearchQueryInfo()` + - query plan details are unavailable (`queryInfo == null`) + +The validator is called in parallel with the routing-map lookup so it does not add a +serial round-trip. + +### Step 5: Query construction + +Query construction is implemented in `ReadManyByPartitionKeyQueryHelper`. The helper: + +- Extracts the table alias from the FROM clause (handles `FROM c`, `FROM root r`, `FROM x WHERE ...`) +- Handles string literals in queries (parens/keywords inside `'...'` are correctly skipped) +- Recognizes SQL keywords: WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING +- Uses parameterized queries (`@__rmPk_` prefix) to prevent SQL injection. The prefix is + reserved: any caller-supplied `SqlParameter` whose name starts with `@__rmPk_` is rejected + to prevent collisions with the auto-generated PK parameters. + +**Single PK (HASH):** + +```sql +{baseQuery} WHERE {alias}["{pkPath}"] IN (@__rmPk_0, @__rmPk_1, @__rmPk_2) +``` + +**Full HPK (MULTI_HASH):** + +```sql +{baseQuery} WHERE ({alias}["{path1}"] = @__rmPk_0 AND {alias}["{path2}"] = @__rmPk_1) + OR ({alias}["{path1}"] = @__rmPk_2 AND {alias}["{path2}"] = @__rmPk_3) +``` + +**Partial HPK (prefix-only):** + +```sql +{baseQuery} WHERE ({alias}["{path1}"] = @__rmPk_0) + OR ({alias}["{path1}"] = @__rmPk_1) +``` + +If the base query already has a WHERE clause: + +```sql +{selectAndFrom} WHERE ({existingWhere}) AND ({pkFilter}) +``` + +### Step 6: Interface wiring + +New method `readManyByPartitionKeys` added directly to `AsyncDocumentClient` interface, +implemented in `RxDocumentClientImpl`. The query executor is `createQueryInternal` (the same +internal entry point used by the streaming query API), invoked once per batch. + +A new `fetchQueryPlanForValidation` static method on `DocumentQueryExecutionContextFactory` +exposes the query-plan retrieval path for the custom-query validator. + +### Step 7: Configuration + +New configurable batch size via system property `COSMOS.READ_MANY_BY_PK_MAX_BATCH_SIZE` or +environment variable `COSMOS_READ_MANY_BY_PK_MAX_BATCH_SIZE` (default: 100, minimum: 1). +Both inputs are parsed via a shared safe-parse helper that logs a WARN and falls back to +the default on `NumberFormatException` or non-positive values. + +## Phase 2 - Spark Connector (`azure-cosmos-spark_3`) + +### Step 8: New UDF - `GetCosmosPartitionKeyValue` + +- Input: partition key value (single value or Seq for hierarchical PKs). +- Output: serialized PK string in format `pk([...json...])`. +- **Null handling:** Null input is serialized as a JSON-null partition key component. If callers + need `PartitionKey.NONE` semantics they must use the schema-matched path with + `spark.cosmos.read.readManyByPk.nullHandling=None`, which is only supported for single-path + partition keys. + +### Step 9: PK-only serialization helper + +`CosmosPartitionKeyHelper`: + +- `getCosmosPartitionKeyValueString(pkValues: List[Object]): String` - serialize to `pk([...])` format. +- `tryParsePartitionKey(serialized: String): Option[PartitionKey]` - deserialize; returns `None` + for malformed input including invalid JSON. The parser regex is anchored + (`^pk\((.*)\)$`) so substrings that merely contain a `pk(...)` literal are not accepted. +- When `spark.cosmos.read.readManyByPk.nullHandling=None` is used, hierarchical partition keys + with null components are rejected with a clear error because `PartitionKey.NONE` cannot be + used with multiple paths. + +### Step 10: `CosmosItemsDataSource.readManyByPartitionKeys` + +Static entry points that accept a DataFrame and Cosmos config. PK extraction supports two modes: + +1. **UDF-produced column**: DataFrame contains `_partitionKeyIdentity` column (from + `GetCosmosPartitionKeyValue` UDF). +2. **Schema-matched columns**: DataFrame columns match the container's PK paths. + +Top-level DataFrame columns may supply a full or prefix hierarchical partition key directly. +Nested partition key paths are not resolved automatically and must use the UDF-produced +`_partitionKeyIdentity` column. + +Falls back with `IllegalArgumentException` if neither mode is possible. + +### Step 11: `CosmosReadManyByPartitionKeyReader` + +Orchestrator that resolves schema, initializes and broadcasts client state to executors, then maps +each Spark partition to an `ItemsPartitionReaderWithReadManyByPartitionKey`. The wrapper iterator +closes the reader deterministically on exhaustion, on failures, and via Spark task-completion +callbacks. + +### Step 12: `ItemsPartitionReaderWithReadManyByPartitionKey` + +Spark `PartitionReader[InternalRow]` that: + +- Preserves the caller's PK list and lets the SDK normalize/dedupe by effective partition key; + callers should still dedupe the DataFrame upstream when practical. Logs a WARN if a single + partition reader receives more than 200000 PKs (large input materializes the full set in + memory plus the SDK's normalized batch metadata). +- Passes the pre-built request options (with throughput control, diagnostics, custom serializer) + to the SDK. +- Wraps each batch in a `TransientIOErrorsRetryingReadManyByPartitionKeyIterator` for retry + handling - on transient I/O failures the iterator re-creates the underlying `CosmosPagedFlux` + using the continuation token of the last fully-drained page (Phase 4). +- Short-circuits empty PK lists to avoid SDK rejection. + +## Phase 3 - Testing + +### Unit tests + +- Query construction: single PK, HPK full/partial, custom query composition, table alias detection. +- Query plan rejection: aggregates, ORDER BY, DISTINCT, GROUP BY (with and without aggregates), + DCOUNT. +- String literal handling: WHERE/parentheses inside string constants. +- Keyword detection: WHERE, ORDER, GROUP, JOIN, OFFSET, LIMIT, HAVING. +- PK serialization/deserialization roundtrip (including malformed JSON handling). +- `findTopLevelWhereIndex` edge cases: subqueries, string literals, case insensitivity. +- Continuation token: roundtrip, version-field presence, version-mismatch rejection, + partition-key-set hash stability across reordered/duplicate inputs, partition scope is NOT + persisted. +- Reserved parameter prefix: caller-supplied `@__rmPk_*` parameter is rejected. + +### Integration tests + +- End-to-end SDK: single PK basic, projections, filters, empty results, HPK full/partial, + request options propagation. +- Batch size validation: temporarily lowered batch size to exercise batching/sequential-merge + logic. +- Continuation tokens: round-trip resume, resume with reordered PK list (must produce same + results because the partitionKeySetHash is order-invariant), resume rejection on changed + query / changed PK set / different collection. +- Null/empty PK list rejection (eager validation). +- Spark connector: `ItemsPartitionReaderWithReadManyByPartitionKey` with known PK values and + non-existent PKs. +- Spark public API: nested partition key containers require `_partitionKeyIdentity` and succeed + when populated via `GetCosmosPartitionKeyValue`. +- `CosmosPartitionKeyHelper`: single/HPK roundtrip, case insensitivity, malformed input. + +## Phase 4 - Continuation Token Support for Retry Safety + +### Motivation + +The Spark `TransientIOErrorsRetryingReadManyByPartitionKeyIterator` replays from a continuation +token on transient I/O failures. For this to work, `readManyByPartitionKeys` must (a) emit a +meaningful continuation token on each `FeedResponse`, and (b) accept a continuation token that +resumes where the previous attempt left off. + +Because the SDK batches PK values across multiple physical partitions, the continuation token +must capture not only the backend query continuation within a single batch but also which +batches remain. Additionally, batches must be processed sequentially so that a continuation +token unambiguously identifies a position in the result stream. + +### Composite continuation token + +A new internal class `ReadManyByPartitionKeyContinuationToken` captures the resume state. Each +remaining/current batch is identified by a single `BatchDefinition` carrying just the +**batchFilter** EPK range (`[minInclusive, maxExclusive)` containing the EPKs of all PKs in that +batch). The **partitionScope** used as the FeedRange at execution time is intentionally NOT +persisted - it is rederived per batch on resume. + +| Field | Type | Description | +|-----------------------|----------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `version` | `int` (wire field `v`) | Wire format version. Current = 1. Tokens with an unknown version are rejected with `IllegalArgumentException`. | +| `remainingBatches` | `List` | EPK FILTER range of each batch not yet started. | +| `currentBatch` | `BatchDefinition` | EPK FILTER range of the batch currently being processed. | +| `backendContinuation` | `String` (nullable) | Backend query continuation within the current batch. `null` means "start from the beginning of this batch". | +| `collectionRid` | `String` | Resource id of the collection. Token is rejected on resume if the rid does not match. | +| `queryHash` | `String` (Murmur3-128 hex) | Stable hash of the custom `SqlQuerySpec` (text + parameter names + parameter values). Token is rejected on resume if the hash does not match. Resume intentionally requires the exact same query shape, not just a semantically equivalent query. | +| `partitionKeySetHash` | `String` (Murmur3-128 hex) | Stable hash of the deduplicated, sorted set of input EPKs. Order- and duplicate-invariant: a caller may shuffle the input PK list across resume attempts and still get the same data back. | + +EPK ranges are represented as `Range` (the existing SDK type used throughout the +routing layer) with `isMinInclusive = true` and `isMaxExclusive = false`. + +### Why batchFilter only, not partitionScope + +Persisting the `partitionScope` (the physical-partition EPK boundary) in the token would have +two drawbacks: + +- **Token bloat**: doubles the number of EPK ranges in the serialized form. +- **Stale-after-split routing**: if the partition splits between the moment the token was + emitted and the moment a caller resumes, the persisted `partitionScope` would no longer match + any current physical partition exactly, so the backend would silently fan the query out + across multiple partitions for the entire range - inflating RU charge until the token is + retired. + +Instead, the SDK rederives `partitionScope` per batch at execution time: + +```text +partitionScope = union(getOverlappingRanges(batchFilter)) + = [ min(minInclusive of overlapping ranges), + max(maxExclusive of overlapping ranges) ) +``` + +When the routing-map cache is fresh, this yields a range that exactly matches one or more +physical partitions, so the backend hits only the intended partition(s) at minimum RU cost. +When the cache is briefly stale immediately after a split, the SDK's `StaleResourceRetryPolicy` +refreshes it on the first miss, so any RU-cost elevation is bounded to a single retried +request. + +### Serialization + +JSON -> Base64, keeping the token opaque to callers. The serialized form uses short field +names (`v`, `rb`, `cb`, `bc`, `cr`, `qh`, `ph`) and a compact `BatchDefinition` shape with one +EPK range per batch (`bf` = batch filter). Example decoded: + +```json +{ + "v": 1, + "rb": [ + {"bf": {"min": "05C1E0", "max": "0BF333"}}, + {"bf": {"min": "0BF333", "max": "FF"}} + ], + "cb": {"bf": {"min": "", "max": "05C1E0"}}, + "bc": "eyJDb21wb3NpdGVUb2", + "cr": "dbs/myDb/colls/myColl", + "qh": "", + "ph": "" +} +``` + +### Sequential batch execution and stamping + +Batches are sorted by `batchFilter.minInclusive` (lexicographic) and processed via +`Flux.mergeSequential` so emission order is strictly batch-by-batch. Each `FeedResponse` +is stamped with a composite token reflecting its position in the stream: + +```text +FeedResponse 1: remaining=[B2, B3], current=B1, backend=null +FeedResponse 2: remaining=[B2, B3], current=B1, backend= + (B1 exhausted) +FeedResponse 3: remaining=[B3], current=B2, backend=null +FeedResponse 4: remaining=[B3], current=B2, backend= + (B2 exhausted) +FeedResponse 5: remaining=[], current=B3, backend=null +FeedResponse 6: remaining=[], current=B3, backend= + (done - final FeedResponse has null continuation token) +``` + +Reactor's `mergeSequential` is allowed to subscribe-eagerly to the next batch (small prefetch) +while the current batch is still draining, so backend-side prefetch of the next batch is +permitted and improves throughput, but it never reorders the emitted items. This guarantees +every `FeedResponse` carries a usable continuation token regardless of whether the caller plans +to resume or not. + +### Resume sequence + +```mermaid +sequenceDiagram + autonumber + participant Caller + participant SDK as RxDocumentClientImpl + participant Cache as PartitionKeyRangeCache + participant Backend + Caller->>SDK: readManyByPartitionKeys(pks, ?customQuery, options(token)) + SDK->>SDK: deserialize token
verify v + collectionRid
+ queryHash + partitionKeySetHash + SDK->>Cache: tryLookupAsync(collectionRid) + Cache-->>SDK: routingMap (post-split if cache refreshed) + loop per persisted batchFilter + SDK->>SDK: filterPksByEpkRange(batchFilter) + SDK->>SDK: resolvePartitionScopeFromBatchFilter(batchFilter, routingMap) + SDK->>Backend: createQueryInternal(SELECT WHERE pkFilter, FeedRange=partitionScope, [backendCont if 1st]) + Backend-->>SDK: FeedResponse pages + SDK->>SDK: stamp each FeedResponse with composite token (current=this, remaining=tail) + SDK-->>Caller: FeedResponse + end +``` + +### API changes + +**`CosmosReadManyByPartitionKeysRequestOptionsImpl`:** Holds `continuationToken`, +`maxConcurrentBatchPrefetch`, and `maxItemCount` fields with corresponding getters and setters; +all other knobs (consistency, session token, regions, throughput control, dedicated gateway, +diagnostics, custom item serializer, keyword identifiers, e2e latency policy) are inherited +from `CosmosQueryRequestOptionsBase`. The public `CosmosReadManyByPartitionKeysRequestOptions` +facade exposes only the surface that is meaningful for this operation - +`MaxBackendContinuationTokenSizeInKb` (renamed from the inherited +`ResponseContinuationTokenLimitInKb` to make clear it limits the *backend* token wrapped inside +the composite token), `MaxItemCount`, `MaxConcurrentBatchPrefetch`, `ContinuationToken`, plus +the standard query-options forwarders. + +**`RxDocumentClientImpl.readManyByPartitionKeys`:** see Step 3 above. + +**`CosmosAsyncContainer.readManyByPartitionKeyInternalFunc`:** Pulls the continuation token +out of the request options (before entering the reactive chain) and threads it into the +cloned `CosmosQueryRequestOptions` so `QueryFeedOperationState` carries it through to +`RxDocumentClientImpl`. + +### Spark reader integration + +**`ItemsPartitionReaderWithReadManyByPartitionKey`:** The `fluxFactory` lambda applies the +continuation token: when non-null, set it on the preconfigured request options before calling +`readManyByPartitionKeys`. + +**`TransientIOErrorsRetryingReadManyByPartitionKeyIterator`:** A lightweight clone of the +existing `TransientIOErrorsRetryingIterator` adapted for `readManyByPartitionKeys`. Captures +`lastContinuationToken` from each fully-drained `FeedResponse` and passes it back to the +factory on retry. The iterator must NEVER call `.subscribe()` on the `CosmosPagedFlux` it +holds (a bug previously fixed for the regular query iterator); it only consumes the +`iterableByPage().iterator()` view. + +### EPK range stability across partition splits + +The token uses EPK ranges, not PkRangeIds. If a split occurs between retries: + +- Routing scope is rederived from the current routing map (Step 3b/5), so the FeedRange used + for the resumed query exactly matches the post-split physical partition boundaries. +- A stale backend continuation is rejected by the backend; the SDK's + `StaleResourceRetryPolicy` retries with a fresh routing map while the EPK range ensures + correct targeting. +- The `partitionKeySetHash` and `queryHash` are unaffected by split (they hash the input PK + set and the user query, not the routing layout).