Add support for AzureDNSZone enabled storage accounts used for deep storage

docs/development/extensions-core/azure.md

@@ -42,6 +42,5 @@ To use this Apache Druid extension, [include](../../configuration/extensions.md#
 |`druid.azure.protocol`|the protocol to use|http or https|https|
 |`druid.azure.maxTries`|Number of tries before canceling an Azure operation.| |3|
 |`druid.azure.maxListingLength`|maximum number of input files matching a given prefix to retrieve at a time| |1024|
-|`druid.azure.endpointSuffix`|The endpoint suffix to use. Override the default value to connect to [Azure Government](https://learn.microsoft.com/en-us/azure/azure-government/documentation-government-get-started-connect-to-storage#getting-started-with-storage-api).|Examples: `core.windows.net`, `core.usgovcloudapi.net`|`core.windows.net`|
-
+|`druid.azure.storageAccountEndpointSuffix`| The endpoint suffix to use. Use this config instead of `druid.azure.endpointSuffix`. Override the default value to connect to [Azure Government](https://learn.microsoft.com/en-us/azure/azure-government/documentation-government-get-started-connect-to-storage#getting-started-with-storage-api). This config supports storage accounts enabled for [AzureDNSZone](https://learn.microsoft.com/en-us/azure/dns/dns-getstarted-portal). Note: do not include the storage account name prefix in this config value. | Examples: `ABCD1234.blob.storage.azure.net`, `blob.core.usgovcloudapi.net`| `blob.core.windows.net`|
 See [Azure Services](http://azure.microsoft.com/en-us/pricing/free-trial/) for more information.

extensions-core/azure-extensions/src/main/java/org/apache/druid/storage/azure/AzureAccountConfig.java

@@ -21,19 +21,24 @@
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 
+import javax.annotation.Nullable;
 import javax.validation.constraints.Min;
+import java.util.Objects;
 
 /**
  * Stores the configuration for an Azure account.
  */
 public class AzureAccountConfig
 {
+  static final String DEFAULT_PROTOCOL = "https";
+  static final int DEFAULT_MAX_TRIES = 3;
+
   @JsonProperty
-  private String protocol = "https";
+  private String protocol = DEFAULT_PROTOCOL;
 
   @JsonProperty
   @Min(1)
-  private int maxTries = 3;
+  private int maxTries = DEFAULT_MAX_TRIES;
 
   @JsonProperty
   private String account;
@@ -50,8 +55,13 @@ public class AzureAccountConfig
   @JsonProperty
   private Boolean useAzureCredentialsChain = Boolean.FALSE;
 
+  @Deprecated
+  @Nullable
+  @JsonProperty
+  private String endpointSuffix = null;
+
   @JsonProperty
-  private String endpointSuffix = AzureUtils.DEFAULT_AZURE_ENDPOINT_SUFFIX;
+  private String storageAccountEndpointSuffix = AzureUtils.AZURE_STORAGE_HOST_ADDRESS;
 
   @SuppressWarnings("unused") // Used by Jackson deserialization?
   public void setProtocol(String protocol)
@@ -82,6 +92,12 @@ public void setEndpointSuffix(String endpointSuffix)
     this.endpointSuffix = endpointSuffix;
   }
 
+  @SuppressWarnings("unused") // Used by Jackson deserialization?
+  public void setStorageAccountEndpointSuffix(String storageAccountEndpointSuffix)
+  {
+    this.storageAccountEndpointSuffix = storageAccountEndpointSuffix;
+  }
+
   public String getProtocol()
   {
     return protocol;
@@ -124,18 +140,77 @@ public void setSharedAccessStorageToken(String sharedAccessStorageToken)
     this.sharedAccessStorageToken = sharedAccessStorageToken;
   }
 
+  @SuppressWarnings("unused") // Used by Jackson deserialization?
+  public void setManagedIdentityClientId(String managedIdentityClientId)
+  {
+    this.managedIdentityClientId = managedIdentityClientId;
+  }
+
   public void setUseAzureCredentialsChain(Boolean useAzureCredentialsChain)
   {
     this.useAzureCredentialsChain = useAzureCredentialsChain;
   }
 
+  @Nullable
+  @Deprecated
   public String getEndpointSuffix()
   {
     return endpointSuffix;
   }
 
+  public String getStorageAccountEndpointSuffix()
+  {
+    return storageAccountEndpointSuffix;
+  }
+
   public String getBlobStorageEndpoint()
   {
-    return "blob." + endpointSuffix;
+    // this is here to support the legacy runtime property.
+    if (endpointSuffix != null) {
+      return AzureUtils.BLOB + "." + endpointSuffix;
+    }
+    return storageAccountEndpointSuffix;
+  }
+  @Override
+  public boolean equals(Object o)
+  {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    AzureAccountConfig that = (AzureAccountConfig) o;
+    return Objects.equals(protocol, that.protocol)
+           && Objects.equals(maxTries, that.maxTries)
+           && Objects.equals(account, that.account)
+           && Objects.equals(key, that.key)
+           && Objects.equals(sharedAccessStorageToken, that.sharedAccessStorageToken)
+           && Objects.equals(managedIdentityClientId, that.managedIdentityClientId)
+           && Objects.equals(useAzureCredentialsChain, that.useAzureCredentialsChain)
+           && Objects.equals(endpointSuffix, that.endpointSuffix)
+           && Objects.equals(storageAccountEndpointSuffix, that.storageAccountEndpointSuffix);
+  }
+
+  @Override
+  public int hashCode()
+  {
+    return Objects.hash(protocol, maxTries, account, key, sharedAccessStorageToken, managedIdentityClientId, useAzureCredentialsChain, endpointSuffix, storageAccountEndpointSuffix);
+  }
+
+  @Override
+  public String toString()
+  {
+    return "AzureAccountConfig{" +
+           "protocol=" + protocol +
+           ", maxTries=" + maxTries +
+           ", account=" + account +
+           ", key=" + key +
+           ", sharedAccessStorageToken=" + sharedAccessStorageToken +
+           ", managedIdentityClientId=" + managedIdentityClientId +
+           ", useAzureCredentialsChain=" + useAzureCredentialsChain +
+           ", endpointSuffix=" + endpointSuffix +
+           ", storageAccountEndpointSuffix=" + storageAccountEndpointSuffix +
+           '}';
   }
 }

extensions-core/azure-extensions/src/main/java/org/apache/druid/storage/azure/AzureClientFactory.java

@@ -36,7 +36,7 @@
 import java.util.Map;
 
 /**
- * Factory class for generating BlobServiceClient objects.
+ * Factory class for generating BlobServiceClient objects used for deep storage.
  */
 public class AzureClientFactory
 {

extensions-core/azure-extensions/src/main/java/org/apache/druid/storage/azure/AzureUtils.java

@@ -42,6 +42,8 @@ public class AzureUtils
   @VisibleForTesting
   static final String AZURE_STORAGE_HOST_ADDRESS = "blob.core.windows.net";
 
+  static final String BLOB = "blob";
+
   // The azure storage hadoop access pattern is:
   // wasb[s]://<containername>@<accountname>.blob.<endpointSuffix>/<path>
   // (from https://docs.microsoft.com/en-us/azure/hdinsight/hdinsight-hadoop-use-blob-storage)

extensions-core/azure-extensions/src/test/java/org/apache/druid/storage/azure/AzureAccountConfigTest.java

@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.storage.azure;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.druid.jackson.DefaultObjectMapper;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class AzureAccountConfigTest
+{
+  private static final ObjectMapper MAPPER = new DefaultObjectMapper();
+
+  @Test
+  public void test_getBlobStorageEndpoint_endpointSuffixNullAndStorageAccountEndpointSuffixNull_expectedDefault()
+      throws JsonProcessingException
+  {
+    AzureAccountConfig config = new AzureAccountConfig();
+    AzureAccountConfig configSerde = MAPPER.readValue("{}", AzureAccountConfig.class);
+    Assert.assertEquals(configSerde, config);
+    Assert.assertEquals(AzureUtils.AZURE_STORAGE_HOST_ADDRESS, config.getBlobStorageEndpoint());
+  }
+
+  @Test
+  public void test_getBlobStorageEndpoint_endpointSuffixNotNullAndStorageAccountEndpointSuffixNull_expectEndpoint()
+      throws JsonProcessingException
+  {
+    String endpointSuffix = "core.usgovcloudapi.net";
+    AzureAccountConfig config = new AzureAccountConfig();
+    config.setEndpointSuffix(endpointSuffix);
+    AzureAccountConfig configSerde = MAPPER.readValue(
+        "{"
+        + "\"endpointSuffix\": \"" + endpointSuffix + "\""
+        + "}",
+        AzureAccountConfig.class);
+    Assert.assertEquals(configSerde, config);
+    Assert.assertEquals(AzureUtils.BLOB + "." + endpointSuffix, config.getBlobStorageEndpoint());
+  }
+
+  @Test
+  public void test_getBlobStorageEndpoint_endpointSuffixNotNullAndStorageAccountEndpointSuffixNotNull_expectEndpoint()
+      throws JsonProcessingException
+  {
+    String endpointSuffix = "core.usgovcloudapi.net";
+    String storageAccountEndpointSuffix = "ABCD1234.blob.storage.azure.net";
+    AzureAccountConfig config = new AzureAccountConfig();
+    config.setEndpointSuffix(endpointSuffix);
+    config.setStorageAccountEndpointSuffix(storageAccountEndpointSuffix);
+    AzureAccountConfig configSerde = MAPPER.readValue(
+        "{"
+        + "\"endpointSuffix\": \"" + endpointSuffix + "\","
+        + " \"storageAccountEndpointSuffix\": \"" + storageAccountEndpointSuffix + "\""
+        + "}",
+        AzureAccountConfig.class);
+    Assert.assertEquals(configSerde, config);
+    Assert.assertEquals(AzureUtils.BLOB + "." + endpointSuffix, config.getBlobStorageEndpoint());
+  }
+
+  @Test
+  public void test_getBlobStorageEndpoint_endpointSuffixNullAndStorageAccountEndpointSuffixNotNull_expectStorageAccountEndpoint()
+      throws JsonProcessingException
+  {
+    String storageAccountEndpointSuffix = "ABCD1234.blob.storage.azure.net";
+    AzureAccountConfig config = new AzureAccountConfig();
+    config.setStorageAccountEndpointSuffix(storageAccountEndpointSuffix);
+    AzureAccountConfig configSerde = MAPPER.readValue(
+        "{"
+        + "\"storageAccountEndpointSuffix\": \"" + storageAccountEndpointSuffix + "\""
+        + "}",
+        AzureAccountConfig.class);
+    Assert.assertEquals(configSerde, config);
+    Assert.assertEquals(storageAccountEndpointSuffix, config.getBlobStorageEndpoint());
+  }
+
+  @Test
+  public void test_getManagedIdentityClientId_withValueForManagedIdentityClientId_expectManagedIdentityClientId()
+      throws JsonProcessingException
+  {
+    String managedIdentityClientId = "blah";
+    AzureAccountConfig config = new AzureAccountConfig();
+    config.setManagedIdentityClientId("blah");
+    AzureAccountConfig configSerde = MAPPER.readValue(
+        "{"
+        + "\"managedIdentityClientId\": \"" + managedIdentityClientId + "\""
+        + "}",
+        AzureAccountConfig.class);
+    Assert.assertEquals(configSerde, config);
+    Assert.assertEquals("blah", config.getManagedIdentityClientId());
+  }
+}

extensions-core/azure-extensions/src/test/java/org/apache/druid/storage/azure/AzureClientFactoryTest.java

@@ -24,7 +24,6 @@
 import com.azure.storage.blob.BlobServiceClient;
 import com.azure.storage.common.StorageSharedKeyCredential;
 import com.google.common.collect.ImmutableMap;
-import org.easymock.EasyMock;
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -123,13 +122,55 @@ public void test_blobServiceClientBuilder_useNewClientForDifferentRetryCount()
   @Test
   public void test_blobServiceClientBuilder_useAzureAccountConfig_asDefaultMaxTries()
   {
-    AzureAccountConfig config = EasyMock.createMock(AzureAccountConfig.class);
-    EasyMock.expect(config.getKey()).andReturn("key").times(2);
-    EasyMock.expect(config.getMaxTries()).andReturn(3);
-    EasyMock.expect(config.getBlobStorageEndpoint()).andReturn(AzureUtils.AZURE_STORAGE_HOST_ADDRESS);
+    AzureAccountConfig config = new AzureAccountConfig();
+    config.setKey("key");
+    azureClientFactory = new AzureClientFactory(config);
+    BlobServiceClient expectedBlobServiceClient = azureClientFactory.getBlobServiceClient(AzureAccountConfig.DEFAULT_MAX_TRIES, ACCOUNT);
+    BlobServiceClient blobServiceClient = azureClientFactory.getBlobServiceClient(null, ACCOUNT);
+    Assert.assertEquals(expectedBlobServiceClient, blobServiceClient);
+  }
+
+  @Test
+  public void test_blobServiceClientBuilder_useAzureAccountConfigWithNonDefaultEndpoint_clientUsesEndpointSpecified()
+      throws MalformedURLException
+  {
+    String endpointSuffix = "core.nonDefault.windows.net";
+    AzureAccountConfig config = new AzureAccountConfig();
+    config.setKey("key");
+    config.setEndpointSuffix(endpointSuffix);
+    URL expectedAccountUrl = new URL(AzureAccountConfig.DEFAULT_PROTOCOL, ACCOUNT + "." + AzureUtils.BLOB + "." + endpointSuffix, "");
     azureClientFactory = new AzureClientFactory(config);
-    EasyMock.replay(config);
-    azureClientFactory.getBlobServiceClient(null, ACCOUNT);
-    EasyMock.verify(config);
+    BlobServiceClient blobServiceClient = azureClientFactory.getBlobServiceClient(null, ACCOUNT);
+    Assert.assertEquals(expectedAccountUrl.toString(), blobServiceClient.getAccountUrl());
+  }
+
+  @Test
+  public void test_blobServiceClientBuilder_useAzureAccountConfigWithStorageAccountEndpointAndNonDefaultEndpoint_clientUsesEndpointSpecified()
+      throws MalformedURLException
+  {
+    String endpointSuffix = "core.nonDefault.windows.net";
+    String storageAccountEndpointSuffix = "ABC123.blob.storage.azure.net";
+    AzureAccountConfig config = new AzureAccountConfig();
+    config.setKey("key");
+    config.setEndpointSuffix(endpointSuffix);
+    config.setStorageAccountEndpointSuffix(storageAccountEndpointSuffix);
+    URL expectedAccountUrl = new URL(AzureAccountConfig.DEFAULT_PROTOCOL, ACCOUNT + "." + AzureUtils.BLOB + "." + endpointSuffix, "");
+    azureClientFactory = new AzureClientFactory(config);
+    BlobServiceClient blobServiceClient = azureClientFactory.getBlobServiceClient(null, ACCOUNT);
+    Assert.assertEquals(expectedAccountUrl.toString(), blobServiceClient.getAccountUrl());
+  }
+
+  @Test
+  public void test_blobServiceClientBuilder_useAzureAccountConfigWithStorageAccountEndpointAndNoEndpoint_clientUsesStorageAccountEndpointSpecified()
+      throws MalformedURLException
+  {
+    String storageAccountEndpointSuffix = "ABC123.blob.storage.azure.net";
+    AzureAccountConfig config = new AzureAccountConfig();
+    config.setKey("key");
+    config.setStorageAccountEndpointSuffix(storageAccountEndpointSuffix);
+    URL expectedAccountUrl = new URL(AzureAccountConfig.DEFAULT_PROTOCOL, ACCOUNT + "." + storageAccountEndpointSuffix, "");
+    azureClientFactory = new AzureClientFactory(config);
+    BlobServiceClient blobServiceClient = azureClientFactory.getBlobServiceClient(null, ACCOUNT);
+    Assert.assertEquals(expectedAccountUrl.toString(), blobServiceClient.getAccountUrl());
   }
 }

extensions-core/azure-extensions/src/test/java/org/apache/druid/storage/azure/AzureStorageDruidModuleTest.java

@@ -285,7 +285,8 @@ public void testGetBlobStorageEndpointWithDefaultProperties()
   {
     Properties properties = initializePropertes();
     AzureAccountConfig config = makeInjectorWithProperties(properties).getInstance(AzureAccountConfig.class);
-    Assert.assertEquals(config.getEndpointSuffix(), AzureUtils.DEFAULT_AZURE_ENDPOINT_SUFFIX);
+    Assert.assertNull(config.getEndpointSuffix());
+    Assert.assertEquals(config.getStorageAccountEndpointSuffix(), AzureUtils.AZURE_STORAGE_HOST_ADDRESS);
     Assert.assertEquals(config.getBlobStorageEndpoint(), AzureUtils.AZURE_STORAGE_HOST_ADDRESS);
   }
 

website/.spelling

@@ -39,6 +39,7 @@ Authorizer
 Avatica
 Avro
 Azul
+AzureDNSZone
 BCP
 Base64
 Base64-encoded