Skip to content

Commit 1d2ffaa

Browse files
Annie LiangAnnie Liang
Annie Liang
authored and
Annie Liang
committed
change
1 parent 6c02664 commit 1d2ffaa

7 files changed

Lines changed: 527 additions & 66 deletions

File tree

‎sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosItemIdEncodingTest.java‎

Lines changed: 192 additions & 61 deletions
Large diffs are not rendered by default.

‎sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosPartitionKeyEncodingTest.java‎

Lines changed: 329 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,329 @@
1+
/*
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License.
4+
*
5+
*/
6+
7+
package com.azure.cosmos;
8+
9+
import com.azure.cosmos.implementation.Configs;
10+
import com.azure.cosmos.implementation.HttpConstants;
11+
import com.azure.cosmos.implementation.Utils;
12+
import com.azure.cosmos.models.CosmosContainerProperties;
13+
import com.azure.cosmos.models.CosmosItemRequestOptions;
14+
import com.azure.cosmos.models.CosmosItemResponse;
15+
import com.azure.cosmos.models.PartitionKey;
16+
import com.azure.cosmos.models.CosmosQueryRequestOptions;
17+
import com.azure.cosmos.models.CosmosPatchOperations;
18+
import com.azure.cosmos.models.FeedResponse;
19+
import com.azure.cosmos.rx.TestSuiteBase;
20+
import com.fasterxml.jackson.core.JsonParseException;
21+
import com.fasterxml.jackson.core.JsonProcessingException;
22+
import com.fasterxml.jackson.core.json.JsonReadFeature;
23+
import com.fasterxml.jackson.databind.ObjectMapper;
24+
import com.fasterxml.jackson.databind.node.ObjectNode;
25+
import org.assertj.core.api.Fail;
26+
import org.testng.annotations.AfterClass;
27+
import org.testng.annotations.BeforeClass;
28+
import org.testng.annotations.Factory;
29+
import org.testng.annotations.Test;
30+
import reactor.core.Exceptions;
31+
32+
import java.util.UUID;
33+
34+
import static org.assertj.core.api.Assertions.assertThat;
35+
import static org.assertj.core.api.Assertions.fail;
36+
37+
public class CosmosPartitionKeyEncodingTest extends TestSuiteBase {
38+
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper().configure(
39+
JsonReadFeature.ALLOW_UNESCAPED_CONTROL_CHARS.mappedFeature(),
40+
true
41+
);
42+
43+
private CosmosClient client;
44+
private CosmosContainer container;
45+
46+
@Factory(dataProvider = "clientBuildersWithDirectSessionIncludeComputeGateway")
47+
public CosmosPartitionKeyEncodingTest(CosmosClientBuilder clientBuilder) {
48+
super(clientBuilder.contentResponseOnWriteEnabled(true));
49+
}
50+
51+
@BeforeClass(groups = { "emulator" }, timeOut = SETUP_TIMEOUT)
52+
public void before_PartitionKeyTest() {
53+
assertThat(this.client).isNull();
54+
CosmosContainerProperties containerProperties = getCollectionDefinitionWithRangeRangeIndex();
55+
try {
56+
this.client = getClientBuilder().buildClient();
57+
getSharedCosmosDatabase(this.client.asyncClient()).createContainer(containerProperties).block();
58+
CosmosAsyncContainer asyncContainer =
59+
getSharedCosmosDatabase(this.client.asyncClient()).getContainer(containerProperties.getId());
60+
this.container = client
61+
.getDatabase(asyncContainer.getDatabase().getId())
62+
.getContainer(asyncContainer.getId());
63+
} catch (Exception error) {
64+
String message = String.format(
65+
"Failed creating separate container %s for PartitionKeyEncoding tests.",
66+
containerProperties.getId());
67+
68+
logger.error(message, error);
69+
70+
fail(message);
71+
}
72+
}
73+
74+
@AfterClass(groups = { "emulator" }, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true)
75+
public void afterClass() {
76+
assertThat(this.client).isNotNull();
77+
if (this.container != null) {
78+
container.delete();
79+
}
80+
this.client.close();
81+
}
82+
83+
@Test(groups = { "emulator" }, timeOut = TIMEOUT)
84+
public void partitionKeyAscii() {
85+
TestScenario scenario = new TestScenario(
86+
"PartitionKeyAscii",
87+
"AsciiPK-" + UUID.randomUUID(),
88+
new TestScenarioExpectations(
89+
ConnectionMode.GATEWAY.toString(),
90+
HttpConstants.StatusCodes.CREATED,
91+
HttpConstants.StatusCodes.OK,
92+
HttpConstants.StatusCodes.OK,
93+
HttpConstants.StatusCodes.NO_CONTENT,
94+
HttpConstants.StatusCodes.OK),
95+
new TestScenarioExpectations(
96+
"COMPUTE_GATEWAY",
97+
HttpConstants.StatusCodes.CREATED,
98+
HttpConstants.StatusCodes.OK,
99+
HttpConstants.StatusCodes.OK,
100+
HttpConstants.StatusCodes.NO_CONTENT,
101+
HttpConstants.StatusCodes.OK),
102+
new TestScenarioExpectations(
103+
ConnectionMode.DIRECT.toString(),
104+
HttpConstants.StatusCodes.CREATED,
105+
HttpConstants.StatusCodes.OK,
106+
HttpConstants.StatusCodes.OK,
107+
HttpConstants.StatusCodes.NO_CONTENT,
108+
HttpConstants.StatusCodes.OK));
109+
110+
this.executeTestCase(scenario);
111+
}
112+
113+
@Test(groups = { "emulator" }, timeOut = TIMEOUT)
114+
public void partitionKeyUnicode() {
115+
TestScenario scenario = new TestScenario(
116+
"PartitionKeyUnicode",
117+
"Unicode鱀-" + UUID.randomUUID(),
118+
new TestScenarioExpectations(
119+
ConnectionMode.GATEWAY.toString(),
120+
HttpConstants.StatusCodes.CREATED,
121+
HttpConstants.StatusCodes.OK,
122+
HttpConstants.StatusCodes.OK,
123+
HttpConstants.StatusCodes.NO_CONTENT,
124+
HttpConstants.StatusCodes.OK),
125+
new TestScenarioExpectations(
126+
"COMPUTE_GATEWAY",
127+
HttpConstants.StatusCodes.CREATED,
128+
HttpConstants.StatusCodes.OK,
129+
HttpConstants.StatusCodes.OK,
130+
HttpConstants.StatusCodes.NO_CONTENT,
131+
HttpConstants.StatusCodes.OK),
132+
new TestScenarioExpectations(
133+
ConnectionMode.DIRECT.toString(),
134+
HttpConstants.StatusCodes.CREATED,
135+
HttpConstants.StatusCodes.OK,
136+
HttpConstants.StatusCodes.OK,
137+
HttpConstants.StatusCodes.NO_CONTENT,
138+
HttpConstants.StatusCodes.OK));
139+
140+
this.executeTestCase(scenario);
141+
}
142+
143+
private void executeTestCase(TestScenario scenario) {
144+
TestScenarioExpectations expected =
145+
this.getConnectionPolicy().getConnectionMode() == ConnectionMode.DIRECT ?
146+
scenario.direct : this.getClientBuilder().getEndpoint().contains(COMPUTE_GATEWAY_EMULATOR_PORT) ?
147+
scenario.computeGateway : scenario.gateway;
148+
149+
logger.info("Scenario: {}, Id/PK: \"{}\"", scenario.name, scenario.id);
150+
151+
try {
152+
try {
153+
CosmosItemResponse<ObjectNode> response = this.container.createItem(
154+
getDocumentDefinition(scenario.id),
155+
new PartitionKey(scenario.id),
156+
null);
157+
158+
deserializeAndValidatePayload(response, scenario.id, expected.ExpectedCreateStatusCode);
159+
} catch (Throwable throwable) {
160+
CosmosException cosmosError = Utils.as(Exceptions.unwrap(throwable), CosmosException.class);
161+
if (cosmosError == null) {
162+
Fail.fail(
163+
"Unexpected exception type " + Exceptions.unwrap(throwable).getClass().getName(),
164+
throwable);
165+
}
166+
167+
logger.error(cosmosError.toString());
168+
169+
assertThat(cosmosError.getStatusCode())
170+
.isEqualTo(expected.ExpectedCreateStatusCode);
171+
172+
return;
173+
}
174+
175+
try {
176+
CosmosItemResponse<ObjectNode> response = this.container.readItem(
177+
scenario.id,
178+
new PartitionKey(scenario.id),
179+
ObjectNode.class);
180+
181+
deserializeAndValidatePayload(response, scenario.id, expected.ExpectedReadStatusCode);
182+
} catch (Throwable throwable) {
183+
CosmosException cosmosError = Utils.as(Exceptions.unwrap(throwable), CosmosException.class);
184+
if (cosmosError == null) {
185+
if (expected.ExpectedReadStatusCode == -1) {
186+
return;
187+
}
188+
189+
Fail.fail(
190+
"Unexpected exception type " + Exceptions.unwrap(throwable).getClass().getName(),
191+
throwable);
192+
}
193+
if (cosmosError.getStatusCode() == 0 &&
194+
cosmosError.getCause() instanceof IllegalArgumentException &&
195+
cosmosError.getCause().getCause() instanceof JsonParseException &&
196+
(cosmosError.getCause().toString().contains("<TITLE>Bad Request</TITLE>") ||
197+
cosmosError.getCause().getCause().toString().contains("<TITLE>Bad Request</TITLE>"))) {
198+
199+
logger.info("HTML BAD REQUEST", cosmosError);
200+
assertThat(expected.ExpectedReadStatusCode).isEqualTo(400);
201+
return;
202+
} else {
203+
logger.info("BAD REQUEST", cosmosError);
204+
assertThat(cosmosError.getStatusCode()).isEqualTo(expected.ExpectedReadStatusCode);
205+
}
206+
}
207+
208+
try {
209+
CosmosItemResponse<ObjectNode> response = this.container.replaceItem(
210+
getDocumentDefinition(scenario.id),
211+
scenario.id,
212+
new PartitionKey(scenario.id),
213+
null);
214+
215+
deserializeAndValidatePayload(response, scenario.id, expected.ExpectedReplaceStatusCode);
216+
} catch (Throwable throwable) {
217+
CosmosException cosmosError = Utils.as(Exceptions.unwrap(throwable), CosmosException.class);
218+
if (cosmosError == null) {
219+
Fail.fail(
220+
"Unexpected exception type " + Exceptions.unwrap(throwable).getClass().getName(),
221+
throwable);
222+
}
223+
assertThat(cosmosError.getStatusCode()).isEqualTo(expected.ExpectedReplaceStatusCode);
224+
}
225+
226+
try {
227+
CosmosItemResponse<Object> response = this.container.deleteItem(
228+
scenario.id,
229+
new PartitionKey(scenario.id),
230+
(CosmosItemRequestOptions) null);
231+
232+
assertThat(response.getStatusCode()).isEqualTo(expected.ExpectedDeleteStatusCode);
233+
} catch (Throwable throwable) {
234+
CosmosException cosmosError = Utils.as(Exceptions.unwrap(throwable), CosmosException.class);
235+
if (cosmosError == null) {
236+
Fail.fail(
237+
"Unexpected exception type " + Exceptions.unwrap(throwable).getClass().getName(),
238+
throwable);
239+
}
240+
assertThat(cosmosError.getStatusCode()).isEqualTo(expected.ExpectedDeleteStatusCode);
241+
}
242+
} finally {
243+
System.clearProperty(Configs.PREVENT_INVALID_ID_CHARS);
244+
}
245+
}
246+
247+
private void deserializeAndValidatePayload(
248+
CosmosItemResponse<ObjectNode> response,
249+
String expectedId,
250+
int expectedStatusCode) {
251+
252+
assertThat(response.getStatusCode()).isEqualTo(expectedStatusCode);
253+
assertThat(response.getItem().get("id").asText()).isEqualTo(expectedId);
254+
assertThat(response.getItem().get("mypk").asText()).isEqualTo(expectedId);
255+
}
256+
257+
private ObjectNode getDocumentDefinition(String documentId) {
258+
String json = String.format(
259+
"{ \"id\": \"%s\", \"mypk\": \"%s\" }",
260+
documentId,
261+
documentId);
262+
263+
try {
264+
return
265+
OBJECT_MAPPER.readValue(json, ObjectNode.class);
266+
} catch (JsonProcessingException jsonError) {
267+
fail("No json processing error expected", jsonError);
268+
269+
throw new IllegalStateException("No json processing error expected", jsonError);
270+
}
271+
}
272+
273+
private static class TestScenarioExpectations {
274+
275+
public TestScenarioExpectations(
276+
String connectionMode,
277+
int expectedCreateStatusCode,
278+
int expectedReadStatusCode,
279+
int expectedReplaceStatusCode,
280+
int expectedDeleteStatusCode,
281+
int expectedQueryStatusCode
282+
) {
283+
this.ConnectionMode = connectionMode;
284+
this.ExpectedCreateStatusCode = expectedCreateStatusCode;
285+
this.ExpectedReadStatusCode = expectedReadStatusCode;
286+
this.ExpectedReplaceStatusCode = expectedReplaceStatusCode;
287+
this.ExpectedDeleteStatusCode = expectedDeleteStatusCode;
288+
this.ExpectedQueryStatusCode = expectedQueryStatusCode;
289+
}
290+
291+
public String ConnectionMode;
292+
293+
public int ExpectedCreateStatusCode;
294+
295+
public int ExpectedReadStatusCode;
296+
297+
public int ExpectedReplaceStatusCode;
298+
299+
public int ExpectedDeleteStatusCode;
300+
301+
public int ExpectedQueryStatusCode;
302+
}
303+
304+
private static class TestScenario {
305+
public TestScenario(
306+
String name,
307+
String pkValue,
308+
TestScenarioExpectations gateway,
309+
TestScenarioExpectations computeGateway,
310+
TestScenarioExpectations direct) {
311+
312+
this.name = name;
313+
this.id = pkValue;
314+
this.gateway = gateway;
315+
this.computeGateway = computeGateway;
316+
this.direct = direct;
317+
}
318+
319+
public String name;
320+
321+
public String id;
322+
323+
public TestScenarioExpectations gateway;
324+
325+
public TestScenarioExpectations computeGateway;
326+
327+
public TestScenarioExpectations direct;
328+
}
329+
}

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -742,7 +742,7 @@ public static CosmosChangeFeedRequestOptions getEffectiveCosmosChangeFeedRequest
742742
cosmosChangeFeedRequestRequestOptions, pagedFluxOptions);
743743
}
744744

