Search in sources :

Example 6 with RequestOptions

use of com.microsoft.azure.cosmosdb.RequestOptions in project ambry by linkedin.

the class CosmosDataAccessor method updateMetadata.

/**
 * Update the blob metadata document in the CosmosDB collection.
 * @param blobId the {@link BlobId} for which metadata is replaced.
 * @param updateFields Map of field names and new values to update.
 * @return the {@link ResourceResponse} returned by the operation, if successful.
 * Returns {@Null} if the field already has the specified value.
 * @throws DocumentClientException if the record was not found or if the operation failed.
 */
ResourceResponse<Document> updateMetadata(BlobId blobId, Map<String, String> updateFields) throws DocumentClientException {
    // Read the existing record
    String docLink = getDocumentLink(blobId.getID());
    RequestOptions options = getRequestOptions(blobId.getPartition().toPathString());
    ResourceResponse<Document> readResponse = executeCosmosAction(() -> asyncDocumentClient.readDocument(docLink, options).toBlocking().single(), azureMetrics.documentReadTime);
    Document doc = readResponse.getResource();
    // Update only if value has changed
    Map<String, String> fieldsToUpdate = updateFields.entrySet().stream().filter(map -> !String.valueOf(updateFields.get(map.getKey())).equals(doc.get(map.getKey()))).collect(Collectors.toMap(Map.Entry::getKey, map -> String.valueOf(map.getValue())));
    if (fieldsToUpdate.size() == 0) {
        logger.debug("No change in value for {} in blob {}", updateFields.keySet(), blobId.getID());
        return null;
    }
    // For testing conflict handling
    if (updateCallback != null) {
        try {
            updateCallback.call();
        } catch (Exception ex) {
            logger.error("Error in update callback", ex);
        }
    }
    // Perform the update
    fieldsToUpdate.forEach((key, value) -> doc.set(key, value));
    // Set condition to ensure we don't clobber a concurrent update
    AccessCondition accessCondition = new AccessCondition();
    accessCondition.setCondition(doc.getETag());
    options.setAccessCondition(accessCondition);
    try {
        return executeCosmosAction(() -> asyncDocumentClient.replaceDocument(doc, options).toBlocking().single(), azureMetrics.documentUpdateTime);
    } catch (DocumentClientException e) {
        if (e.getStatusCode() == HttpConstants.StatusCodes.PRECONDITION_FAILED) {
            azureMetrics.blobUpdateConflictCount.inc();
        }
        throw e;
    }
}
Also used : HttpURLConnection(java.net.HttpURLConnection) SqlParameterCollection(com.microsoft.azure.cosmosdb.SqlParameterCollection) ConnectionMode(com.microsoft.azure.cosmosdb.ConnectionMode) DocumentClientException(com.microsoft.azure.cosmosdb.DocumentClientException) Date(java.util.Date) LoggerFactory(org.slf4j.LoggerFactory) VcrMetrics(com.github.ambry.cloud.VcrMetrics) ConsistencyLevel(com.microsoft.azure.cosmosdb.ConsistencyLevel) FeedResponse(com.microsoft.azure.cosmosdb.FeedResponse) JSONObject(org.json.JSONObject) Map(java.util.Map) SecretClient(com.azure.security.keyvault.secrets.SecretClient) Container(com.github.ambry.account.Container) StoredProcedureResponse(com.microsoft.azure.cosmosdb.StoredProcedureResponse) Set(java.util.Set) Utils(com.github.ambry.utils.Utils) Collectors(java.util.stream.Collectors) PartitionKey(com.microsoft.azure.cosmosdb.PartitionKey) List(java.util.List) CloudRequestAgent(com.github.ambry.cloud.CloudRequestAgent) SqlQuerySpec(com.microsoft.azure.cosmosdb.SqlQuerySpec) Document(com.microsoft.azure.cosmosdb.Document) Timer(com.codahale.metrics.Timer) BlobId(com.github.ambry.commons.BlobId) ResourceResponse(com.microsoft.azure.cosmosdb.ResourceResponse) RequestOptions(com.microsoft.azure.cosmosdb.RequestOptions) DocumentCollection(com.microsoft.azure.cosmosdb.DocumentCollection) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) Callable(java.util.concurrent.Callable) ArrayList(java.util.ArrayList) CloudConfig(com.github.ambry.config.CloudConfig) HashSet(java.util.HashSet) BiConsumer(java.util.function.BiConsumer) CloudStorageException(com.github.ambry.cloud.CloudStorageException) Properties(java.util.Properties) Logger(org.slf4j.Logger) AsyncDocumentClient(com.microsoft.azure.cosmosdb.rx.AsyncDocumentClient) Iterator(java.util.Iterator) VerifiableProperties(com.github.ambry.config.VerifiableProperties) ChangeFeedOptions(com.microsoft.azure.cosmosdb.ChangeFeedOptions) AccessCondition(com.microsoft.azure.cosmosdb.AccessCondition) ConnectionPolicy(com.microsoft.azure.cosmosdb.ConnectionPolicy) SecretClientBuilder(com.azure.security.keyvault.secrets.SecretClientBuilder) HttpConstants(com.microsoft.azure.cosmosdb.internal.HttpConstants) RetryOptions(com.microsoft.azure.cosmosdb.RetryOptions) FeedOptions(com.microsoft.azure.cosmosdb.FeedOptions) StoredProcedure(com.microsoft.azure.cosmosdb.StoredProcedure) SqlParameter(com.microsoft.azure.cosmosdb.SqlParameter) CloudBlobMetadata(com.github.ambry.cloud.CloudBlobMetadata) BlockingObservable(rx.observables.BlockingObservable) RequestOptions(com.microsoft.azure.cosmosdb.RequestOptions) AccessCondition(com.microsoft.azure.cosmosdb.AccessCondition) Document(com.microsoft.azure.cosmosdb.Document) Map(java.util.Map) DocumentClientException(com.microsoft.azure.cosmosdb.DocumentClientException) CloudStorageException(com.github.ambry.cloud.CloudStorageException) DocumentClientException(com.microsoft.azure.cosmosdb.DocumentClientException)

