Skip to content

Commit 8e581dc

Browse files
xinlian12annie-macannie-mac
authored
Resource token fix - 1 (Azure#19547)
* Fix resourceToke in V4- part 1 Co-authored-by: annie-mac <annie-mac@s-pkges.redmond.corp.microsoft.com> Co-authored-by: annie-mac <annie-mac@Gurpreet.fareast.corp.microsoft.com>
1 parent f3cb50e commit 8e581dc

File tree

13 files changed

+625
-23
lines changed

13 files changed

+625
-23
lines changed

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import com.azure.cosmos.implementation.CosmosAuthorizationTokenResolver;
1414
import com.azure.cosmos.implementation.Database;
1515
import com.azure.cosmos.implementation.HttpConstants;
16+
import com.azure.cosmos.implementation.Permission;
1617
import com.azure.cosmos.implementation.TracerProvider;
1718
import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdMetrics;
1819
import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal;
@@ -32,13 +33,14 @@
3233
import reactor.core.publisher.Mono;
3334

3435
import java.io.Closeable;
36+
import java.util.ArrayList;
3537
import java.util.Iterator;
3638
import java.util.List;
3739
import java.util.ServiceLoader;
40+
import java.util.stream.Collectors;
3841

3942
import static com.azure.core.util.FluxUtil.withContext;
4043
import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount;
41-
import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument;
4244
import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
4345

4446
/**
@@ -93,6 +95,17 @@ public final class CosmosAsyncClient implements Closeable {
9395
this.clientTelemetryEnabled = builder.isClientTelemetryEnabled();
9496
this.contentResponseOnWriteEnabled = builder.isContentResponseOnWriteEnabled();
9597
this.tracerProvider = new TracerProvider(TRACER);
98+
99+
List<Permission> permissionList = new ArrayList<>();
100+
if (this.permissions != null) {
101+
permissionList =
102+
this.permissions
103+
.stream()
104+
.map(permissionProperties -> ModelBridgeInternal.getPermission(permissionProperties))
105+
.filter(permission -> permission != null)
106+
.collect(Collectors.toList());
107+
}
108+
96109
this.asyncDocumentClient = new AsyncDocumentClient.Builder()
97110
.withServiceEndpoint(this.serviceEndpoint)
98111
.withMasterKeyOrResourceToken(this.keyOrResourceToken)
@@ -105,6 +118,7 @@ public final class CosmosAsyncClient implements Closeable {
105118
.withTransportClientSharing(this.enableTransportClientSharing)
106119
.withContentResponseOnWriteEnabled(this.contentResponseOnWriteEnabled)
107120
.withTokenCredential(this.tokenCredential)
121+
.withPermissionFeed(permissionList)
108122
.build();
109123
}
110124

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,4 +139,4 @@ public <T> T getObject(Class<T> klass) throws IOException {
139139

140140
return MAPPER.treeToValue(this.getPropertyBag(), klass);
141141
}
142-
}
142+
}

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,16 @@
44
package com.azure.cosmos.implementation;
55

66
import com.azure.cosmos.BridgeInternal;
7+
import com.azure.cosmos.models.ModelBridgeInternal;
78
import com.azure.cosmos.models.PartitionKey;
89
import com.azure.cosmos.models.PermissionMode;
910
import com.fasterxml.jackson.databind.node.ArrayNode;
1011
import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
1112

1213
import java.util.Locale;
1314

15+
import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
16+
1417
/**
1518
* Represents a per-User Permission to access a specific resource e.g. item or container in the Azure Cosmos DB
1619
* database service.
@@ -112,9 +115,10 @@ public PartitionKey getResourcePartitionKey() {
112115
* @param partitionkey the partition key.
113116
*/
114117
public void setResourcePartitionKey(PartitionKey partitionkey) {
118+
checkNotNull(partitionkey, "Partition key can not be null");
119+
115120
BridgeInternal.setProperty(this,
116-
Constants.Properties.RESOURCE_PARTITION_KEY,
117-
BridgeInternal.getPartitionKeyInternal(partitionkey)
118-
.toJson());
121+
Constants.Properties.RESOURCE_PARTITION_KEY,
122+
new Object[]{ModelBridgeInternal.getPartitionKeyObject(partitionkey)});
119123
}
120124
}

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

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ public static String getAuthorizationTokenUsingResourceTokens(
123123
}
124124

125125
if (resourceToken == null) {
126-
throw new IllegalArgumentException(RMResources.ResourceTokenNotFound);
126+
throw new UnauthorizedException(RMResources.ResourceTokenNotFound);
127127
}
128128

