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

Example 1 with PartitionKey

use of com.microsoft.azure.cosmosdb.PartitionKey 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 PartitionKey

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

the class CosmosDataAccessor method getDeadBlobs.

/**
 * Get the list of blobs in the specified partition that have been deleted or expired for at least the
 * configured retention period.
 * @param partitionPath the partition to query.
 * @param fieldName the field name to query on. Allowed values are {@link CloudBlobMetadata#FIELD_DELETION_TIME} and
 *                  {@link CloudBlobMetadata#FIELD_EXPIRATION_TIME}.
 * @param startTime the start of the query time range.
 * @param endTime the end of the query time range.
 * @param maxEntries the max number of metadata records to return.
 * @return a List of {@link CloudBlobMetadata} referencing the dead blobs found.
 * @throws DocumentClientException
 */
List<CloudBlobMetadata> getDeadBlobs(String partitionPath, String fieldName, long startTime, long endTime, int maxEntries) throws DocumentClientException {
    String deadBlobsQuery;
    if (fieldName.equals(CloudBlobMetadata.FIELD_DELETION_TIME)) {
        deadBlobsQuery = DELETED_BLOBS_QUERY;
    } else if (fieldName.equals(CloudBlobMetadata.FIELD_EXPIRATION_TIME)) {
        deadBlobsQuery = EXPIRED_BLOBS_QUERY;
    } else {
        throw new IllegalArgumentException("Invalid field: " + fieldName);
    }
    SqlQuerySpec querySpec = new SqlQuerySpec(deadBlobsQuery, new SqlParameterCollection(new SqlParameter(LIMIT_PARAM, maxEntries), new SqlParameter(START_TIME_PARAM, startTime), new SqlParameter(END_TIME_PARAM, endTime)));
    FeedOptions feedOptions = new FeedOptions();
    feedOptions.setMaxItemCount(maxEntries);
    feedOptions.setResponseContinuationTokenLimitInKb(continuationTokenLimitKb);
    feedOptions.setPartitionKey(new PartitionKey(partitionPath));
    try {
        Iterator<FeedResponse<Document>> iterator = executeCosmosQuery(partitionPath, querySpec, feedOptions, azureMetrics.deadBlobsQueryTime).getIterator();
        List<CloudBlobMetadata> deadBlobsList = new ArrayList<>();
        double requestCharge = 0.0;
        while (iterator.hasNext()) {
            FeedResponse<Document> response = iterator.next();
            requestCharge += response.getRequestCharge();
            response.getResults().iterator().forEachRemaining(doc -> deadBlobsList.add(createMetadataFromDocument(doc)));
        }
        if (requestCharge >= requestChargeThreshold) {
            logger.info("Dead blobs query partition {} endTime {} request charge {} for {} records", partitionPath, new Date(endTime), requestCharge, deadBlobsList.size());
        }
        return deadBlobsList;
    } catch (RuntimeException rex) {
        if (rex.getCause() instanceof DocumentClientException) {
            logger.warn("Dead blobs query {} partition {} got {}", deadBlobsQuery, partitionPath, ((DocumentClientException) rex.getCause()).getStatusCode());
            throw (DocumentClientException) rex.getCause();
        }
        throw rex;
    }
}

Example 3 with PartitionKey

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

the class CosmosDataAccessor method getContainerBlobs.

/**
 * Get the list of blobs in the specified partition that belong to the specified container.
 * @param partitionPath the partition to query.
 * @param accountId account id of the container.
 * @param containerId container id of the container.
 * @param queryLimit max number of blobs to return.
 * @return a List of {@link CloudBlobMetadata} referencing the blobs belonging to the deprecated containers.
 * @throws DocumentClientException in case of any error.
 */