Example 7 with RequestOptions

use of com.microsoft.azure.cosmosdb.RequestOptions in project ambry by linkedin.

the class CosmosDataAccessor method getRequestOptions.

private RequestOptions getRequestOptions(String partitionPath) {
    RequestOptions options = new RequestOptions();
    options.setPartitionKey(new PartitionKey(partitionPath));
    return options;
}
Also used : RequestOptions(com.microsoft.azure.cosmosdb.RequestOptions) PartitionKey(com.microsoft.azure.cosmosdb.PartitionKey)

Example 8 with RequestOptions

use of com.microsoft.azure.cosmosdb.RequestOptions in project ambry by linkedin.

the class CosmosDataAccessor method testConnectivity.

/**
 * Test connectivity to Azure CosmosDB
 */
void testConnectivity() {
    ResourceResponse<DocumentCollection> response = asyncDocumentClient.readCollection(cosmosCollectionLink, new RequestOptions()).toBlocking().single();
    if (response.getResource() == null) {
        throw new IllegalStateException("CosmosDB collection not found: " + cosmosCollectionLink);
    }
    logger.info("CosmosDB connection test succeeded.");
    if (purgeBatchSize > 1) {
        // Check for existence of BulkDelete stored procedure.
        // Source: https://github.com/Azure/azure-cosmosdb-js-server/blob/master/samples/stored-procedures/bulkDelete.js
        String sprocLink = cosmosCollectionLink + BULK_DELETE_SPROC;
        try {
            ResourceResponse<StoredProcedure> spResponse = asyncDocumentClient.readStoredProcedure(sprocLink, null).toBlocking().single();
            if (spResponse.getResource() == null) {
                logger.error("Did not find stored procedure {}, falling back to individual record deletions", sprocLink);
            } else {
                logger.info("Found stored procedure {}, will use it with batch size {}.", sprocLink, purgeBatchSize);
                bulkDeleteEnabled = true;
            }
        } catch (Exception ex) {
            logger.error("Did not find stored procedure {}, falling back to individual record deletions", sprocLink, ex);
        }
    } else {
        logger.info("Cosmos purge batch size = 1 ==> disabling bulk deletes.");
        bulkDeleteEnabled = false;
    }
}
Also used : StoredProcedure(com.microsoft.azure.cosmosdb.StoredProcedure) RequestOptions(com.microsoft.azure.cosmosdb.RequestOptions) DocumentCollection(com.microsoft.azure.cosmosdb.DocumentCollection) DocumentClientException(com.microsoft.azure.cosmosdb.DocumentClientException) CloudStorageException(com.github.ambry.cloud.CloudStorageException)

