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

Example 1 with Document

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

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

the class AzureCloudDestinationTest method testFindEntriesSinceWithUniqueUpdateTimes.

/**
 * Test findEntriesSince with all entries having unique updateTimes.
 * @throws Exception
 */
private void testFindEntriesSinceWithUniqueUpdateTimes(AzureCloudDestination azureDest) throws Exception {
    long chunkSize = 110000;
    // between 9 and 10 chunks
    long maxTotalSize = 1000000;
    long startTime = System.currentTimeMillis() - TimeUnit.DAYS.toMillis(1);
    int totalBlobs = 20;
    // create metadata list where total size > maxTotalSize
    List<Document> docList = new ArrayList<>();
    List<String> blobIdList = new ArrayList<>();
    for (int j = 0; j < totalBlobs; j++) {
        BlobId blobId = generateBlobId();
        blobIdList.add(blobId.getID());
        CloudBlobMetadata inputMetadata = new CloudBlobMetadata(blobId, creationTime, Utils.Infinite_Time, chunkSize, CloudBlobMetadata.EncryptionOrigin.NONE);
        inputMetadata.setUploadTime(startTime + j);
        docList.add(AzureTestUtils.createDocumentFromCloudBlobMetadata(inputMetadata, startTime + j));
    }
    Observable<FeedResponse<Document>> mockResponse = mock(Observable.class);
    mockObservableForQuery(docList, mockResponse);
    when(mockumentClient.queryDocuments(anyString(), any(SqlQuerySpec.class), any(FeedOptions.class))).thenReturn(mockResponse);
    CosmosUpdateTimeFindToken findToken = new CosmosUpdateTimeFindToken();
    // Run the query
    FindResult findResult = azureDest.findEntriesSince(blobId.getPartition().toPathString(), findToken, maxTotalSize);
    List<CloudBlobMetadata> firstResult = findResult.getMetadataList();
    findToken = (CosmosUpdateTimeFindToken) findResult.getUpdatedFindToken();
    assertEquals("Did not get expected doc count", maxTotalSize / chunkSize, firstResult.size());
    docList = docList.subList(firstResult.size(), docList.size());
    assertEquals("Find token has wrong last update time", findToken.getLastUpdateTime(), firstResult.get(firstResult.size() - 1).getLastUpdateTime());
    assertEquals("Find token has wrong lastUpdateTimeReadBlobIds", findToken.getLastUpdateTimeReadBlobIds(), new HashSet<>(Collections.singletonList(firstResult.get(firstResult.size() - 1).getId())));
    mockObservableForQuery(docList, mockResponse);
    findResult = azureDest.findEntriesSince(blobId.getPartition().toPathString(), findToken, maxTotalSize);
    List<CloudBlobMetadata> secondResult = findResult.getMetadataList();
    findToken = (CosmosUpdateTimeFindToken) findResult.getUpdatedFindToken();
    assertEquals("Unexpected doc count", maxTotalSize / chunkSize, secondResult.size());
    assertEquals("Unexpected first blobId", blobIdList.get(firstResult.size()), secondResult.get(0).getId());
    mockObservableForQuery(docList, mockResponse);
    findResult = azureDest.findEntriesSince(blobId.getPartition().toPathString(), findToken, chunkSize / 2);
    // Rerun with max size below blob size, and make sure it returns one result
    assertEquals("Expected one result", 1, findResult.getMetadataList().size());
}

Example 3 with Document

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

the class AzureCloudDestinationTest method testFindEntriesSinceWithNonUniqueUpdateTimes.

/**
 * Test findEntriesSince with entries having non unique updateTimes.
 * @throws Exception
 */
