Examples with RequestOptions com.microsoft.azure.cosmosdb.RequestOptions used on opensource projects

Example 1 with RequestOptions

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

the class AzureContainerCompactorIntegrationTest method cleanup.

/**
 * Cleanup entries in the cosmos db container deletion table.
 * @throws DocumentClientException in case of any exception.
 */
private void cleanup() throws DocumentClientException {
    ConnectionPolicy connectionPolicy = new ConnectionPolicy();
    connectionPolicy.setRequestTimeoutInMillis(cloudConfig.cloudQueryRequestTimeout);
    // Note: retry decisions are made at CloudBlobStore level.  Configure Cosmos with no retries.
    RetryOptions noRetries = new RetryOptions();
    noRetries.setMaxRetryAttemptsOnThrottledRequests(0);
    connectionPolicy.setRetryOptions(noRetries);
    if (azureCloudConfig.cosmosDirectHttps) {
        logger.info("Using CosmosDB DirectHttps connection mode");
        connectionPolicy.setConnectionMode(ConnectionMode.Direct);
    }
    AsyncDocumentClient asyncDocumentClient = new AsyncDocumentClient.Builder().withServiceEndpoint(azureCloudConfig.cosmosEndpoint).withMasterKeyOrResourceToken(azureCloudConfig.cosmosKey).withConnectionPolicy(connectionPolicy).withConsistencyLevel(ConsistencyLevel.Session).build();
    Set<CosmosContainerDeletionEntry> entries = cloudDestination.getCosmosDataAccessor().getDeprecatedContainers(100);
    AtomicBoolean error = new AtomicBoolean(false);
    while (!entries.isEmpty() && !error.get()) {
        entries.stream().forEach(entry -> {
            try {
                RequestOptions requestOptions = new RequestOptions();
                requestOptions.setPartitionKey(new PartitionKey(entry.getId()));
                cloudRequestAgent.doWithRetries(() -> CosmosDataAccessor.executeCosmosAction(() -> asyncDocumentClient.deleteDocument(azureCloudConfig.cosmosDeletedContainerCollectionLink + "/docs/" + entry.getId(), requestOptions).toBlocking().single(), null), "Test Cleanup", entry.getId());
            } catch (CloudStorageException ex) {
                logger.warn("Failed to delete container deprecation entry {}. Unable to cleanup", entry);
                error.set(true);
            }
        });
        entries = cloudDestination.getCosmosDataAccessor().getDeprecatedContainers(100);
    }
}

Example 2 with RequestOptions

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

the class CosmosDataAccessor method updateContainerDeletionEntry.

/**
 * Update the container deletion entry document in the CosmosDB collection.
 * @param containerId the container id for which document is replaced.
 * @param accountId the account id for which document is replaced.
 * @param updateFields {@link BiConsumer} object to use as callback to update the required fields.
 * @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> updateContainerDeletionEntry(short containerId, short accountId, BiConsumer<Document, AtomicBoolean> updateFields) throws DocumentClientException {
    // Read the existing record
    String id = CosmosContainerDeletionEntry.generateContainerDeletionEntryId(accountId, containerId);
    String docLink = getContainerDeletionEntryDocumentLink(id);
    RequestOptions options = getRequestOptions(id);
    ResourceResponse<Document> readResponse = executeCosmosAction(() -> asyncDocumentClient.readDocument(docLink, options).toBlocking().single(), azureMetrics.continerDeletionEntryReadTime);
    Document doc = readResponse.getResource();
    AtomicBoolean fieldsChanged = new AtomicBoolean(false);
    updateFields.accept(doc, fieldsChanged);
    if (!fieldsChanged.get()) {
        logger.debug("No change in value for container deletion entry {}", doc.toJson());
        return null;
    }
    // For testing conflict handling
    if (updateCallback != null) {
        try {
            updateCallback.call();
        } catch (Exception ex) {
            logger.error("Error in update callback", ex);
        }
    }
    // 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;
    }
}

Example 3 with RequestOptions

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

the class CosmosDataAccessor method deleteMetadata.

/**
 * Delete the blob metadata document in the CosmosDB collection, if it exists.
 * @param blobMetadata the blob metadata document to delete.
 * @return {@code true} if the record was deleted, {@code false} if it was not found.
 * @throws DocumentClientException if the operation failed.
 */