Example 9 with RequestOptions

use of com.microsoft.azure.cosmosdb.RequestOptions in project ambry by linkedin.

the class CloudTestUtil method cleanupPartition.

/**
 * Cleanup the specified partition in azure by deleting all the blobs of the partition.
 * @param azureCloudConfig Properties containing the credentials needed for connection to azure.
 * @param partitionId partition to be deleted.
 */
static void cleanupPartition(AzureCloudConfig azureCloudConfig, PartitionId partitionId) {
    ConnectionPolicy connectionPolicy = new ConnectionPolicy();
    AsyncDocumentClient asyncDocumentClient = new AsyncDocumentClient.Builder().withServiceEndpoint(azureCloudConfig.cosmosEndpoint).withMasterKeyOrResourceToken(azureCloudConfig.cosmosKey).withConnectionPolicy(connectionPolicy).withConsistencyLevel(ConsistencyLevel.Session).build();
    SqlQuerySpec sqlQuerySpec = new SqlQuerySpec("select * from c where c.partitionId=\"" + partitionId.toPathString() + "\"");
    FeedOptions feedOptions = new FeedOptions();
    feedOptions.setPartitionKey(new PartitionKey(partitionId.toPathString()));
    Iterator<FeedResponse<Document>> iterator = asyncDocumentClient.queryDocuments(azureCloudConfig.cosmosCollectionLink, sqlQuerySpec, feedOptions).toBlocking().getIterator();
    RequestOptions requestOptions = new RequestOptions();
    requestOptions.setPartitionKey(new PartitionKey(partitionId.toPathString()));
    while (iterator.hasNext()) {
        FeedResponse<Document> response = iterator.next();
        response.getResults().forEach(document -> asyncDocumentClient.deleteDocument(azureCloudConfig.cosmosCollectionLink + "/docs/" + document.getId(), requestOptions).toBlocking().single());
    }
}
Also used : SqlQuerySpec(com.microsoft.azure.cosmosdb.SqlQuerySpec) RequestOptions(com.microsoft.azure.cosmosdb.RequestOptions) FeedOptions(com.microsoft.azure.cosmosdb.FeedOptions) PartitionKey(com.microsoft.azure.cosmosdb.PartitionKey) FeedResponse(com.microsoft.azure.cosmosdb.FeedResponse) ConnectionPolicy(com.microsoft.azure.cosmosdb.ConnectionPolicy) Document(com.microsoft.azure.cosmosdb.Document) AsyncDocumentClient(com.microsoft.azure.cosmosdb.rx.AsyncDocumentClient)

Aggregations

RequestOptions (com.microsoft.azure.cosmosdb.RequestOptions)9 DocumentClientException (com.microsoft.azure.cosmosdb.DocumentClientException)6 CloudStorageException (com.github.ambry.cloud.CloudStorageException)5 Document (com.microsoft.azure.cosmosdb.Document)5 PartitionKey (com.microsoft.azure.cosmosdb.PartitionKey)5 ConnectionPolicy (com.microsoft.azure.cosmosdb.ConnectionPolicy)4 AsyncDocumentClient (com.microsoft.azure.cosmosdb.rx.AsyncDocumentClient)4 AccessCondition (com.microsoft.azure.cosmosdb.AccessCondition)3 DocumentCollection (com.microsoft.azure.cosmosdb.DocumentCollection)3 FeedOptions (com.microsoft.azure.cosmosdb.FeedOptions)3 FeedResponse (com.microsoft.azure.cosmosdb.FeedResponse)3 RetryOptions (com.microsoft.azure.cosmosdb.RetryOptions)3 SqlQuerySpec (com.microsoft.azure.cosmosdb.SqlQuerySpec)3 StoredProcedure (com.microsoft.azure.cosmosdb.StoredProcedure)3 SecretClient (com.azure.security.keyvault.secrets.SecretClient)2 SecretClientBuilder (com.azure.security.keyvault.secrets.SecretClientBuilder)2 Timer (com.codahale.metrics.Timer)2 Container (com.github.ambry.account.Container)2 CloudBlobMetadata (com.github.ambry.cloud.CloudBlobMetadata)2 CloudRequestAgent (com.github.ambry.cloud.CloudRequestAgent)2