private void testFindEntriesSinceWithNonUniqueUpdateTimes(AzureCloudDestination azureDest) throws Exception {
    long chunkSize = 110000;
    // between 9 and 10 chunks
    long maxTotalSize = 1000000;
    long startTime = System.currentTimeMillis() - TimeUnit.DAYS.toMillis(1);
    int totalBlobs = 20;
    // create metadata list where total size > maxTotalSize
    List<Document> docList = new ArrayList<>();
    List<String> blobIdList = new ArrayList<>();
    for (int j = 0; j < totalBlobs - 1; j++) {
        BlobId blobId = generateBlobId();
        blobIdList.add(blobId.getID());
        CloudBlobMetadata inputMetadata = new CloudBlobMetadata(blobId, creationTime, Utils.Infinite_Time, chunkSize, CloudBlobMetadata.EncryptionOrigin.NONE);
        docList.add(AzureTestUtils.createDocumentFromCloudBlobMetadata(inputMetadata, startTime));
    }
    BlobId blobId = generateBlobId();
    blobIdList.add(blobId.getID());
    CloudBlobMetadata inputMetadata = new CloudBlobMetadata(blobId, creationTime, Utils.Infinite_Time, chunkSize, CloudBlobMetadata.EncryptionOrigin.NONE);
    docList.add(AzureTestUtils.createDocumentFromCloudBlobMetadata(inputMetadata, startTime + 1));
    Observable<FeedResponse<Document>> mockResponse = mock(Observable.class);
    mockObservableForQuery(docList, mockResponse);
    when(mockumentClient.queryDocuments(anyString(), any(SqlQuerySpec.class), any(FeedOptions.class))).thenReturn(mockResponse);
    CosmosUpdateTimeFindToken findToken = new CosmosUpdateTimeFindToken();
    // Run the query
    FindResult findResult = azureDest.findEntriesSince(blobId.getPartition().toPathString(), findToken, maxTotalSize);
    List<CloudBlobMetadata> firstResult = findResult.getMetadataList();
    findToken = (CosmosUpdateTimeFindToken) findResult.getUpdatedFindToken();
    assertEquals("Did not get expected doc count", maxTotalSize / chunkSize, firstResult.size());
    assertEquals("Find token has wrong last update time", findToken.getLastUpdateTime(), firstResult.get(firstResult.size() - 1).getLastUpdateTime());
    Set<String> resultBlobIdSet = firstResult.stream().map(CloudBlobMetadata::getId).collect(Collectors.toSet());
    assertEquals("Find token has wrong lastUpdateTimeReadBlobIds", findToken.getLastUpdateTimeReadBlobIds(), resultBlobIdSet);
    mockObservableForQuery(docList, mockResponse);
    findResult = azureDest.findEntriesSince(blobId.getPartition().toPathString(), findToken, maxTotalSize);
    List<CloudBlobMetadata> secondResult = findResult.getMetadataList();
    CosmosUpdateTimeFindToken secondFindToken = (CosmosUpdateTimeFindToken) findResult.getUpdatedFindToken();
    assertEquals("Unexpected doc count", maxTotalSize / chunkSize, secondResult.size());
    assertEquals("Unexpected first blobId", blobIdList.get(firstResult.size()), secondResult.get(0).getId());
    assertEquals("Find token has wrong last update time", secondFindToken.getLastUpdateTime(), firstResult.get(firstResult.size() - 1).getLastUpdateTime());
    resultBlobIdSet.addAll(secondResult.stream().map(CloudBlobMetadata::getId).collect(Collectors.toSet()));
    assertEquals("Find token has wrong lastUpdateTimeReadBlobIds", secondFindToken.getLastUpdateTimeReadBlobIds(), resultBlobIdSet);
    mockObservableForQuery(docList, mockResponse);
    // Rerun with max size below blob size, and make sure it returns one result
    findResult = azureDest.findEntriesSince(blobId.getPartition().toPathString(), findToken, chunkSize / 2);
    List<CloudBlobMetadata> finalResult = findResult.getMetadataList();
    assertEquals("Expected one result", 1, finalResult.size());
    mockObservableForQuery(docList, mockResponse);
    // Rerun final time, and make sure that it returns all the remaining blobs
    findResult = azureDest.findEntriesSince(blobId.getPartition().toPathString(), secondFindToken, maxTotalSize);
    List<CloudBlobMetadata> thirdResult = findResult.getMetadataList();
    CosmosUpdateTimeFindToken thirdFindToken = (CosmosUpdateTimeFindToken) findResult.getUpdatedFindToken();
    assertEquals("Unexpected doc count", totalBlobs - (firstResult.size() + secondResult.size()), thirdResult.size());
    assertEquals("Unexpected first blobId", blobIdList.get(firstResult.size() + secondResult.size()), thirdResult.get(0).getId());
    assertEquals("Find token has wrong last update time", thirdFindToken.getLastUpdateTime(), startTime + 1);
    assertEquals("Find token has wrong lastUpdateTimeReadBlobIds", thirdFindToken.getLastUpdateTimeReadBlobIds(), new HashSet<>(Collections.singletonList(thirdResult.get(thirdResult.size() - 1).getId())));
}

