Search in sources :

Example 1 with BlobContainerAsyncClient

use of com.azure.storage.blob.BlobContainerAsyncClient in project samza by apache.

the class TestAzureBlobAvroWriter method testRecordLimitExceeded.

@Test
public void testRecordLimitExceeded() throws Exception {
    String blobUrlPrefix = "test";
    String blobNameRegex = "test/[0-9]{4}/[0-9]{2}/[0-9]{2}/[0-9]{2}/[0-9]{2}-[0-9]{2}-.{8}.avro.gz";
    AzureBlobWriterMetrics mockMetrics = mock(AzureBlobWriterMetrics.class);
    long maxBlobSize = AzureBlobAvroWriter.DATAFILEWRITER_OVERHEAD + 1000;
    long maxRecordsPerBlob = 10;
    BlobContainerAsyncClient mockContainerClient = PowerMockito.mock(BlobContainerAsyncClient.class);
    azureBlobAvroWriter = spy(new AzureBlobAvroWriter(mockContainerClient, mockMetrics, threadPool, THRESHOLD, 60000, blobUrlPrefix, null, null, null, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, maxBlobSize, maxRecordsPerBlob, mockCompression, true));
    DataFileWriter mockDataFileWriter1 = mock(DataFileWriter.class);
    PowerMockito.whenNew(DataFileWriter.class).withAnyArguments().thenReturn(mockDataFileWriter1);
    BlobAsyncClient mockBlobAsyncClient1 = mock(BlobAsyncClient.class);
    doReturn(mockBlobAsyncClient1).when(mockContainerClient).getBlobAsyncClient(Matchers.matches(blobNameRegex));
    BlockBlobAsyncClient mockBlockBlobAsyncClient1 = mock(BlockBlobAsyncClient.class);
    doReturn(mockBlockBlobAsyncClient1).when(mockBlobAsyncClient1).getBlockBlobAsyncClient();
    AzureBlobOutputStream mockAzureBlobOutputStream1 = mock(AzureBlobOutputStream.class);
    PowerMockito.whenNew(AzureBlobOutputStream.class).withArguments(mockBlockBlobAsyncClient1, threadPool, mockMetrics, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, (long) 60000, THRESHOLD, mockCompression).thenReturn(mockAzureBlobOutputStream1);
    when(mockAzureBlobOutputStream1.getSize()).thenReturn((long) 1);
    for (int i = 0; i < maxRecordsPerBlob; i++) {
        azureBlobAvroWriter.write(ome);
    }
    OutgoingMessageEnvelope ome2 = createOME("Topic2");
    DataFileWriter mockDataFileWriter2 = mock(DataFileWriter.class);
    PowerMockito.whenNew(DataFileWriter.class).withAnyArguments().thenReturn(mockDataFileWriter2);
    BlobAsyncClient mockBlobAsyncClient2 = mock(BlobAsyncClient.class);
    doReturn(mockBlobAsyncClient2).when(mockContainerClient).getBlobAsyncClient(Matchers.matches(blobNameRegex));
    BlockBlobAsyncClient mockBlockBlobAsyncClient2 = mock(BlockBlobAsyncClient.class);
    doReturn(mockBlockBlobAsyncClient2).when(mockBlobAsyncClient2).getBlockBlobAsyncClient();
    AzureBlobOutputStream mockAzureBlobOutputStream2 = mock(AzureBlobOutputStream.class);
    PowerMockito.whenNew(AzureBlobOutputStream.class).withArguments(mockBlockBlobAsyncClient2, threadPool, mockMetrics, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, (long) 60000, THRESHOLD, mockCompression).thenReturn(mockAzureBlobOutputStream2);
    when(mockAzureBlobOutputStream2.getSize()).thenReturn((long) 1);
    azureBlobAvroWriter.write(ome2);
    ArgumentCaptor<String> argument = ArgumentCaptor.forClass(String.class);
    verify(mockContainerClient, times(2)).getBlobAsyncClient(argument.capture());
    argument.getAllValues().forEach(blobName -> {
        Assert.assertTrue(blobName.contains(blobUrlPrefix));
    });
    List<String> allBlobNames = argument.getAllValues();
    Assert.assertNotEquals(allBlobNames.get(0), allBlobNames.get(1));
    verify(mockDataFileWriter1, times((int) maxRecordsPerBlob)).appendEncoded(ByteBuffer.wrap(encodeRecord((IndexedRecord) ome.getMessage())));
    verify(mockDataFileWriter2).appendEncoded(ByteBuffer.wrap(encodeRecord((IndexedRecord) ome2.getMessage())));
    verify(mockDataFileWriter1).create(((IndexedRecord) ome.getMessage()).getSchema(), mockAzureBlobOutputStream1);
    verify(mockDataFileWriter2).create(((IndexedRecord) ome2.getMessage()).getSchema(), mockAzureBlobOutputStream2);
}
Also used : BlobContainerAsyncClient(com.azure.storage.blob.BlobContainerAsyncClient) DataFileWriter(org.apache.avro.file.DataFileWriter) Mockito.anyString(org.mockito.Mockito.anyString) AzureBlobWriterMetrics(org.apache.samza.system.azureblob.producer.AzureBlobWriterMetrics) BlobAsyncClient(com.azure.storage.blob.BlobAsyncClient) BlockBlobAsyncClient(com.azure.storage.blob.specialized.BlockBlobAsyncClient) BlockBlobAsyncClient(com.azure.storage.blob.specialized.BlockBlobAsyncClient) OutgoingMessageEnvelope(org.apache.samza.system.OutgoingMessageEnvelope) PrepareForTest(org.powermock.core.classloader.annotations.PrepareForTest) Test(org.junit.Test)