boolean deleteMetadata(CloudBlobMetadata blobMetadata) throws DocumentClientException {
    String docLink = getDocumentLink(blobMetadata.getId());
    RequestOptions options = getRequestOptions(blobMetadata.getPartitionId());
    try {
        // Note: not timing here since bulk deletions are timed.
        executeCosmosAction(() -> asyncDocumentClient.deleteDocument(docLink, options).toBlocking().single(), null);
        return true;
    } catch (DocumentClientException dex) {
        if (dex.getStatusCode() == HttpURLConnection.HTTP_NOT_FOUND) {
            // Can happen on retry
            logger.debug("Could not find metadata for blob {} to delete", blobMetadata.getId());
            return false;
        } else {
            throw dex;
        }
    }
}

Example 4 with RequestOptions

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

the class CosmosDataAccessor method getMetadataOrNull.

/**
 * Get the metadata record for a single blob.
 * @param blobId the blob to read.
 * @return the {@link CloudBlobMetadata} for the blob if it is found, otherwise null.
 * @throws DocumentClientException on any other error.
 */
CloudBlobMetadata getMetadataOrNull(BlobId blobId) throws DocumentClientException {
    String docLink = getDocumentLink(blobId.getID());
    RequestOptions options = getRequestOptions(blobId.getPartition().toPathString());
    try {
        ResourceResponse<Document> readResponse = executeCosmosAction(() -> asyncDocumentClient.readDocument(docLink, options).toBlocking().single(), azureMetrics.documentReadTime);
        return createMetadataFromDocument(readResponse.getResource());
    } catch (DocumentClientException dex) {
        if (dex.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) {
            return null;
        } else {
            throw dex;
        }
    }
}

Example 5 with RequestOptions

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

the class CosmosDataAccessor method bulkDeleteMetadata.

/**
 * Delete the blob metadata documents from CosmosDB using the BulkDelete stored procedure.
 * @param blobMetadataList the list of blob metadata documents to delete.
 * @return the number of documents deleted.
 * @throws DocumentClientException if the operation failed.
 */
private int bulkDeleteMetadata(List<CloudBlobMetadata> blobMetadataList) throws DocumentClientException {
    String partitionPath = blobMetadataList.get(0).getPartitionId();
    RequestOptions options = getRequestOptions(partitionPath);
    // stored proc link provided in config.  Test for it at startup and use if available.
    String quotedBlobIds = blobMetadataList.stream().map(metadata -> '"' + metadata.getId() + '"').collect(Collectors.joining(","));
    String query = String.format(BULK_DELETE_QUERY, quotedBlobIds);
    String sprocLink = cosmosCollectionLink + BULK_DELETE_SPROC;
    boolean more = true;
    int deleteCount = 0;
    double requestCharge = 0;
    try {
        while (more) {
            StoredProcedureResponse response = asyncDocumentClient.executeStoredProcedure(sprocLink, options, new String[] { query }).toBlocking().single();
            requestCharge += response.getRequestCharge();
            Document responseDoc = response.getResponseAsDocument();
            more = responseDoc.getBoolean(PROPERTY_CONTINUATION);
            deleteCount += responseDoc.getInt(PROPERTY_DELETED);
        }
        if (requestCharge >= requestChargeThreshold) {
            logger.info("Bulk delete partition {} request charge {} for {} records", partitionPath, requestCharge, deleteCount);
        }
        return deleteCount;
    } catch (RuntimeException rex) {
        if (rex.getCause() instanceof DocumentClientException) {
            throw (DocumentClientException) rex.getCause();
        } else {
            throw rex;
        }
    }
}