129129
logger.debug("returned token for resourceAddress [{}] = [{}]",
@@ -171,21 +171,22 @@ public static String getAuthorizationTokenUsingResourceTokens(
171171
ResourceId tokenRid;
172172
final String key = entry.getKey();
173173
Pair<Boolean, ResourceId> pair = ResourceId.tryParse(key);
174-
ResourceId test1= pair.getRight().getDocumentCollectionId();
175-
boolean test = test1.equals(resourceId);
176-
if (!PathsHelper.isNameBased(key) && pair.getLeft()
174+
if (pair.getLeft()) {
175+
ResourceId test1= pair.getRight().getDocumentCollectionId();
176+
boolean test = test1.equals(resourceId);
177+
if (!PathsHelper.isNameBased(key) && pair.getLeft()
177178
&& pair.getRight().getDocumentCollectionId().equals(resourceId)) {
178-
List<PartitionKeyAndResourceTokenPair> resourceTokens = entry.getValue();
179-
if (resourceTokens != null && resourceTokens.size() > 0) {
180-
resourceToken = resourceTokens.get(0).getResourceToken();
179+
List<PartitionKeyAndResourceTokenPair> resourceTokens = entry.getValue();
180+
if (resourceTokens != null && resourceTokens.size() > 0) {
181+
resourceToken = resourceTokens.get(0).getResourceToken();
182+
}
181183
}
182184
}
183185
}
184-
185186
}
186187

187188
if (resourceToken == null) {
188-
throw new IllegalArgumentException(RMResources.ResourceTokenNotFound);
189+
throw new UnauthorizedException(RMResources.ResourceTokenNotFound);
189190
}
190191

191192
logger.debug("returned token for resourceAddress [{}] = [{}]",

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1515,6 +1515,7 @@ public String getUserAuthorizationToken(String resourceName,
15151515
if(resourceType.equals(ResourceType.DatabaseAccount)) {
15161516
return this.firstResourceTokenFromPermissionFeed;
15171517
}
1518+
15181519
return ResourceTokenAuthorizationHelper.getAuthorizationTokenUsingResourceTokens(resourceTokensMap, requestVerb, resourceName, headers);
15191520
}
15201521
}

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

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import com.azure.cosmos.implementation.ResourceType;
2828
import com.azure.cosmos.implementation.RxDocumentServiceRequest;
2929
import com.azure.cosmos.implementation.RxDocumentServiceResponse;
30+
import com.azure.cosmos.implementation.UnauthorizedException;
3031
import com.azure.cosmos.implementation.UserAgentContainer;
3132
import com.azure.cosmos.implementation.Utils;
3233
import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
@@ -311,13 +312,22 @@ public Mono<List<Address>> getServerAddressesViaGatewayAsync(
311312
headers.put(HttpConstants.HttpHeaders.X_DATE, Utils.nowAsRFC1123());
312313

313314
if (tokenProvider.getAuthorizationTokenType() != AuthorizationTokenType.AadToken) {
314-
String token = this.tokenProvider.getUserAuthorizationToken(
315+
String token = null;
316+
try {
317+
token = this.tokenProvider.getUserAuthorizationToken(
315318
collectionRid,
316319
ResourceType.Document,
317320
RequestVerb.GET,
318321
headers,
319322
AuthorizationTokenType.PrimaryMasterKey,
320323
request.properties);
324+
} catch (UnauthorizedException e) {
325+
// User doesn't have rid based resource token. Maybe user has name based.
326+
327+
if (logger.isDebugEnabled()) {
328+
logger.debug("User doesn't have resource token for collection rid {}", collectionRid);
329+
}
330+
}
321331

322332
if (token == null && request.getIsNameBased()) {
323333
// User doesn't have rid based resource token. Maybe user has name based.

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -104,9 +104,13 @@ private static <T extends Resource> Mono<Pair<List<Range<String>>,QueryInfo>> ge
104104
}
105105
}
106106

107-
queryExecutionInfoMono = QueryPlanRetriever
108-
.getQueryPlanThroughGatewayAsync(diagnosticsClientContext, client, query,
109-
resourceLink);
107+
queryExecutionInfoMono =
108+
QueryPlanRetriever.getQueryPlanThroughGatewayAsync(
109+
diagnosticsClientContext,
110+
client,
111+
query,
112+
resourceLink,
113+
cosmosQueryRequestOptions != null ? cosmosQueryRequestOptions.getPartitionKey() : null);
110114

111115
return queryExecutionInfoMono.flatMap(
112116
partitionedQueryExecutionInfo -> {

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,11 @@
33

44
package com.azure.cosmos.implementation.query;
55

6+
import com.azure.cosmos.BridgeInternal;
67
import com.azure.cosmos.implementation.DiagnosticsClientContext;
8+
import com.azure.cosmos.implementation.routing.PartitionKeyInternal;
79
import com.azure.cosmos.models.ModelBridgeInternal;
10+
import com.azure.cosmos.models.PartitionKey;
811
import com.azure.cosmos.models.SqlQuerySpec;
912
import com.azure.cosmos.implementation.BackoffRetryUtility;
1013
import com.azure.cosmos.implementation.DocumentClientRetryPolicy;
@@ -35,13 +38,19 @@ class QueryPlanRetriever {
3538
static Mono<PartitionedQueryExecutionInfo> getQueryPlanThroughGatewayAsync(DiagnosticsClientContext diagnosticsClientContext,
3639
IDocumentQueryClient queryClient,
3740
SqlQuerySpec sqlQuerySpec,
38-
String resourceLink) {
41+
String resourceLink,
42+
PartitionKey partitionKey) {
3943
final Map<String, String> requestHeaders = new HashMap<>();
4044
requestHeaders.put(HttpConstants.HttpHeaders.CONTENT_TYPE, RuntimeConstants.MediaTypes.JSON);
4145
requestHeaders.put(HttpConstants.HttpHeaders.IS_QUERY_PLAN_REQUEST, TRUE);
4246
requestHeaders.put(HttpConstants.HttpHeaders.SUPPORTED_QUERY_FEATURES, SUPPORTED_QUERY_FEATURES);
4347
requestHeaders.put(HttpConstants.HttpHeaders.QUERY_VERSION, HttpConstants.Versions.QUERY_VERSION);
4448

49+
if (partitionKey != null && partitionKey != PartitionKey.NONE) {
50+
PartitionKeyInternal partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(partitionKey);
51+
requestHeaders.put(HttpConstants.HttpHeaders.PARTITION_KEY, partitionKeyInternal.toJson());
52+
}
53+
4554
final RxDocumentServiceRequest request = RxDocumentServiceRequest.create(diagnosticsClientContext,
4655
OperationType.QueryPlan,
4756
ResourceType.Document,

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosPermissionProperties.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,7 @@ String getResourcePath(String databaseName) {
221221
if (databaseName == null || databaseName.isEmpty()) {
222222
throw new IllegalArgumentException("databaseName");
223223
}
224-
resourcePrefixPath.append("/").append(Paths.DATABASES_PATH_SEGMENT)
224+
resourcePrefixPath.append(Paths.DATABASES_PATH_SEGMENT)
225225
.append("/").append(databaseName);
226226

227227
if (this.containerName == null || this.containerName.isEmpty()) {
@@ -264,6 +264,10 @@ Permission getResource(String databaseName) {
264264
permission.setPermissionMode(this.permissionMode);
265265
permission.setResourceLink(getResourcePath(databaseName));
266266

267+
if (this.resourcePartitionKey != null) {
268+
permission.setResourcePartitionKey(this.resourcePartitionKey);
269+
}
270+
267271
return permission;
268272
}
269273

@@ -320,6 +324,10 @@ Permission getPermission(String databaseName) {
320324
return getResource(databaseName);
321325
}
322326

327+
Permission getPermission() {
328+
return this.permission;
329+
}
330+
323331
static List<CosmosPermissionProperties> getPermissions(List<Permission> results) {
324332
return results.stream().map(permission -> new CosmosPermissionProperties(permission.toJson()))
325333
.collect(Collectors.toList());

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,11 @@ public static Permission getPermission(CosmosPermissionProperties permissionProp
170170
return permissionProperties.getPermission(databaseName);
171171
}
172172

173+
@Warning(value = INTERNAL_USE_ONLY_WARNING)
174+
public static Permission getPermission(CosmosPermissionProperties permissionProperties) {
175+
return permissionProperties.getPermission();
176+
}
177+
173178
@Warning(value = INTERNAL_USE_ONLY_WARNING)
174179
public static List<CosmosPermissionProperties> getCosmosPermissionPropertiesFromResults(List<Permission> results) {
175180
return CosmosPermissionProperties.getPermissions(results);

0 commit comments

Comments
 (0)