Example 2 with BlobContainerAsyncClient

use of com.azure.storage.blob.BlobContainerAsyncClient in project samza by apache.

the class TestAzureBlobAvroWriter method testMaxBlobSizeExceeded.

@Test
public void testMaxBlobSizeExceeded() throws Exception {
    String blobUrlPrefix = "test";
    String blobNameRegex = "test/[0-9]{4}/[0-9]{2}/[0-9]{2}/[0-9]{2}/[0-9]{2}-[0-9]{2}-.{8}.avro.gz";
    long maxBlobSize = 1000;
    AzureBlobWriterMetrics mockMetrics = mock(AzureBlobWriterMetrics.class);
    BlobContainerAsyncClient mockContainerClient = PowerMockito.mock(BlobContainerAsyncClient.class);
    azureBlobAvroWriter = spy(new AzureBlobAvroWriter(mockContainerClient, mockMetrics, threadPool, THRESHOLD, 60000, blobUrlPrefix, null, null, null, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, maxBlobSize, 10, mockCompression, true));
    DataFileWriter mockDataFileWriter1 = mock(DataFileWriter.class);
    PowerMockito.whenNew(DataFileWriter.class).withAnyArguments().thenReturn(mockDataFileWriter1);
    BlobAsyncClient mockBlobAsyncClient1 = mock(BlobAsyncClient.class);
    doReturn(mockBlobAsyncClient1).when(mockContainerClient).getBlobAsyncClient(Matchers.matches(blobNameRegex));
    BlockBlobAsyncClient mockBlockBlobAsyncClient1 = mock(BlockBlobAsyncClient.class);
    doReturn(mockBlockBlobAsyncClient1).when(mockBlobAsyncClient1).getBlockBlobAsyncClient();
    AzureBlobOutputStream mockAzureBlobOutputStream1 = mock(AzureBlobOutputStream.class);
    PowerMockito.whenNew(AzureBlobOutputStream.class).withArguments(mockBlockBlobAsyncClient1, threadPool, mockMetrics, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, (long) 60000, THRESHOLD, mockCompression).thenReturn(mockAzureBlobOutputStream1);
    when(mockAzureBlobOutputStream1.getSize()).thenReturn((long) maxBlobSize - 1);
    // first OME creates the first blob
    azureBlobAvroWriter.write(ome);
    OutgoingMessageEnvelope ome2 = createOME("Topic2");
    DataFileWriter mockDataFileWriter2 = mock(DataFileWriter.class);
    PowerMockito.whenNew(DataFileWriter.class).withAnyArguments().thenReturn(mockDataFileWriter2);
    BlobAsyncClient mockBlobAsyncClient2 = mock(BlobAsyncClient.class);
    doReturn(mockBlobAsyncClient2).when(mockContainerClient).getBlobAsyncClient(Matchers.matches(blobNameRegex));
    BlockBlobAsyncClient mockBlockBlobAsyncClient2 = mock(BlockBlobAsyncClient.class);
    doReturn(mockBlockBlobAsyncClient2).when(mockBlobAsyncClient2).getBlockBlobAsyncClient();
    AzureBlobOutputStream mockAzureBlobOutputStream2 = mock(AzureBlobOutputStream.class);
    PowerMockito.whenNew(AzureBlobOutputStream.class).withArguments(mockBlockBlobAsyncClient2, threadPool, mockMetrics, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, (long) 60000, THRESHOLD, mockCompression).thenReturn(mockAzureBlobOutputStream2);
    when(mockAzureBlobOutputStream2.getSize()).thenReturn((long) maxBlobSize - 1);
    // Second OME creates the second blob because maxBlobSize is 1000 and mockAzureBlobOutputStream.getSize is 999.
    azureBlobAvroWriter.write(ome2);
    ArgumentCaptor<String> argument = ArgumentCaptor.forClass(String.class);
    verify(mockContainerClient, times(2)).getBlobAsyncClient(argument.capture());
    argument.getAllValues().forEach(blobName -> {
        Assert.assertTrue(blobName.contains(blobUrlPrefix));
    });
    List<String> allBlobNames = argument.getAllValues();
    Assert.assertNotEquals(allBlobNames.get(0), allBlobNames.get(1));
    verify(mockDataFileWriter1).appendEncoded(ByteBuffer.wrap(encodeRecord((IndexedRecord) ome.getMessage())));
    verify(mockDataFileWriter2).appendEncoded(ByteBuffer.wrap(encodeRecord((IndexedRecord) ome2.getMessage())));
    verify(mockDataFileWriter1).create(((IndexedRecord) ome.getMessage()).getSchema(), mockAzureBlobOutputStream1);
    verify(mockDataFileWriter2).create(((IndexedRecord) ome2.getMessage()).getSchema(), mockAzureBlobOutputStream2);
}
Also used : BlobContainerAsyncClient(com.azure.storage.blob.BlobContainerAsyncClient) DataFileWriter(org.apache.avro.file.DataFileWriter) BlobAsyncClient(com.azure.storage.blob.BlobAsyncClient) BlockBlobAsyncClient(com.azure.storage.blob.specialized.BlockBlobAsyncClient) Mockito.anyString(org.mockito.Mockito.anyString) BlockBlobAsyncClient(com.azure.storage.blob.specialized.BlockBlobAsyncClient) OutgoingMessageEnvelope(org.apache.samza.system.OutgoingMessageEnvelope) AzureBlobWriterMetrics(org.apache.samza.system.azureblob.producer.AzureBlobWriterMetrics) PrepareForTest(org.powermock.core.classloader.annotations.PrepareForTest) Test(org.junit.Test)