List<CloudBlobMetadata> getContainerBlobs(String partitionPath, short accountId, short containerId, int queryLimit) throws DocumentClientException {
    SqlQuerySpec querySpec = new SqlQuerySpec(CONTAINER_BLOBS_QUERY, new SqlParameterCollection(new SqlParameter(LIMIT_PARAM, queryLimit), new SqlParameter(CONTAINER_ID_PARAM, containerId), new SqlParameter(ACCOUNT_ID_PARAM, accountId)));
    FeedOptions feedOptions = new FeedOptions();
    feedOptions.setMaxItemCount(queryLimit);
    feedOptions.setResponseContinuationTokenLimitInKb(continuationTokenLimitKb);
    feedOptions.setPartitionKey(new PartitionKey(partitionPath));
    try {
        Iterator<FeedResponse<Document>> iterator = executeCosmosQuery(partitionPath, querySpec, feedOptions, azureMetrics.deletedContainerBlobsQueryTime).getIterator();
        List<CloudBlobMetadata> containerBlobsList = new ArrayList<>();
        double requestCharge = 0.0;
        while (iterator.hasNext()) {
            FeedResponse<Document> response = iterator.next();
            requestCharge += response.getRequestCharge();
            response.getResults().iterator().forEachRemaining(doc -> containerBlobsList.add(createMetadataFromDocument(doc)));
        }
        if (requestCharge >= requestChargeThreshold) {
            logger.info("Deleted container blobs query partition {} containerId {} accountId {} request charge {} for {} records", partitionPath, containerId, accountId, requestCharge, containerBlobsList.size());
        }
        return containerBlobsList;
    } catch (RuntimeException rex) {
        if (rex.getCause() instanceof DocumentClientException) {
            logger.warn("Dead blobs query {} partition {} got {}", querySpec.getQueryText(), partitionPath, ((DocumentClientException) rex.getCause()).getStatusCode());
            throw (DocumentClientException) rex.getCause();
        }
        throw rex;
    }
}

Example 4 with PartitionKey

use of com.microsoft.azure.cosmosdb.PartitionKey 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());
    }
}

Example 5 with PartitionKey

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

the class CosmosDataAccessor method queryChangeFeed.

/**
 * Query Cosmos change feed to get the next set of {@code CloudBlobMetadata} objects in specified {@code partitionPath}
 * after {@code requestContinationToken}, capped by specified {@code maxFeedSize} representing the max number of items to
 * be queried from the change feed.
 * @param requestContinuationToken Continuation token after which change feed is requested.
 * @param maxFeedSize max item count to be requested in the feed query.
 * @param changeFeed {@link CloudBlobMetadata} {@code List} to be populated with the next set of entries returned by change feed query.
 * @param partitionPath partition for which the change feed is requested.
 * @param timer the {@link Timer} to use to record query time (excluding waiting).
 * @return next continuation token.
 * @throws DocumentClientException
 */
public String queryChangeFeed(String requestContinuationToken, int maxFeedSize, List<CloudBlobMetadata> changeFeed, String partitionPath, Timer timer) throws DocumentClientException {
    azureMetrics.changeFeedQueryCount.inc();
    ChangeFeedOptions changeFeedOptions = new ChangeFeedOptions();
    changeFeedOptions.setPartitionKey(new PartitionKey(partitionPath));
    changeFeedOptions.setMaxItemCount(maxFeedSize);
    if (Utils.isNullOrEmpty(requestContinuationToken)) {
        changeFeedOptions.setStartFromBeginning(true);
    } else {
        changeFeedOptions.setRequestContinuation(requestContinuationToken);
    }
    try {
        FeedResponse<Document> feedResponse = executeCosmosChangeFeedQuery(changeFeedOptions, timer);
        feedResponse.getResults().stream().map(this::createMetadataFromDocument).forEach(changeFeed::add);
        return feedResponse.getResponseContinuation();
    } catch (RuntimeException rex) {
        azureMetrics.changeFeedQueryFailureCount.inc();
        if (rex.getCause() instanceof DocumentClientException) {
            throw (DocumentClientException) rex.getCause();
        }
        throw rex;
    } catch (Exception ex) {
        azureMetrics.changeFeedQueryFailureCount.inc();
        throw ex;
    }
}