745-
static String escapeNonAscii(String partitionKeyJson) {
745+
public static String escapeNonAscii(String partitionKeyJson) {
746746
// if all are ascii original string will be returned, and avoids copying data.
747747
StringBuilder sb = null;
748748
for (int i = 0; i < partitionKeyJson.length(); i++) {

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/feedranges/FeedRangePartitionKeyImpl.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,7 @@ public Mono<RxDocumentServiceRequest> populateFeedRangeFilteringHeaders(
156156

157157
request.getHeaders().put(
158158
HttpConstants.HttpHeaders.PARTITION_KEY,
159-
this.partitionKey.toJson());
159+
Utils.escapeNonAscii(this.partitionKey.toJson()));
160160
request.setPartitionKeyInternal(this.partitionKey);
161161

162162
MetadataDiagnosticsContext metadataDiagnosticsCtx =

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextBase.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -330,7 +330,7 @@ private void populatePartitionKeyInfo(RxDocumentServiceRequest request, Partitio
330330
if (partitionKey != null) {
331331
request.setPartitionKeyInternal(partitionKey);
332332
request.setPartitionKeyDefinition(partitionKeyDefinition);
333-
request.getHeaders().put(HttpConstants.HttpHeaders.PARTITION_KEY, partitionKey.toJson());
333+
request.getHeaders().put(HttpConstants.HttpHeaders.PARTITION_KEY, Utils.escapeNonAscii(partitionKey.toJson()));
334334
}
335335
}
336336
}

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/ParallelDocumentQueryExecutionContextBase.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import com.azure.cosmos.implementation.ResourceType;
1414
import com.azure.cosmos.implementation.RxDocumentServiceRequest;
1515
import com.azure.cosmos.implementation.Strings;
16+
import com.azure.cosmos.implementation.Utils;
1617
import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl;
1718
import com.azure.cosmos.implementation.routing.PartitionKeyInternal;
1819
import com.azure.cosmos.models.CosmosQueryRequestOptions;
@@ -83,7 +84,7 @@ protected void initialize(
8384
// partitionKeyInternal gets passed as null which avoids the feedRange normalization.
8485
if (!PartitionKeyInternal.isPartialPartitionKeyQuery(collection, cosmosQueryRequestOptions.getPartitionKey())) {
8586
partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(cosmosQueryRequestOptions.getPartitionKey());
86-
headers.put(HttpConstants.HttpHeaders.PARTITION_KEY, partitionKeyInternal.toJson());
87+
headers.put(HttpConstants.HttpHeaders.PARTITION_KEY, Utils.escapeNonAscii(partitionKeyInternal.toJson()));
8788
}
8889
}
8990

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/QueryPlanRetriever.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ static Mono<PartitionedQueryExecutionInfo> getQueryPlanThroughGatewayAsync(Diagn
9696

9797
if (partitionKey != null && partitionKey != PartitionKey.NONE) {
9898
PartitionKeyInternal partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(partitionKey);
99-
requestHeaders.put(HttpConstants.HttpHeaders.PARTITION_KEY, partitionKeyInternal.toJson());
99+
requestHeaders.put(HttpConstants.HttpHeaders.PARTITION_KEY, Utils.escapeNonAscii(partitionKeyInternal.toJson()));
100100
}
101101

102102
final RxDocumentServiceRequest queryPlanRequest = RxDocumentServiceRequest.create(diagnosticsClientContext,

0 commit comments

Comments
 (0)