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

Example 1 with StoredProcedureResponse

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

the class AzureTestUtils method getMockBulkDeleteResponse.

static Observable<StoredProcedureResponse> getMockBulkDeleteResponse(int count) {
    Observable<StoredProcedureResponse> mockObservable = mock(Observable.class);
    BlockingObservable<StoredProcedureResponse> mockBlockingObservable = mock(BlockingObservable.class);
    when(mockObservable.toBlocking()).thenReturn(mockBlockingObservable);
    StoredProcedureResponse mockResponse = mock(StoredProcedureResponse.class);
    when(mockBlockingObservable.single()).thenReturn(mockResponse);
    Document responseDoc = new Document();
    responseDoc.set(CosmosDataAccessor.PROPERTY_CONTINUATION, "false");
    responseDoc.set(CosmosDataAccessor.PROPERTY_DELETED, String.valueOf(count));
    when(mockResponse.getResponseAsDocument()).thenReturn(responseDoc);
    return mockObservable;
}

Example 2 with StoredProcedureResponse

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

the class AzureStorageCompactorTest method buildCompactor.

private void buildCompactor(Properties configProps) throws Exception {
    CloudConfig cloudConfig = new CloudConfig(new VerifiableProperties(configProps));
    VcrMetrics vcrMetrics = new VcrMetrics(new MetricRegistry());
    azureBlobDataAccessor = new AzureBlobDataAccessor(mockServiceClient, mockBlobBatchClient, clusterName, azureMetrics, new AzureCloudConfig(new VerifiableProperties(configProps)));
    cosmosDataAccessor = new CosmosDataAccessor(mockumentClient, collectionLink, cosmosDeletedContainerCollectionLink, vcrMetrics, azureMetrics);
    azureStorageCompactor = new AzureStorageCompactor(azureBlobDataAccessor, cosmosDataAccessor, cloudConfig, vcrMetrics, azureMetrics);
    // Mocks for getDeadBlobs query
    List<Document> docList = new ArrayList<>();
    for (int j = 0; j < numBlobsPerQuery; j++) {
        BlobId blobId = generateBlobId();
        CloudBlobMetadata inputMetadata = new CloudBlobMetadata(blobId, testTime, Utils.Infinite_Time, blobSize, CloudBlobMetadata.EncryptionOrigin.NONE);
        blobMetadataList.add(inputMetadata);
        docList.add(AzureTestUtils.createDocumentFromCloudBlobMetadata(inputMetadata));
    }
    Observable<FeedResponse<Document>> mockResponse = mock(Observable.class);
    mockObservableForQuery(docList, mockResponse);
    when(mockumentClient.queryDocuments(anyString(), any(SqlQuerySpec.class), any(FeedOptions.class))).thenReturn(mockResponse);
    // Mocks for purge
    BlobBatch mockBatch = mock(BlobBatch.class);
    when(mockBlobBatchClient.getBlobBatch()).thenReturn(mockBatch);
    Response<Void> okResponse = mock(Response.class);
    when(okResponse.getStatusCode()).thenReturn(202);
    when(mockBatch.deleteBlob(anyString(), anyString())).thenReturn(okResponse);
    Observable<StoredProcedureResponse> mockBulkDeleteResponse = getMockBulkDeleteResponse(1);
    when(mockumentClient.executeStoredProcedure(anyString(), any(RequestOptions.class), any())).thenReturn(mockBulkDeleteResponse);
    String checkpointJson = objectMapper.writeValueAsString(AzureStorageCompactor.emptyCheckpoints);
    mockCheckpointDownload(true, checkpointJson);
}

Example 3 with StoredProcedureResponse

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