Example 4 with Document

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

the class AzureCloudDestinationTest method setup.

@Before
public void setup() throws Exception {
    long partition = 666;
    PartitionId partitionId = new MockPartitionId(partition, MockClusterMap.DEFAULT_PARTITION_CLASS);
    blobId = new BlobId(BLOB_ID_V6, BlobIdType.NATIVE, dataCenterId, accountId, containerId, partitionId, false, BlobDataType.DATACHUNK);
    CloudBlobMetadata blobMetadata = new CloudBlobMetadata(blobId, 0, Utils.Infinite_Time, blobSize, CloudBlobMetadata.EncryptionOrigin.NONE);
    mockServiceClient = mock(BlobServiceClient.class);
    mockBlobBatchClient = mock(BlobBatchClient.class);
    mockBlockBlobClient = AzureBlobDataAccessorTest.setupMockBlobClient(mockServiceClient);
    mockBlobExistence(false);
    mockumentClient = mock(AsyncDocumentClient.class);
    Observable<ResourceResponse<Document>> mockResponse = getMockedObservableForSingleResource(blobMetadata);
    when(mockumentClient.readDocument(anyString(), any(RequestOptions.class))).thenReturn(mockResponse);
    when(mockumentClient.upsertDocument(anyString(), any(Object.class), any(RequestOptions.class), anyBoolean())).thenReturn(mockResponse);
    when(mockumentClient.replaceDocument(any(Document.class), any(RequestOptions.class))).thenReturn(mockResponse);
    when(mockumentClient.deleteDocument(anyString(), any(RequestOptions.class))).thenReturn(mockResponse);
    configProps.setProperty(AzureCloudConfig.AZURE_STORAGE_CONNECTION_STRING, storageConnection);
    configProps.setProperty(AzureCloudConfig.COSMOS_ENDPOINT, "http://ambry.beyond-the-cosmos.com:443");
    configProps.setProperty(AzureCloudConfig.COSMOS_COLLECTION_LINK, "ambry/metadata");
    configProps.setProperty(AzureCloudConfig.COSMOS_DELETED_CONTAINER_COLLECTION_LINK, "ambry/deletedContainer");
    configProps.setProperty(AzureCloudConfig.COSMOS_KEY, "cosmos-key");
    configProps.setProperty("clustermap.cluster.name", "main");
    configProps.setProperty("clustermap.datacenter.name", "uswest");
    configProps.setProperty("clustermap.host.name", "localhost");
    configProps.setProperty(AzureCloudConfig.AZURE_STORAGE_AUTHORITY, "https://login.microsoftonline.com/test-account/");
    configProps.setProperty(AzureCloudConfig.AZURE_STORAGE_CLIENTID, "client-id");
    configProps.setProperty(AzureCloudConfig.AZURE_STORAGE_SECRET, "client-secret");
    configProps.setProperty(AzureCloudConfig.AZURE_STORAGE_ENDPOINT, "https://azure_storage.blob.core.windows.net");
    configProps.setProperty(AzureCloudConfig.AZURE_STORAGE_CLIENT_CLASS, "com.github.ambry.cloud.azure.ConnectionStringBasedStorageClient");
    vcrMetrics = new VcrMetrics(new MetricRegistry());
    azureMetrics = new AzureMetrics(new MetricRegistry());
    clusterMap = mock(ClusterMap.class);
    azureDest = new AzureCloudDestination(mockServiceClient, mockBlobBatchClient, mockumentClient, "foo", "bar", clusterName, azureMetrics, defaultAzureReplicationFeedType, clusterMap, false, configProps);
}

Example 5 with Document

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