Azure client upgrade to allow identity options

docs/development/extensions-core/azure.md

@@ -33,8 +33,10 @@ To use this Apache Druid extension, [include](../../configuration/extensions.md#
 |--------|---------------|-----------|-------|
 |`druid.storage.type`|azure||Must be set.|
 |`druid.azure.account`||Azure Storage account name.|Must be set.|
-|`druid.azure.key`||Azure Storage account key.|Optional. Either set key or sharedAccessStorageToken but not both.|
-|`druid.azure.sharedAccessStorageToken`||Azure Shared Storage access token|Optional. Either set key or sharedAccessStorageToken but not both.| 
+|`druid.azure.key`||Azure Storage account key.|Optional. Set one of key, sharedAccessStorageToken or useAzureCredentialsChain.|
+|`druid.azure.sharedAccessStorageToken`||Azure Shared Storage access token|Optional. Set one of key, sharedAccessStorageToken or useAzureCredentialsChain..| 
+|`druid.azure.useAzureCredentialsChain`|Use [DefaultAzureCredential](https://learn.microsoft.com/en-us/java/api/overview/azure/identity-readme?view=azure-java-stable) for authentication|Optional. Set one of key, sharedAccessStorageToken or useAzureCredentialsChain.|False|
+|`druid.azure.managedIdentityClientId`|If you want to use managed identity authentication in the `DefaultAzureCredential`, `useAzureCredentialsChain` must be true.||Optional.| 
 |`druid.azure.container`||Azure Storage container name.|Must be set.|
 |`druid.azure.prefix`|A prefix string that will be prepended to the blob names for the segments published to Azure deep storage| |""|
 |`druid.azure.protocol`|the protocol to use|http or https|https|

extensions-core/azure-extensions/pom.xml

@@ -33,36 +33,43 @@
         <relativePath>../../pom.xml</relativePath>
     </parent>
 
+    <dependencyManagement>
+        <dependencies>
+            <dependency>
+                <groupId>com.azure</groupId>
+                <artifactId>azure-sdk-bom</artifactId>
+                <version>1.2.19</version>
+                <type>pom</type>
+                <scope>import</scope>
+            </dependency>
+        </dependencies>
+    </dependencyManagement>
     <dependencies>
         <dependency>
             <groupId>org.apache.druid</groupId>
             <artifactId>druid-processing</artifactId>
             <version>${project.parent.version}</version>
             <scope>provided</scope>
         </dependency>
-
         <dependency>
-            <groupId>com.microsoft.azure</groupId>
-            <artifactId>azure-storage</artifactId>
-            <version>8.6.0</version>
-            <exclusions>
-                <exclusion>
-                    <groupId>org.slf4j</groupId>
-                    <artifactId>slf4j-api</artifactId>
-                </exclusion>
-                <exclusion>
-                    <groupId>com.fasterxml.jackson.core</groupId>
-                    <artifactId>jackson-core</artifactId>
-                </exclusion>
-                <exclusion>
-                    <groupId>org.apache.commons</groupId>
-                    <artifactId>commons-lang3</artifactId>
-                </exclusion>
-                <exclusion>
-                    <groupId>com.google.guava</groupId>
-                    <artifactId>guava</artifactId>
-                </exclusion>
-            </exclusions>
+            <groupId>com.azure</groupId>
+            <artifactId>azure-identity</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.azure</groupId>
+            <artifactId>azure-storage-blob</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.azure</groupId>
+            <artifactId>azure-storage-blob-batch</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.azure</groupId>
+            <artifactId>azure-storage-common</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.azure</groupId>
+            <artifactId>azure-core</artifactId>
         </dependency>
         <dependency>
             <groupId>com.fasterxml.jackson.module</groupId>
@@ -129,7 +136,11 @@
             <artifactId>commons-lang</artifactId>
             <scope>provided</scope>
         </dependency>
-
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-lang3</artifactId>
+            <scope>provided</scope>
+        </dependency>
         <!-- Tests -->
         <dependency>
             <groupId>junit</groupId>

extensions-core/azure-extensions/src/main/java/org/apache/druid/data/input/azure/AzureInputSource.java

@@ -19,14 +19,14 @@
 
 package org.apache.druid.data.input.azure;
 
+import com.azure.storage.blob.models.BlobStorageException;
+import com.azure.storage.blob.specialized.BlockBlobClient;
 import com.fasterxml.jackson.annotation.JacksonInject;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Iterators;
-import com.microsoft.azure.storage.StorageException;
-import com.microsoft.azure.storage.blob.CloudBlob;
 import org.apache.druid.data.input.InputEntity;
 import org.apache.druid.data.input.InputSplit;
 import org.apache.druid.data.input.impl.CloudObjectInputSource;
@@ -42,7 +42,6 @@
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import java.net.URI;
-import java.net.URISyntaxException;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
@@ -150,7 +149,7 @@ public Iterator<LocationWithSize> getDescriptorIteratorForPrefixes(List<URI> pre
                     blob.getBlobLength()
                 );
               }
-              catch (URISyntaxException | StorageException e) {
+              catch (BlobStorageException e) {
                 throw new RuntimeException(e);
               }
             }
@@ -161,14 +160,14 @@ public Iterator<LocationWithSize> getDescriptorIteratorForPrefixes(List<URI> pre
       public long getObjectSize(CloudObjectLocation location)
       {
         try {
-          final CloudBlob blobWithAttributes = storage.getBlockBlobReferenceWithAttributes(
+          final BlockBlobClient blobWithAttributes = storage.getBlockBlobReferenceWithAttributes(
               location.getBucket(),
               location.getPath()
           );
 
-          return blobWithAttributes.getProperties().getLength();
+          return blobWithAttributes.getProperties().getBlobSize();
         }
-        catch (URISyntaxException | StorageException e) {
+        catch (BlobStorageException e) {
           throw new RuntimeException(e);
         }
       }

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

@@ -46,6 +46,12 @@ public class AzureAccountConfig
   @JsonProperty
   private String sharedAccessStorageToken;
 
+  @JsonProperty
+  private String managedIdentityClientId;
+
+  @JsonProperty
+  private Boolean useAzureCredentialsChain = Boolean.FALSE;
+
   @SuppressWarnings("unused") // Used by Jackson deserialization?
   public void setProtocol(String protocol)
   {
@@ -94,9 +100,25 @@ public String getSharedAccessStorageToken()
     return sharedAccessStorageToken;
   }
 
+  public Boolean getUseAzureCredentialsChain()
+  {
+    return useAzureCredentialsChain;
+  }
+
+  public String getManagedIdentityClientId()
+  {
+    return managedIdentityClientId;
+  }
+
+
   @SuppressWarnings("unused") // Used by Jackson deserialization?
   public void setSharedAccessStorageToken(String sharedAccessStorageToken)
   {
     this.sharedAccessStorageToken = sharedAccessStorageToken;
   }
+
+  public void setUseAzureCredentialsChain(Boolean useAzureCredentialsChain)
+  {
+    this.useAzureCredentialsChain = useAzureCredentialsChain;
+  }
 }

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

@@ -19,15 +19,14 @@
 
 package org.apache.druid.storage.azure;
 
+import com.azure.storage.blob.models.BlobStorageException;
 import com.google.common.io.ByteSource;
 import com.google.inject.assistedinject.Assisted;
 import com.google.inject.assistedinject.AssistedInject;
-import com.microsoft.azure.storage.StorageException;
 import org.apache.druid.java.util.common.logger.Logger;
 
 import java.io.IOException;
 import java.io.InputStream;
-import java.net.URISyntaxException;
 
 /**
  * Used for getting an {@link InputStream} to an azure resource.
@@ -62,7 +61,7 @@ public InputStream openStream(long offset) throws IOException
     try {
       return azureStorage.getBlockBlobInputStream(offset, containerName, blobPath);
     }
-    catch (StorageException | URISyntaxException e) {
+    catch (BlobStorageException e) {
       if (AzureUtils.AZURE_RETRY.apply(e)) {
         throw new IOException("Recoverable exception", e);
       }

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

@@ -0,0 +1,84 @@
+/*
+ * 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.azure.core.http.policy.ExponentialBackoffOptions;
+import com.azure.core.http.policy.RetryOptions;
+import com.azure.identity.DefaultAzureCredentialBuilder;
+import com.azure.storage.blob.BlobServiceClient;
+import com.azure.storage.blob.BlobServiceClientBuilder;
+import com.azure.storage.common.StorageSharedKeyCredential;
+
+import java.time.Duration;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Factory class for generating BlobServiceClient objects.
+ */
+public class AzureClientFactory
+{
+
+  private final AzureAccountConfig config;
+  private final Map<Integer, BlobServiceClient> cachedBlobServiceClients;
+
+  public AzureClientFactory(AzureAccountConfig config)
+  {
+    this.config = config;
+    this.cachedBlobServiceClients = new HashMap<>();
+  }
+
+  // It's okay to store clients in a map here because all the configs for specifying azure retries are static, and there are only 2 of them.
+  // The 2 configs are AzureAccountConfig.maxTries and AzureOutputConfig.maxRetrr.
+  // We will only ever have at most 2 clients in cachedBlobServiceClients.
+  public BlobServiceClient getBlobServiceClient(Integer retryCount)
+  {
+    if (!cachedBlobServiceClients.containsKey(retryCount)) {
+      BlobServiceClientBuilder clientBuilder = getAuthenticatedBlobServiceClientBuilder()
+          .retryOptions(new RetryOptions(
+              new ExponentialBackoffOptions()
+                  .setMaxRetries(retryCount != null ? retryCount : config.getMaxTries())
+                  .setBaseDelay(Duration.ofMillis(1000))
+                  .setMaxDelay(Duration.ofMillis(60000))
+          ));
+      cachedBlobServiceClients.put(retryCount, clientBuilder.buildClient());
+    }
+
+    return cachedBlobServiceClients.get(retryCount);
+  }
+
+  private BlobServiceClientBuilder getAuthenticatedBlobServiceClientBuilder()
+  {
+    BlobServiceClientBuilder clientBuilder = new BlobServiceClientBuilder()
+        .endpoint("https://" + config.getAccount() + ".blob.core.windows.net");
+
+    if (config.getKey() != null) {
+      clientBuilder.credential(new StorageSharedKeyCredential(config.getAccount(), config.getKey()));
+    } else if (config.getSharedAccessStorageToken() != null) {
+      clientBuilder.sasToken(config.getSharedAccessStorageToken());
+    } else if (config.getUseAzureCredentialsChain()) {
+      // We might not use the managed identity client id in the credential chain but we can just set it here and it will no-op.
+      DefaultAzureCredentialBuilder defaultAzureCredentialBuilder = new DefaultAzureCredentialBuilder()
+          .managedIdentityClientId(config.getManagedIdentityClientId());
+      clientBuilder.credential(defaultAzureCredentialBuilder.build());
+    }
+    return clientBuilder;
+  }
+}