From 56155ab5d00525decee4f3c02812388011ea59ac Mon Sep 17 00:00:00 2001 From: Tom Longhurst <30480171+thomhurst@users.noreply.github.com> Date: Thu, 10 Sep 2026 10:21:57 +0100 Subject: [PATCH 1/3] fix(cosmos): keep point reads and deletes within the requested partition --- .../CosmosPointPartitionCompatibilityTests.cs | 56 +++++++++++++++ .../az/services/cosmos/CosmosHandler.java | 6 +- .../cosmos/CosmosPointPartitionTest.java | 68 +++++++++++++++++++ 3 files changed, 128 insertions(+), 2 deletions(-) create mode 100644 compatibility-tests/sdk-test-dotnet/CosmosPointPartitionCompatibilityTests.cs create mode 100644 src/test/java/io/floci/az/services/cosmos/CosmosPointPartitionTest.java diff --git a/compatibility-tests/sdk-test-dotnet/CosmosPointPartitionCompatibilityTests.cs b/compatibility-tests/sdk-test-dotnet/CosmosPointPartitionCompatibilityTests.cs new file mode 100644 index 00000000..515e23ac --- /dev/null +++ b/compatibility-tests/sdk-test-dotnet/CosmosPointPartitionCompatibilityTests.cs @@ -0,0 +1,56 @@ +using System.Net; +using Microsoft.Azure.Cosmos; + +namespace FlociAz.Compatibility; + +[NotInParallel] +public sealed class CosmosPointPartitionCompatibilityTests +{ + [Test] + [Timeout(60_000)] + public async Task PointReadsAndDeletesStayWithinTheRequestedPartition(CancellationToken cancellationToken) + { + string endpoint = Environment.GetEnvironmentVariable("FLOCI_AZ_ENDPOINT") ?? "http://localhost:4577"; + using var client = new CosmosClient($"{endpoint}/devstoreaccount1-cosmos/", + "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==", + new CosmosClientOptions { ConnectionMode = ConnectionMode.Gateway, LimitToEndpoint = true }); + Database database = await client.CreateDatabaseAsync( + $"dotnet-point-partition-{Guid.NewGuid():N}", cancellationToken: cancellationToken); + try + { + Container container = await database.CreateContainerAsync( + "items", "/pk", cancellationToken: cancellationToken); + await container.CreateItemAsync(new { id = "same-id", pk = "alice" }, + new PartitionKey("alice"), cancellationToken: cancellationToken); + + foreach (PartitionKey missing in new[] { new PartitionKey("bob"), new PartitionKey(""), PartitionKey.Null, PartitionKey.None }) + { + using ResponseMessage read = await container.ReadItemStreamAsync("same-id", missing, + cancellationToken: cancellationToken); + await Assert.That(read.StatusCode).IsEqualTo(HttpStatusCode.NotFound); + using ResponseMessage delete = await container.DeleteItemStreamAsync("same-id", missing, + cancellationToken: cancellationToken); + await Assert.That(delete.StatusCode).IsEqualTo(HttpStatusCode.NotFound); + } + + await container.CreateItemAsync(new { id = "same-id", pk = "bob" }, + new PartitionKey("bob"), cancellationToken: cancellationToken); + using ResponseMessage bobRead = await container.ReadItemStreamAsync("same-id", new PartitionKey("bob"), + cancellationToken: cancellationToken); + await Assert.That(bobRead.StatusCode).IsEqualTo(HttpStatusCode.OK); + using ResponseMessage bobDelete = await container.DeleteItemStreamAsync("same-id", new PartitionKey("bob"), + cancellationToken: cancellationToken); + await Assert.That(bobDelete.StatusCode).IsEqualTo(HttpStatusCode.NoContent); + using ResponseMessage bobMissing = await container.ReadItemStreamAsync("same-id", new PartitionKey("bob"), + cancellationToken: cancellationToken); + await Assert.That(bobMissing.StatusCode).IsEqualTo(HttpStatusCode.NotFound); + using ResponseMessage alice = await container.ReadItemStreamAsync("same-id", new PartitionKey("alice"), + cancellationToken: cancellationToken); + await Assert.That(alice.StatusCode).IsEqualTo(HttpStatusCode.OK); + } + finally + { + await database.DeleteAsync(cancellationToken: cancellationToken); + } + } +} diff --git a/src/main/java/io/floci/az/services/cosmos/CosmosHandler.java b/src/main/java/io/floci/az/services/cosmos/CosmosHandler.java index b1f08c2a..92ba524d 100644 --- a/src/main/java/io/floci/az/services/cosmos/CosmosHandler.java +++ b/src/main/java/io/floci/az/services/cosmos/CosmosHandler.java @@ -1276,9 +1276,11 @@ private StoredObject findDoc(AzureRequest req, String dbId, String collId, Strin String pkEnc = encodeKey(pk); String exact = docKey(req.accountName(), dbId, collId, pkEnc, docId); Optional found = liveDoc(store.get(exact), defaultTtl); - if (found.isPresent()) return found.get(); + if (found.isPresent() || req.headers().getHeaderString("x-ms-documentdb-partitionkey") != null) { + return found.orElse(null); + } } - // Fallback: scan (handles missing PK header or cross-partition reads) + // Only requests without a partition header may use the legacy unscoped lookup. String prefix = req.accountName() + K_DOC + dbId + "|" + collId + "|"; return liveDoc(store.scan(k -> k.startsWith(prefix) && k.endsWith("|" + docId)) .stream().findFirst(), defaultTtl).orElse(null); diff --git a/src/test/java/io/floci/az/services/cosmos/CosmosPointPartitionTest.java b/src/test/java/io/floci/az/services/cosmos/CosmosPointPartitionTest.java new file mode 100644 index 00000000..cc49313e --- /dev/null +++ b/src/test/java/io/floci/az/services/cosmos/CosmosPointPartitionTest.java @@ -0,0 +1,68 @@ +package io.floci.az.services.cosmos; + +import io.quarkus.test.junit.QuarkusTest; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import java.util.Map; + +import static io.restassured.RestAssured.given; +import static org.hamcrest.Matchers.is; + +@QuarkusTest +class CosmosPointPartitionTest { + private static final String BASE = "/pointpartition-cosmos/dbs/db"; + private static final String DOCS = BASE + "/colls/items/docs"; + + @BeforeEach + void setup() { + given().post("/_admin/reset").then().statusCode(204); + given().contentType("application/json").body(Map.of("id", "db")) + .post("/pointpartition-cosmos/dbs").then().statusCode(201); + given().contentType("application/json") + .body("{\"id\":\"items\",\"partitionKey\":{\"paths\":[\"/pk\"],\"kind\":\"Hash\"}}") + .post(BASE + "/colls").then().statusCode(201); + create("alice"); + } + + @ParameterizedTest + @ValueSource(strings = {"[\"bob\"]", "[\"\"]", "[null]", "[{}]"}) + void readWithMissingPartitionDoesNotFindAnotherPartitionsDocument(String partition) { + given().header("x-ms-documentdb-partitionkey", partition) + .get(DOCS + "/same-id").then().statusCode(404); + assertAliceExists(); + } + + @ParameterizedTest + @ValueSource(strings = {"[\"bob\"]", "[\"\"]", "[null]", "[{}]"}) + void deleteWithMissingPartitionDoesNotDeleteAnotherPartitionsDocument(String partition) { + given().header("x-ms-documentdb-partitionkey", partition) + .delete(DOCS + "/same-id").then().statusCode(404); + assertAliceExists(); + } + + @Test + void sameIdInDifferentPartitionsRemainsIndependent() { + create("bob"); + given().header("x-ms-documentdb-partitionkey", "[\"bob\"]") + .get(DOCS + "/same-id").then().statusCode(200).body("pk", is("bob")); + given().header("x-ms-documentdb-partitionkey", "[\"bob\"]") + .delete(DOCS + "/same-id").then().statusCode(204); + given().header("x-ms-documentdb-partitionkey", "[\"bob\"]") + .get(DOCS + "/same-id").then().statusCode(404); + assertAliceExists(); + } + + private void create(String partition) { + given().contentType("application/json") + .body(Map.of("id", "same-id", "pk", partition)) + .post(DOCS).then().statusCode(201); + } + + private void assertAliceExists() { + given().header("x-ms-documentdb-partitionkey", "[\"alice\"]") + .get(DOCS + "/same-id").then().statusCode(200).body("pk", is("alice")); + } +} From 5b5f88d23dc7a5d5e6e951369a576a9f72561b5e Mon Sep 17 00:00:00 2001 From: Tom Longhurst <30480171+thomhurst@users.noreply.github.com> Date: Thu, 10 Sep 2026 11:15:09 +0100 Subject: [PATCH 2/3] fix(cosmos): validate logical partition values during point lookup --- .../CosmosPointPartitionCompatibilityTests.cs | 8 +++++ .../az/services/cosmos/CosmosHandler.java | 29 +++++++++++++++++-- .../cosmos/CosmosPointPartitionTest.java | 24 +++++++++++++++ 3 files changed, 58 insertions(+), 3 deletions(-) diff --git a/compatibility-tests/sdk-test-dotnet/CosmosPointPartitionCompatibilityTests.cs b/compatibility-tests/sdk-test-dotnet/CosmosPointPartitionCompatibilityTests.cs index 515e23ac..98e76b7c 100644 --- a/compatibility-tests/sdk-test-dotnet/CosmosPointPartitionCompatibilityTests.cs +++ b/compatibility-tests/sdk-test-dotnet/CosmosPointPartitionCompatibilityTests.cs @@ -47,6 +47,14 @@ await container.CreateItemAsync(new { id = "same-id", pk = "bob" }, using ResponseMessage alice = await container.ReadItemStreamAsync("same-id", new PartitionKey("alice"), cancellationToken: cancellationToken); await Assert.That(alice.StatusCode).IsEqualTo(HttpStatusCode.OK); + await container.CreateItemAsync(new { id = "empty", pk = "" }, new PartitionKey(""), + cancellationToken: cancellationToken); + using ResponseMessage nullRead = await container.ReadItemStreamAsync("empty", PartitionKey.Null, + cancellationToken: cancellationToken); + await Assert.That(nullRead.StatusCode).IsEqualTo(HttpStatusCode.NotFound); + using ResponseMessage nullDelete = await container.DeleteItemStreamAsync("empty", PartitionKey.Null, + cancellationToken: cancellationToken); + await Assert.That(nullDelete.StatusCode).IsEqualTo(HttpStatusCode.NotFound); } finally { diff --git a/src/main/java/io/floci/az/services/cosmos/CosmosHandler.java b/src/main/java/io/floci/az/services/cosmos/CosmosHandler.java index 92ba524d..9a7ab2a5 100644 --- a/src/main/java/io/floci/az/services/cosmos/CosmosHandler.java +++ b/src/main/java/io/floci/az/services/cosmos/CosmosHandler.java @@ -1270,15 +1270,38 @@ private String encodeContinuationToken(int skip) { private StoredObject findDoc(AzureRequest req, String dbId, String collId, String docId) { Optional collFound = store.get(collKey(req.accountName(), dbId, collId)); Object defaultTtl = containerDefaultTtl(collFound); + String partitionHeader = req.headers().getHeaderString("x-ms-documentdb-partitionkey"); + if (partitionHeader != null) { + if (collFound.isEmpty()) { + return null; + } + try { + var partition = CosmosQueryPartition.parse(partitionHeader, parseData(collFound.get())); + String exact = docKey(req.accountName(), dbId, collId, + encodeKey(extractPartitionKeyValue(req)), docId); + Optional found = liveDoc(store.get(exact), defaultTtl) + .filter(object -> partition.test(parseData(object))); + if (found.isPresent()) { + return found.get(); + } + // Legacy storage keys stringify partition values. Compare logical values too: + // numbers may have equivalent encodings, while null and strings remain distinct. + String prefix = req.accountName() + K_DOC + dbId + "|" + collId + "|"; + return store.scan(key -> key.startsWith(prefix) && key.endsWith("|" + docId)).stream() + .filter(object -> partition.test(parseData(object))) + .map(object -> liveDoc(Optional.of(object), defaultTtl)) + .flatMap(Optional::stream).findFirst().orElse(null); + } catch (IllegalArgumentException e) { + throw new jakarta.ws.rs.WebApplicationException(errorResponse(400, "BadRequest", e.getMessage())); + } + } // Fast path: construct exact key using partition key from header if (collFound.isPresent()) { String pk = extractPartitionKeyValue(req); String pkEnc = encodeKey(pk); String exact = docKey(req.accountName(), dbId, collId, pkEnc, docId); Optional found = liveDoc(store.get(exact), defaultTtl); - if (found.isPresent() || req.headers().getHeaderString("x-ms-documentdb-partitionkey") != null) { - return found.orElse(null); - } + if (found.isPresent()) return found.get(); } // Only requests without a partition header may use the legacy unscoped lookup. String prefix = req.accountName() + K_DOC + dbId + "|" + collId + "|"; diff --git a/src/test/java/io/floci/az/services/cosmos/CosmosPointPartitionTest.java b/src/test/java/io/floci/az/services/cosmos/CosmosPointPartitionTest.java index cc49313e..54cc34d3 100644 --- a/src/test/java/io/floci/az/services/cosmos/CosmosPointPartitionTest.java +++ b/src/test/java/io/floci/az/services/cosmos/CosmosPointPartitionTest.java @@ -61,6 +61,30 @@ private void create(String partition) { .post(DOCS).then().statusCode(201); } + @Test + void typedPartitionValuesCannotReadOrDeleteStringPartitions() { + create(""); + create("42"); + for (String partition : new String[] {"[null]", "[42]"}) { + given().header("x-ms-documentdb-partitionkey", partition) + .get(DOCS + "/same-id").then().statusCode(404); + given().header("x-ms-documentdb-partitionkey", partition) + .delete(DOCS + "/same-id").then().statusCode(404); + } + given().header("x-ms-documentdb-partitionkey", "[\"\"]") + .get(DOCS + "/same-id").then().statusCode(200); + given().header("x-ms-documentdb-partitionkey", "[\"42\"]") + .get(DOCS + "/same-id").then().statusCode(200); + } + + @Test + void equivalentNumericPartitionRepresentationsStillMatch() { + given().contentType("application/json").body(Map.of("id", "numeric", "pk", 42)) + .post(DOCS).then().statusCode(201); + given().header("x-ms-documentdb-partitionkey", "[42.0]") + .get(DOCS + "/numeric").then().statusCode(200); + } + private void assertAliceExists() { given().header("x-ms-documentdb-partitionkey", "[\"alice\"]") .get(DOCS + "/same-id").then().statusCode(200).body("pk", is("alice")); From 0708484faa4e52612640b956277d5fede6c64c10 Mon Sep 17 00:00:00 2001 From: Tom Longhurst <30480171+thomhurst@users.noreply.github.com> Date: Thu, 10 Sep 2026 11:18:46 +0100 Subject: [PATCH 3/3] fix(cosmos): require complete document IDs in scoped fallback --- .../io/floci/az/services/cosmos/CosmosHandler.java | 5 ++++- .../az/services/cosmos/CosmosPointPartitionTest.java | 10 ++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/main/java/io/floci/az/services/cosmos/CosmosHandler.java b/src/main/java/io/floci/az/services/cosmos/CosmosHandler.java index 9a7ab2a5..4e4a8d6b 100644 --- a/src/main/java/io/floci/az/services/cosmos/CosmosHandler.java +++ b/src/main/java/io/floci/az/services/cosmos/CosmosHandler.java @@ -1288,7 +1288,10 @@ private StoredObject findDoc(AzureRequest req, String dbId, String collId, Strin // numbers may have equivalent encodings, while null and strings remain distinct. String prefix = req.accountName() + K_DOC + dbId + "|" + collId + "|"; return store.scan(key -> key.startsWith(prefix) && key.endsWith("|" + docId)).stream() - .filter(object -> partition.test(parseData(object))) + .filter(object -> { + Map document = parseData(object); + return docId.equals(document.get("id")) && partition.test(document); + }) .map(object -> liveDoc(Optional.of(object), defaultTtl)) .flatMap(Optional::stream).findFirst().orElse(null); } catch (IllegalArgumentException e) { diff --git a/src/test/java/io/floci/az/services/cosmos/CosmosPointPartitionTest.java b/src/test/java/io/floci/az/services/cosmos/CosmosPointPartitionTest.java index 54cc34d3..c68684c7 100644 --- a/src/test/java/io/floci/az/services/cosmos/CosmosPointPartitionTest.java +++ b/src/test/java/io/floci/az/services/cosmos/CosmosPointPartitionTest.java @@ -85,6 +85,16 @@ void equivalentNumericPartitionRepresentationsStillMatch() { .get(DOCS + "/numeric").then().statusCode(200); } + @Test + void scopedFallbackRequiresTheWholeDocumentId() { + given().contentType("application/json").body(Map.of("id", "prefix|missing", "pk", "alice")) + .post(DOCS).then().statusCode(201); + given().header("x-ms-documentdb-partitionkey", "[\"alice\"]") + .get(DOCS + "/missing").then().statusCode(404); + given().header("x-ms-documentdb-partitionkey", "[\"alice\"]") + .delete(DOCS + "/missing").then().statusCode(404); + } + private void assertAliceExists() { given().header("x-ms-documentdb-partitionkey", "[\"alice\"]") .get(DOCS + "/same-id").then().statusCode(200).body("pk", is("alice"));