Example 3 with BlobContainerAsyncClient

use of com.azure.storage.blob.BlobContainerAsyncClient in project ambry by linkedin.

the class StorageClient method getAsyncContainerClient.

/**
 * Get a reference to an Azure container async client, creating it if necessary.
 * @param containerName the container name.
 * @param autoCreate flag indicating whether to create the container if it does not exist.
 * @return the created {@link BlobContainerClient}.
 */
private BlobContainerAsyncClient getAsyncContainerClient(String containerName, boolean autoCreate) throws ExecutionException, InterruptedException {
    BlobContainerAsyncClient containerAsyncClient = storageAsyncClientRef.get().getBlobContainerAsyncClient(containerName);
    if (autoCreate) {
        if (!knownContainers.contains(containerName)) {
            try {
                if (!containerAsyncClient.exists().toFuture().get()) {
                    containerAsyncClient.create().block();
                    logger.info("Created container {}", containerName);
                }
            } catch (BlobStorageException ex) {
                if (ex.getErrorCode() != BlobErrorCode.CONTAINER_ALREADY_EXISTS) {
                    logger.error("Failed to create container {}", containerName);
                    throw ex;
                }
            }
            knownContainers.add(containerName);
        }
    }
    return containerAsyncClient;
}
Also used : BlobContainerAsyncClient(com.azure.storage.blob.BlobContainerAsyncClient) BlobStorageException(com.azure.storage.blob.models.BlobStorageException)

Example 4 with BlobContainerAsyncClient

use of com.azure.storage.blob.BlobContainerAsyncClient in project samza by apache.

the class TestAzureBlobAvroWriter method setup.

@Before
public void setup() throws Exception {
    threadPool = new ThreadPoolExecutor(1, 1, 60, TimeUnit.SECONDS, new LinkedBlockingDeque<>());
    ome = createOME("Topic1");
    encodedRecord = new byte[100];
    BlobContainerAsyncClient mockContainerAsyncClient = PowerMockito.mock(BlobContainerAsyncClient.class);
    mockDataFileWriter = mock(DataFileWriter.class);
    mockAzureBlobOutputStream = mock(AzureBlobOutputStream.class);
    mockBlockBlobAsyncClient = PowerMockito.mock(BlockBlobAsyncClient.class);
    when(mockBlockBlobAsyncClient.getBlobUrl()).thenReturn("https://samza.blob.core.windows.net/fake-blob-url");
    mockCompression = CompressionFactory.getInstance().getCompression(CompressionType.GZIP);
    azureBlobAvroWriter = spy(new AzureBlobAvroWriter(mockContainerAsyncClient, mock(AzureBlobWriterMetrics.class), threadPool, THRESHOLD, 60000, "test", mockDataFileWriter, mockAzureBlobOutputStream, mockBlockBlobAsyncClient, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, Long.MAX_VALUE, Long.MAX_VALUE, mockCompression, // keeping blob size and number of records unlimited
    false));
    doReturn(encodedRecord).when(azureBlobAvroWriter).encodeRecord((IndexedRecord) ome.getMessage());
}
Also used : LinkedBlockingDeque(java.util.concurrent.LinkedBlockingDeque) BlobContainerAsyncClient(com.azure.storage.blob.BlobContainerAsyncClient) DataFileWriter(org.apache.avro.file.DataFileWriter) ThreadPoolExecutor(java.util.concurrent.ThreadPoolExecutor) BlockBlobAsyncClient(com.azure.storage.blob.specialized.BlockBlobAsyncClient) AzureBlobWriterMetrics(org.apache.samza.system.azureblob.producer.AzureBlobWriterMetrics) Before(org.junit.Before)

