Skip to content

Commit

Permalink
Make Azure filesystem tests extensible for different auth types
Browse files Browse the repository at this point in the history
  • Loading branch information
@charlesjmorgan
charlesjmorgan committed Mar 4, 2024
1 parent 131f61a commit df1e21c
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 37 deletions.
80 changes: 45 additions & 35 deletions ...filesystem-azure/src/test/java/io/trino/filesystem/azure/AbstractTestAzureFileSystem.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,10 @@
package io.trino.filesystem.azure;

import com.azure.storage.blob.BlobContainerClient;
import com.azure.storage.blob.BlobServiceClient;
import com.azure.storage.blob.BlobServiceClientBuilder;
import com.azure.storage.blob.models.StorageAccountInfo;
import com.azure.storage.common.StorageSharedKeyCredential;
import com.azure.storage.blob.BlobContainerClientBuilder;
import com.azure.storage.file.datalake.DataLakeFileSystemClient;
import com.azure.storage.file.datalake.DataLakeFileSystemClientBuilder;
import com.azure.storage.file.datalake.DataLakeServiceClient;
import com.azure.storage.file.datalake.DataLakeServiceClientBuilder;
import com.azure.storage.file.datalake.models.PathItem;
import com.azure.storage.file.datalake.options.DataLakePathDeleteOptions;
import io.opentelemetry.api.OpenTelemetry;
Expand Down Expand Up @@ -51,65 +49,73 @@ protected static String getRequiredEnvironmentVariable(String name)
return requireNonNull(System.getenv(name), "Environment variable not set: " + name);
}

enum AccountKind
protected enum AccountKind
{
HIERARCHICAL, FLAT
}

private String account;
private StorageSharedKeyCredential credential;
private AzureAuth azureAuth;
private AccountKind accountKind;
private String containerName;
private Location rootLocation;
private BlobContainerClient blobContainerClient;
private TrinoFileSystem fileSystem;

protected void initialize(String account, String accountKey, AccountKind expectedAccountKind)
protected void initializeWithAccessKey(String account, String accountKey, AccountKind accountKind)
throws IOException
{
this.account = account;
credential = new StorageSharedKeyCredential(account, accountKey);

String blobEndpoint = "https://%s.blob.core.windows.net".formatted(account);
BlobServiceClient blobServiceClient = new BlobServiceClientBuilder()
.endpoint(blobEndpoint)
.credential(credential)
.buildClient();
accountKind = getAccountKind(blobServiceClient);
checkState(accountKind == expectedAccountKind, "Expected %s account, but found %s".formatted(expectedAccountKind, accountKind));
initialize(account, new AzureAuthAccessKey(accountKey), accountKind);
}

private void initialize(String account, AzureAuth azureAuth, AccountKind accountKind)
throws IOException
{
this.account = requireNonNull(account, "account is null");
this.azureAuth = requireNonNull(azureAuth, "azureAuth is null");
this.accountKind = requireNonNull(accountKind, "accountKind is null");
containerName = "test-%s-%s".formatted(accountKind.name().toLowerCase(ROOT), randomUUID());
rootLocation = Location.of("abfs://%s@%s.dfs.core.windows.net/".formatted(containerName, account));

blobContainerClient = blobServiceClient.getBlobContainerClient(containerName);
BlobContainerClientBuilder builder = new BlobContainerClientBuilder()
.endpoint("https://%s.blob.core.windows.net".formatted(account))
.containerName(containerName);
azureAuth.setAuth(account, builder);
blobContainerClient = builder.buildClient();
// this will fail if the container already exists, which is what we want
blobContainerClient.create();
boolean isHierarchicalNamespaceEnabled = isHierarchicalNamespaceEnabled();
if (accountKind == AccountKind.HIERARCHICAL) {
checkState(isHierarchicalNamespaceEnabled, "Expected hierarchical namespaces to be enabled for storage account %s and container %s with account kind %s".formatted(account, containerName, accountKind));
}
else {
checkState(!isHierarchicalNamespaceEnabled, "Expected hierarchical namespaces to not be enabled for storage account %s and container %s with account kind %s".formatted(account, containerName, accountKind));
}

fileSystem = new AzureFileSystemFactory(
OpenTelemetry.noop(),
new AzureAuthAccessKey(accountKey),
azureAuth,
new AzureFileSystemConfig()).create(ConnectorIdentity.ofUser("test"));

cleanupFiles();
}

