com.aventstack.chaintest.storage.AzureBlobClient Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of chaintest-core Show documentation
Core Java client library for ChainTest framework
The newest version!
package com.aventstack.chaintest.storage;

import com.aventstack.chaintest.domain.Embed;
import com.aventstack.chaintest.domain.Test;
import com.azure.identity.DefaultAzureCredential;
import com.azure.identity.DefaultAzureCredentialBuilder;
import com.azure.storage.blob.BlobContainerClient;
import com.azure.storage.blob.BlobServiceClient;
import com.azure.storage.blob.BlobServiceClientBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.util.Base64;
import java.util.Map;

/**
 * Azure Blob Storage client provides the ability to upload files to Azure Blob Storage
 *
 * Further reading:
 * 
 *  
 *      
 *      BlobServiceClientBuilder
 *  
 *  
 *      
 *      DefaultAzureCredential
 *  
 *  
 *      
 *      Setting up your environment for authentication
 *  
 * 
 */
public class AzureBlobClient implements StorageService {

    private static final Logger log = LoggerFactory.getLogger(AzureBlobClient.class);
    private static final String AZURE_STORAGE_CONNECTION_STRING = "AZURE_STORAGE_CONNECTION_STRING";
    private static final String AZURE_CONTAINER_NAME = "AZURE_CONTAINER_NAME";

    private BlobContainerClient _containerClient;
    private String _prefix;

    @Override
    public boolean create(final Map config) {
        _containerClient = getContainerClient(config);
        return null != _containerClient;
    }

    private BlobServiceClient getServiceClient() {
        final BlobServiceClientBuilder builder = new BlobServiceClientBuilder();
        String connectionString = System.getenv(AZURE_STORAGE_CONNECTION_STRING);
        if (connectionString == null) {
            connectionString = System.getProperty(AZURE_STORAGE_CONNECTION_STRING);
        }
        if (connectionString != null) {
            return builder.connectionString(connectionString).buildClient();
        }
        try {
            log.debug("Will use the DefaultAzureCredentialBuilder to authenticate");
            final DefaultAzureCredential defaultAzureCredential = new DefaultAzureCredentialBuilder().build();
            return builder.credential(defaultAzureCredential).buildClient();
        } catch (final Exception e) {
            log.error("Failed to create Azure Blob Storage client", e);
            return null;
        }
    }

    private BlobContainerClient getContainerClient(final Map config) {
        final BlobServiceClient serviceClient = getServiceClient();
        if (serviceClient == null) {
            return null;
        }
        String containerName = System.getenv(AZURE_CONTAINER_NAME);
        if (containerName == null) {
            containerName = System.getProperty(AZURE_CONTAINER_NAME);
        }
        if (containerName == null && config.containsKey(STORAGE_CONTAINER_NAME)) {
            containerName = config.get(STORAGE_CONTAINER_NAME);
        }
        containerName = containerName != null && !containerName.isBlank() ? containerName : DEFAULT_CONTAINER_NAME;
        serviceClient.createBlobContainerIfNotExists(containerName);
        return serviceClient.getBlobContainerClient(containerName);
    }

    public void withPrefix(final String prefix) {
        _prefix = prefix;
    }

    @Override
    public String upload(final Test test, final String key, final byte[] data) {
        final String prefixKey = getPrefixKey(_prefix, key);
        try (final ByteArrayInputStream inputStream = new ByteArrayInputStream(data)) {
            _containerClient.getBlobClient(prefixKey).upload(inputStream, data.length);
            return _containerClient.getBlobClient(key).getBlobUrl();
        } catch (Exception e) {
            log.error("Failed to upload key {} to Azure Blob Storage", key, e);
        }
        return prefixKey;
    }

    @Override
    public String upload(final Test test, final String key, final String base64) {
        final byte[] data = Base64.getDecoder().decode(base64.getBytes());
        return upload(test, key, data);
    }

    @Override
    public String upload(final Test test, final String key, final File file) {
        final String prefixKey = getPrefixKey(_prefix, key);
        try {
            _containerClient.getBlobClient(prefixKey)
                    .uploadFromFile(file.getAbsolutePath());
            return _containerClient.getBlobClient(key + ".png").getBlobUrl();
        } catch (final Exception e) {
            log.error("Failed to upload key {} to Azure Blob Storage", key, e);
        }
        return null;
    }

    @Override
    public void upload(final Test test, final Embed embed) {
        String url = null;
        if (null != embed.getBytes()) {
            url = upload(test, embed.getName(), embed.getBytes());
        } else if (null != embed.getBase64() && !embed.getBase64().isBlank()) {
            url = upload(test, embed.getName(), embed.getBase64());
        } else if (null != embed.getFile()) {
            url = upload(test, embed.getName(), embed.getFile());
        } else {
            log.error("Unable to upload Embed to Azure Blob Storage. Source missing");
        }
        embed.setUrl(url);
    }

    @Override
    public void close() { }

}