Example 5 with BlobContainerAsyncClient

use of com.azure.storage.blob.BlobContainerAsyncClient in project samza by apache.

the class TestAzureBlobAvroWriter method testMultipleBlobClose.

@Test
public void testMultipleBlobClose() throws Exception {
    String blobUrlPrefix = "test";
    long maxBlobSize = AzureBlobAvroWriter.DATAFILEWRITER_OVERHEAD + 1000;
    long maxRecordsPerBlob = 10;
    BlobContainerAsyncClient mockContainerClient = PowerMockito.mock(BlobContainerAsyncClient.class);
    azureBlobAvroWriter = spy(new AzureBlobAvroWriter(mockContainerClient, mock(AzureBlobWriterMetrics.class), threadPool, THRESHOLD, 60000, blobUrlPrefix, mockDataFileWriter, mockAzureBlobOutputStream, mockBlockBlobAsyncClient, blobMetadataGeneratorFactory, blobMetadataGeneratorConfig, STREAM_NAME, maxBlobSize, maxRecordsPerBlob, mockCompression, false));
    DataFileWriter<IndexedRecord> mockDataFileWriter2 = mock(DataFileWriter.class);
    AzureBlobOutputStream mockAzureBlobOutputStream2 = mock(AzureBlobOutputStream.class);
    when(mockAzureBlobOutputStream.getSize()).thenReturn((long) 1);
    BlobAsyncClient mockBlobAsyncClient = mock(BlobAsyncClient.class);
    doReturn(mockBlobAsyncClient).when(mockContainerClient).getBlobAsyncClient(anyString());
    doReturn(mockBlockBlobAsyncClient).when(mockBlobAsyncClient).getBlockBlobAsyncClient();
    PowerMockito.whenNew(AzureBlobOutputStream.class).withAnyArguments().thenReturn(mockAzureBlobOutputStream2);
    PowerMockito.whenNew(DataFileWriter.class).withAnyArguments().thenReturn(mockDataFileWriter2);
    for (int i = 0; i <= maxRecordsPerBlob; i++) {
        azureBlobAvroWriter.write(ome);
    }
    // first OME creates the first blob and 11th OME creates the second blob.
    azureBlobAvroWriter.close();
    verify(mockDataFileWriter).close();
    verify(mockDataFileWriter2).close();
}
Also used : BlobContainerAsyncClient(com.azure.storage.blob.BlobContainerAsyncClient) IndexedRecord(org.apache.avro.generic.IndexedRecord) BlobAsyncClient(com.azure.storage.blob.BlobAsyncClient) BlockBlobAsyncClient(com.azure.storage.blob.specialized.BlockBlobAsyncClient) Mockito.anyString(org.mockito.Mockito.anyString) AzureBlobWriterMetrics(org.apache.samza.system.azureblob.producer.AzureBlobWriterMetrics) PrepareForTest(org.powermock.core.classloader.annotations.PrepareForTest) Test(org.junit.Test)

Aggregations

BlobContainerAsyncClient (com.azure.storage.blob.BlobContainerAsyncClient)5 BlockBlobAsyncClient (com.azure.storage.blob.specialized.BlockBlobAsyncClient)4 AzureBlobWriterMetrics (org.apache.samza.system.azureblob.producer.AzureBlobWriterMetrics)4 BlobAsyncClient (com.azure.storage.blob.BlobAsyncClient)3 DataFileWriter (org.apache.avro.file.DataFileWriter)3 Test (org.junit.Test)3 Mockito.anyString (org.mockito.Mockito.anyString)3 PrepareForTest (org.powermock.core.classloader.annotations.PrepareForTest)3 OutgoingMessageEnvelope (org.apache.samza.system.OutgoingMessageEnvelope)2 BlobStorageException (com.azure.storage.blob.models.BlobStorageException)1 LinkedBlockingDeque (java.util.concurrent.LinkedBlockingDeque)1 ThreadPoolExecutor (java.util.concurrent.ThreadPoolExecutor)1 IndexedRecord (org.apache.avro.generic.IndexedRecord)1 Before (org.junit.Before)1