Skip to content

Commit 404e93e

Browse files
MehaKaushikMeha Kaushik
MehaKaushik
and
Meha Kaushik
authored
Changes to allow specifying apiType in CosmosClient (Azure#25088)
* changes to include apitype in cosmosclientbuilder * introducing apitype enum * adding unit test * updating hostname in client * address comments * changing default apitype to null Co-authored-by: Meha Kaushik <[email protected]>
1 parent 8e0db5e commit 404e93e

File tree

13 files changed

+144
-15
lines changed

13 files changed

+144
-15
lines changed

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import com.azure.core.credential.TokenCredential;
88
import com.azure.core.util.Context;
99
import com.azure.core.util.tracing.Tracer;
10+
import com.azure.cosmos.implementation.ApiType;
1011
import com.azure.cosmos.implementation.AsyncDocumentClient;
1112
import com.azure.cosmos.implementation.Configs;
1213
import com.azure.cosmos.implementation.ConnectionPolicy;
@@ -69,6 +70,7 @@ public final class CosmosAsyncClient implements Closeable {
6970
private final TracerProvider tracerProvider;
7071
private final boolean contentResponseOnWriteEnabled;
7172
private static final Tracer TRACER;
73+
private final ApiType apiType;
7274

7375
static {
7476
ServiceLoader<Tracer> serviceLoader = ServiceLoader.load(Tracer.class);
@@ -95,6 +97,7 @@ public final class CosmosAsyncClient implements Closeable {
9597
this.clientTelemetryEnabled = builder.isClientTelemetryEnabled();
9698
this.contentResponseOnWriteEnabled = builder.isContentResponseOnWriteEnabled();
9799
this.tracerProvider = new TracerProvider(TRACER);
100+
this.apiType = builder.apiType();
98101

99102
List<Permission> permissionList = new ArrayList<>();
100103
if (this.permissions != null) {
@@ -120,6 +123,7 @@ public final class CosmosAsyncClient implements Closeable {
120123
.withTokenCredential(this.tokenCredential)
121124
.withState(builder.metadataCaches())
122125
.withPermissionFeed(permissionList)
126+
.withApiType(apiType)
123127
.build();
124128
}
125129

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import com.azure.core.annotation.ServiceClientBuilder;
66
import com.azure.core.credential.AzureKeyCredential;
77
import com.azure.core.credential.TokenCredential;
8+
import com.azure.cosmos.implementation.ApiType;
89
import com.azure.cosmos.implementation.Configs;
910
import com.azure.cosmos.implementation.ConnectionPolicy;
1011
import com.azure.cosmos.implementation.CosmosAuthorizationTokenResolver;
@@ -108,6 +109,7 @@ public class CosmosClientBuilder {
108109
private boolean multipleWriteRegionsEnabled = true;
109110
private boolean readRequestsFallbackEnabled = true;
110111
private boolean clientTelemetryEnabled = false;
112+
private ApiType apiType = null;
111113

112114
/**
113115
* Instantiates a new Cosmos client builder.
@@ -129,6 +131,22 @@ CosmosClientMetadataCachesSnapshot metadataCaches() {
129131
return this.state;
130132
}
131133

134+
/**
135+
* Sets an apiType for the builder.
136+
* @param apiType
137+
* @return current cosmosClientBuilder
138+
*/
139+
CosmosClientBuilder setApiType(ApiType apiType){
140+
this.apiType = apiType;
141+
return this;
142+
}
143+
144+
/**
145+
* Returns apiType for the Builder.
146+
* @return
147+
*/
148+
ApiType apiType(){ return this.apiType; }
149+
132150
/**
133151
* Session capturing is enabled by default for {@link ConsistencyLevel#SESSION}.
134152
* For other consistency levels, it is not needed, unless if you need occasionally send requests with Session
@@ -845,6 +863,16 @@ public void setCosmosClientMetadataCachesSnapshot(CosmosClientBuilder builder,
845863
public CosmosClientMetadataCachesSnapshot getCosmosClientMetadataCachesSnapshot(CosmosClientBuilder builder) {
846864
return builder.metadataCaches();
847865
}
866+
867+
@Override
868+
public void setCosmosClientApiType(CosmosClientBuilder builder, ApiType apiType) {
869+
builder.setApiType(apiType);
870+
}
871+
872+
@Override
873+
public ApiType getCosmosClientApiType(CosmosClientBuilder builder) {
874+
return builder.apiType();
875+
}
848876
});
849877
}
850878
}

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
package com.azure.cosmos.implementation;
5+
6+
/**
7+
* ApiTypes in the Azure Cosmos DB database service.
8+
*/
9+
public enum ApiType {
10+
NONE("None"),
11+
12+
MONGODB("MongoDB"),
13+
14+
GREMLIN("Gremlin"),
15+
16+
CASSANDRA("Cassandra"),
17+
18+
TABLE("Table"),
19+
20+
ETCD("Etcd"),
21+
22+
SQL("Sql"),
23+
24+
GREMLINV2("GremlinV2");
25+
26+
ApiType(String overWireValue) {
27+
this.overWireValue = overWireValue;
28+
}
29+
30+
private final String overWireValue;
31+
32+
@Override
33+
public String toString() {
34+
return this.overWireValue;
35+
}
36+
}

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ class Builder {
9090
boolean transportClientSharing;
9191
boolean contentResponseOnWriteEnabled;
9292
private CosmosClientMetadataCachesSnapshot state;
93+
private ApiType apiType;
9394

9495
public Builder withServiceEndpoint(String serviceEndpoint) {
9596
try {
@@ -105,6 +106,11 @@ public Builder withState(CosmosClientMetadataCachesSnapshot state) {
105106
return this;
106107
}
107108

109+
public Builder withApiType(ApiType apiType) {
110+
this.apiType = apiType;
111+
return this;
112+
}
113+
108114
/**
109115
* New method withMasterKeyOrResourceToken will take either master key or resource token
110116
* and perform authentication for accessing resource.
@@ -232,7 +238,8 @@ public AsyncDocumentClient build() {
232238
sessionCapturingOverride,
233239
transportClientSharing,
234240
contentResponseOnWriteEnabled,
235-
state);
241+
state,
242+
apiType);
236243

237244
client.init(state, null);
238245
return client;

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ public static void setCosmosClientBuilderAccessor(final CosmosClientBuilderAcces
6262
accessor = newAccessor;
6363
}
6464

65-
static CosmosClientBuilderAccessor getCosmosClientBuilderAccessor() {
65+
public static CosmosClientBuilderAccessor getCosmosClientBuilderAccessor() {
6666
if (accessor == null) {
6767
throw new IllegalStateException("CosmosClientBuilder accessor is not initialized yet!");
6868
}
@@ -73,8 +73,12 @@ static CosmosClientBuilderAccessor getCosmosClientBuilderAccessor() {
7373
public interface CosmosClientBuilderAccessor {
7474
void setCosmosClientMetadataCachesSnapshot(CosmosClientBuilder builder,
7575
CosmosClientMetadataCachesSnapshot metadataCache);
76+
7677
CosmosClientMetadataCachesSnapshot getCosmosClientMetadataCachesSnapshot(CosmosClientBuilder builder);
7778

79+
void setCosmosClientApiType(CosmosClientBuilder builder, ApiType apiType);
80+
81+
ApiType getCosmosClientApiType(CosmosClientBuilder builder);
7882
}
7983
}
8084

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

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,7 @@ public class RxDocumentClientImpl implements AsyncDocumentClient, IAuthorization
148148
private final AtomicBoolean closed = new AtomicBoolean(false);
149149
private final int clientId;
150150
private ClientTelemetry clientTelemetry;
151+
private ApiType apiType;
151152

152153
// RetryPolicy retries a request when it encounters session unavailable (see ClientRetryPolicy).
153154
// Once it exhausts all write regions it clears the session container, then it uses RxClientCollectionCache
@@ -188,9 +189,10 @@ public RxDocumentClientImpl(URI serviceEndpoint,
188189
boolean sessionCapturingOverride,
189190
boolean connectionSharingAcrossClientsEnabled,
190191
boolean contentResponseOnWriteEnabled,
191-
CosmosClientMetadataCachesSnapshot metadataCachesSnapshot) {
192+
CosmosClientMetadataCachesSnapshot metadataCachesSnapshot,
193+
ApiType apiType) {
192194
this(serviceEndpoint, masterKeyOrResourceToken, permissionFeed, connectionPolicy, consistencyLevel, configs,
193-
credential, null, sessionCapturingOverride, connectionSharingAcrossClientsEnabled, contentResponseOnWriteEnabled, metadataCachesSnapshot);
195+
credential, null, sessionCapturingOverride, connectionSharingAcrossClientsEnabled, contentResponseOnWriteEnabled, metadataCachesSnapshot, apiType);
194196
this.cosmosAuthorizationTokenResolver = cosmosAuthorizationTokenResolver;
195197
}
196198

@@ -206,9 +208,10 @@ public RxDocumentClientImpl(URI serviceEndpoint,
206208
boolean sessionCapturingOverride,
207209
boolean connectionSharingAcrossClientsEnabled,
208210
boolean contentResponseOnWriteEnabled,
209-
CosmosClientMetadataCachesSnapshot metadataCachesSnapshot) {
211+
CosmosClientMetadataCachesSnapshot metadataCachesSnapshot,
212+
ApiType apiType) {
210213
this(serviceEndpoint, masterKeyOrResourceToken, permissionFeed, connectionPolicy, consistencyLevel, configs,
211-
credential, tokenCredential, sessionCapturingOverride, connectionSharingAcrossClientsEnabled, contentResponseOnWriteEnabled, metadataCachesSnapshot);
214+
credential, tokenCredential, sessionCapturingOverride, connectionSharingAcrossClientsEnabled, contentResponseOnWriteEnabled, metadataCachesSnapshot, apiType);
212215
this.cosmosAuthorizationTokenResolver = cosmosAuthorizationTokenResolver;
213216
}
214217

@@ -223,9 +226,10 @@ private RxDocumentClientImpl(URI serviceEndpoint,
223226
boolean sessionCapturingOverrideEnabled,
224227
boolean connectionSharingAcrossClientsEnabled,
225228
boolean contentResponseOnWriteEnabled,
226-
CosmosClientMetadataCachesSnapshot metadataCachesSnapshot) {
229+
CosmosClientMetadataCachesSnapshot metadataCachesSnapshot,
230+
ApiType apiType) {
227231
this(serviceEndpoint, masterKeyOrResourceToken, connectionPolicy, consistencyLevel, configs,
228-
credential, tokenCredential, sessionCapturingOverrideEnabled, connectionSharingAcrossClientsEnabled, contentResponseOnWriteEnabled, metadataCachesSnapshot);
232+
credential, tokenCredential, sessionCapturingOverrideEnabled, connectionSharingAcrossClientsEnabled, contentResponseOnWriteEnabled, metadataCachesSnapshot, apiType);
229233

230234
if (permissionFeed != null && permissionFeed.size() > 0) {
231235
this.resourceTokensMap = new HashMap<>();
@@ -279,7 +283,8 @@ private RxDocumentClientImpl(URI serviceEndpoint,
279283
boolean sessionCapturingOverrideEnabled,
280284
boolean connectionSharingAcrossClientsEnabled,
281285
boolean contentResponseOnWriteEnabled,
282-
CosmosClientMetadataCachesSnapshot metadataCachesSnapshot) {
286+
CosmosClientMetadataCachesSnapshot metadataCachesSnapshot,
287+
ApiType apiType) {
283288

284289
activeClientsCnt.incrementAndGet();
285290
this.clientId = clientIdGenerator.getAndDecrement();
@@ -364,6 +369,7 @@ private RxDocumentClientImpl(URI serviceEndpoint,
364369
this.resetSessionTokenRetryPolicy = retryPolicy;
365370
CpuMemoryMonitor.register(this);
366371
this.queryPlanCache = Collections.synchronizedMap(new SizeLimitingLRUCache(Constants.QUERYPLAN_CACHE_SIZE));
372+
this.apiType = apiType;
367373
} catch (RuntimeException e) {
368374
logger.error("unexpected failure in initializing client.", e);
369375
close();
@@ -1505,6 +1511,10 @@ private Mono<RxDocumentServiceRequest> populateHeaders(RxDocumentServiceRequest
15051511
request.getHeaders().put(HttpConstants.HttpHeaders.AUTHORIZATION, authorization);
15061512
}
15071513

1514+
if (this.apiType != null) {
1515+
request.getHeaders().put(HttpConstants.HttpHeaders.API_TYPE, this.apiType.toString());
1516+
}
1517+
15081518
if ((RequestVerb.POST.equals(httpMethod) || RequestVerb.PUT.equals(httpMethod))
15091519
&& !request.getHeaders().containsKey(HttpConstants.HttpHeaders.CONTENT_TYPE)) {
15101520
request.getHeaders().put(HttpConstants.HttpHeaders.CONTENT_TYPE, RuntimeConstants.MediaTypes.JSON);

‎sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/ClientUnderTestBuilder.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
// Licensed under the MIT License.
33
package com.azure.cosmos;
44

5+
import com.azure.cosmos.implementation.ApiType;
6+
import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
57
import com.azure.cosmos.implementation.RxDocumentClientUnderTest;
68
import com.azure.cosmos.implementation.Strings;
79
import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils;
@@ -36,6 +38,9 @@ public ClientUnderTestBuilder(CosmosClientBuilder builder) {
3638
this.multipleWriteRegionsEnabled(builder.isMultipleWriteRegionsEnabled());
3739
this.readRequestsFallbackEnabled(builder.isReadRequestsFallbackEnabled());
3840
this.clientTelemetryEnabled(builder.isClientTelemetryEnabled());
41+
ImplementationBridgeHelpers.CosmosClientBuilderHelper.CosmosClientBuilderAccessor accessor =
42+
ImplementationBridgeHelpers.CosmosClientBuilderHelper.getCosmosClientBuilderAccessor();
43+
accessor.setCosmosClientApiType(this, builder.apiType());
3944
}
4045

4146
@Override
@@ -50,7 +55,8 @@ public CosmosAsyncClient buildAsyncClient() {
5055
this.getConsistencyLevel(),
5156
this.configs(),
5257
this.getCredential(),
53-
this.isContentResponseOnWriteEnabled());
58+
this.isContentResponseOnWriteEnabled(),
59+
this.apiType());
5460
} catch (URISyntaxException e) {
5561
throw new IllegalArgumentException(e.getMessage());
5662
}

‎sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosAsyncClientTest.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@
33

44
package com.azure.cosmos;
55

6+
import com.azure.cosmos.implementation.ApiType;
67
import com.azure.cosmos.implementation.ConnectionPolicy;
8+
import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
79
import com.azure.cosmos.implementation.guava27.Strings;
810
import org.testng.ITest;
911
import org.testng.annotations.AfterMethod;

‎sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
22
// Licensed under the MIT License.
33
package com.azure.cosmos;
44

5-
import com.azure.cosmos.implementation.TestConfigurations;
5+
import com.azure.cosmos.implementation.*;
6+
import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils;
67
import com.google.common.collect.ImmutableList;
78
import org.testng.annotations.Test;
89

@@ -44,4 +45,27 @@ public void validateBadPreferredRegions2() {
4445
assertThat(e.getMessage()).isEqualTo("preferredRegion can't be empty");
4546
}
4647
}
48+
49+
@Test(groups = "emulator")
50+
public void validateApiTypePresent() {
51+
ApiType apiType = ApiType.TABLE;
52+
DirectConnectionConfig directConnectionConfig = DirectConnectionConfig.getDefaultConfig();
53+
54+
CosmosClientBuilder cosmosClientBuilder = new CosmosClientBuilder()
55+
.endpoint(TestConfigurations.HOST)
56+
.key(TestConfigurations.MASTER_KEY)
57+
.directMode(directConnectionConfig)
58+
.userAgentSuffix("custom-direct-client")
59+
.multipleWriteRegionsEnabled(false)
60+
.endpointDiscoveryEnabled(false)
61+
.readRequestsFallbackEnabled(true);
62+
63+
ImplementationBridgeHelpers.CosmosClientBuilderHelper.CosmosClientBuilderAccessor accessor =
64+
ImplementationBridgeHelpers.CosmosClientBuilderHelper.getCosmosClientBuilderAccessor();
65+
accessor.setCosmosClientApiType(cosmosClientBuilder, apiType);
66+
67+
RxDocumentClientImpl documentClient =
68+
(RxDocumentClientImpl) ReflectionUtils.getAsyncDocumentClient(new CosmosAsyncClient(cosmosClientBuilder));
69+
assertThat(ReflectionUtils.getApiType(documentClient)).isEqualTo(apiType);
70+
}
4771
}

‎sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/RxDocumentClientUnderTest.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,10 @@ public RxDocumentClientUnderTest(URI serviceEndpoint,
3737
ConsistencyLevel consistencyLevel,
3838
Configs configs,
3939
AzureKeyCredential credential,
40-
boolean contentResponseOnWriteEnabled) {
40+
boolean contentResponseOnWriteEnabled,
41+
ApiType apiType) {
4142
super(serviceEndpoint, masterKey, connectionPolicy, consistencyLevel, configs, credential, null, false,
42-
false, contentResponseOnWriteEnabled, null);
43+
false, contentResponseOnWriteEnabled, null, apiType);
4344
init(null, null);
4445
}
4546

0 commit comments

Comments
 (0)