private static AccountKind getAccountKind(BlobServiceClient blobServiceClient)
private boolean isHierarchicalNamespaceEnabled()
throws IOException
{
StorageAccountInfo accountInfo = blobServiceClient.getAccountInfo();
if (accountInfo.getAccountKind() == com.azure.storage.blob.models.AccountKind.STORAGE_V2) {
if (accountInfo.isHierarchicalNamespaceEnabled()) {
return AccountKind.HIERARCHICAL;
}
return AccountKind.FLAT;
DataLakeFileSystemClient fileSystemClient = createDataLakeFileSystemClient();
try {
return fileSystemClient.getDirectoryClient("/").exists();
}
catch (RuntimeException e) {
throw new IOException("Failed to check whether hierarchical namespaces is enabled for the storage account %s and container %s".formatted(account, containerName));
}
throw new IOException("Unsupported account kind '%s'".formatted(accountInfo.getAccountKind()));
}

@AfterAll
void tearDown()
{
credential = null;
azureAuth = null;
fileSystem = null;
if (blobContainerClient != null) {
blobContainerClient.deleteIfExists();
Expand All @@ -126,12 +132,7 @@ void afterEach()
private void cleanupFiles()
{
if (accountKind == AccountKind.HIERARCHICAL) {
DataLakeFileSystemClient fileSystemClient = new DataLakeFileSystemClientBuilder()
.endpoint("https://%s.dfs.core.windows.net".formatted(account))
.fileSystemName(containerName)
.credential(credential)
.buildClient();

DataLakeFileSystemClient fileSystemClient = createDataLakeFileSystemClient();
DataLakePathDeleteOptions deleteRecursiveOptions = new DataLakePathDeleteOptions().setIsRecursive(true);
for (PathItem pathItem : fileSystemClient.listPaths()) {
if (pathItem.isDirectory()) {
Expand All @@ -147,6 +148,15 @@ private void cleanupFiles()
}
}

private DataLakeFileSystemClient createDataLakeFileSystemClient()
{
DataLakeServiceClientBuilder serviceClientBuilder = new DataLakeServiceClientBuilder()
.endpoint("https://%s.dfs.core.windows.net".formatted(account));
azureAuth.setAuth(account, serviceClientBuilder);
DataLakeServiceClient serviceClient = serviceClientBuilder.buildClient();
return serviceClient.getFileSystemClient(containerName);
}

@Override
protected final boolean isHierarchical()
{
Expand Down
2 changes: 1 addition & 1 deletion ...filesystem-azure/src/test/java/io/trino/filesystem/azure/TestAzureFileSystemGen2Flat.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,6 @@ class TestAzureFileSystemGen2Flat
void setup()
throws IOException
{
initialize(getRequiredEnvironmentVariable("ABFS_FLAT_ACCOUNT"), getRequiredEnvironmentVariable("ABFS_FLAT_ACCESS_KEY"), FLAT);
initializeWithAccessKey(getRequiredEnvironmentVariable("ABFS_FLAT_ACCOUNT"), getRequiredEnvironmentVariable("ABFS_FLAT_ACCESS_KEY"), FLAT);
}
}
2 changes: 1 addition & 1 deletion ...em-azure/src/test/java/io/trino/filesystem/azure/TestAzureFileSystemGen2Hierarchical.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,6 @@ class TestAzureFileSystemGen2Hierarchical
void setup()
throws IOException
{
initialize(getRequiredEnvironmentVariable("ABFS_HIERARCHICAL_ACCOUNT"), getRequiredEnvironmentVariable("ABFS_HIERARCHICAL_ACCESS_KEY"), HIERARCHICAL);
initializeWithAccessKey(getRequiredEnvironmentVariable("ABFS_HIERARCHICAL_ACCOUNT"), getRequiredEnvironmentVariable("ABFS_HIERARCHICAL_ACCESS_KEY"), HIERARCHICAL);
}
}

0 comments on commit df1e21c

Please sign in to comment.