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

Example 16 with Document

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

Example 17 with Document

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

Example 18 with Document

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

the class CosmosDataAccessor method queryMetadata.

/**
 * Get the list of blobs in the specified partition matching the specified DocumentDB query spec.
 * @param partitionPath the partition to query.
 * @param querySpec the DocumentDB {@link SqlQuerySpec} to execute.
 * @param timer the {@link Timer} to use to record query time (excluding waiting).
 * @return a List of {@link CloudBlobMetadata} referencing the matching blobs.
 */
List<CloudBlobMetadata> queryMetadata(String partitionPath, SqlQuerySpec querySpec, Timer timer) throws DocumentClientException {
    FeedOptions feedOptions = new FeedOptions();
    // TODO: set maxItemCount
    feedOptions.setResponseContinuationTokenLimitInKb(continuationTokenLimitKb);
    feedOptions.setPartitionKey(new PartitionKey(partitionPath));
    // TODO: consolidate error count here
    try {
        Iterator<FeedResponse<Document>> iterator = executeCosmosQuery(partitionPath, querySpec, feedOptions, timer).getIterator();
        List<CloudBlobMetadata> metadataList = new ArrayList<>();
        double requestCharge = 0.0;
        while (iterator.hasNext()) {
            FeedResponse<Document> response = iterator.next();
            requestCharge += response.getRequestCharge();
            response.getResults().iterator().forEachRemaining(doc -> metadataList.add(createMetadataFromDocument(doc)));
        }
        if (requestCharge >= requestChargeThreshold) {
            logger.info("Query partition {} request charge {} for {} records", partitionPath, requestCharge, metadataList.size());
        }
        return metadataList;
    } catch (RuntimeException rex) {
        if (rex.getCause() instanceof DocumentClientException) {
            logger.warn("Query {} on partition {} got {}", querySpec.getQueryText(), partitionPath, ((DocumentClientException) rex.getCause()).getStatusCode());
            throw (DocumentClientException) rex.getCause();
        }
        throw rex;
    }
}

Example 19 with Document

use of com.microsoft.azure.cosmosdb.Document 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 20 with Document

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

the class AzureCloudDestinationTest method testGetOneMetadata.

/**
 * Test to make sure that getting metadata for single blob calls ABS when not vcr and Cosmos when vcr.
 */
@Test
public void testGetOneMetadata() throws Exception {
    // 
    // Test 1: isVcr = false (already setup)
    // 
    // Get for existing blob
    Response<BlobProperties> mockResponse = mock(Response.class);
    BlobProperties mockProperties = mock(BlobProperties.class);
    CloudBlobMetadata blobMetadata = new CloudBlobMetadata(blobId, 0, -1, 0, null);
    Map<String, String> propertyMap = blobMetadata.toMap();
    when(mockProperties.getMetadata()).thenReturn(propertyMap);
    when(mockResponse.getValue()).thenReturn(mockProperties);
    when(mockBlockBlobClient.getPropertiesWithResponse(any(), any(), any())).thenReturn(mockResponse);
    List<BlobId> singleBlobList = Collections.singletonList(blobId);
    Map<String, CloudBlobMetadata> metadataMap = azureDest.getBlobMetadata(singleBlobList);
    assertEquals("Expected map of one", 1, metadataMap.size());
    verify(mockBlockBlobClient).getPropertiesWithResponse(any(), any(), any());
    verifyZeroInteractions(mockumentClient);
    // Get for nonexistent blob
    BlobStorageException ex = mockStorageException(BlobErrorCode.BLOB_NOT_FOUND);
    when(mockBlockBlobClient.getPropertiesWithResponse(any(), any(), any())).thenThrow(ex);
    metadataMap = azureDest.getBlobMetadata(singleBlobList);
    assertTrue("Expected empty map", metadataMap.isEmpty());
    verify(mockBlockBlobClient, times(2)).getPropertiesWithResponse(any(), any(), any());
    verifyZeroInteractions(mockumentClient);
    // 
    // Test 2: isVcr = true
    // 
    azureDest.close();
    azureDest = new AzureCloudDestination(mockServiceClient, mockBlobBatchClient, mockumentClient, "foo", "bar", clusterName, azureMetrics, defaultAzureReplicationFeedType, clusterMap, true, configProps);
    // Existing blob
    List<Document> docList = Collections.singletonList(createDocumentFromCloudBlobMetadata(blobMetadata));
    Observable<FeedResponse<Document>> feedResponse = mock(Observable.class);
    mockObservableForQuery(docList, feedResponse);
    when(mockumentClient.queryDocuments(anyString(), any(SqlQuerySpec.class), any(FeedOptions.class))).thenReturn(feedResponse);
    metadataMap = azureDest.getBlobMetadata(singleBlobList);
    assertEquals("Expected map of one", 1, metadataMap.size());
    verify(mockumentClient).queryDocuments(anyString(), any(SqlQuerySpec.class), any(FeedOptions.class));
    verify(mockBlockBlobClient, times(2)).getPropertiesWithResponse(any(), any(), any());
}