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 used managed identity auth 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

@@ -40,29 +40,30 @@
             <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>
+            <version>1.10.1</version>
+        </dependency>
+        <dependency>
+            <groupId>com.azure</groupId>
+            <artifactId>azure-storage-blob</artifactId>
+            <version>12.24.0</version>
+        </dependency>
+        <dependency>
+            <groupId>com.azure</groupId>
+            <artifactId>azure-storage-blob-batch</artifactId>
+            <version>12.20.0</version>
+        </dependency>
+        <dependency>
+            <groupId>com.azure</groupId>
+            <artifactId>azure-storage-common</artifactId>
+            <version>12.23.0</version>
+        </dependency>
+        <dependency>
+            <groupId>com.azure</groupId>
+            <artifactId>azure-core</artifactId>
+            <version>1.43.0</version>
         </dependency>
         <dependency>
             <groupId>com.fasterxml.jackson.module</groupId>
@@ -129,7 +130,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,6 +100,17 @@ 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)
   {

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,93 @@
+/*
+ * 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.BlobContainerClient;
+import com.azure.storage.blob.BlobContainerClientBuilder;
+import com.azure.storage.blob.BlobServiceClient;
+import com.azure.storage.blob.BlobServiceClientBuilder;
+import com.azure.storage.blob.models.CustomerProvidedKey;
+import com.azure.storage.common.StorageSharedKeyCredential;
+
+import java.time.Duration;
+
+/**
+ * Factory class for generating BlobServiceClient and BlobContainerClient objects. This is necessary instead of using
+ * BlobServiceClient.createBlobContainerIfNotExists because sometimes we need different retryOptions on our container
+ * clients and Azure doesn't let us override this setting on the default BlobServiceClient.
+ */
+public class AzureClientFactory
+{
+
+  private final AzureAccountConfig config;
+
+  public AzureClientFactory(AzureAccountConfig config)
+  {
+    this.config = config;
+  }
+
+  public BlobServiceClient getBlobServiceClient()
+  {
+    BlobServiceClientBuilder clientBuilder = new BlobServiceClientBuilder()
+        .endpoint("https://" + config.getAccount() + ".blob.core.windows.net")
+        .retryOptions(new RetryOptions(
+            new ExponentialBackoffOptions().setMaxRetries(config.getMaxTries()).setBaseDelay(Duration.ofMillis(1000)).setMaxDelay(Duration.ofMillis(60000))
+        ));
+
+    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.buildClient();
+  }
+
+  public BlobContainerClient getBlobContainerClient(String containerName, Integer maxRetries)
+  {
+    BlobContainerClientBuilder clientBuilder = new BlobContainerClientBuilder()
+        .endpoint("https://" + config.getAccount() + ".blob.core.windows.net")
+        .containerName(containerName)
+        .retryOptions(new RetryOptions(
+            new ExponentialBackoffOptions()
+                .setMaxRetries(maxRetries != null ? maxRetries : config.getMaxTries())
+                .setBaseDelay(Duration.ofMillis(1000))
+                .setMaxDelay(Duration.ofMillis(60000))
+        ));
+    if (config.getKey() != null) {
+      clientBuilder.customerProvidedKey(new CustomerProvidedKey(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.buildClient();
+  }
+}

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

@@ -19,16 +19,12 @@
 
 package org.apache.druid.storage.azure;
 
+import com.azure.storage.blob.models.BlobItem;
 import com.google.inject.assistedinject.Assisted;
 import com.google.inject.assistedinject.AssistedInject;
-import com.microsoft.azure.storage.ResultContinuation;
-import com.microsoft.azure.storage.ResultSegment;
-import com.microsoft.azure.storage.blob.ListBlobItem;
 import org.apache.druid.java.util.common.RE;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.storage.azure.blob.CloudBlobHolder;
-import org.apache.druid.storage.azure.blob.ListBlobItemHolder;
-import org.apache.druid.storage.azure.blob.ListBlobItemHolderFactory;
 
 import java.net.URI;
 import java.util.Iterator;
@@ -42,36 +38,28 @@ public class AzureCloudBlobIterator implements Iterator<CloudBlobHolder>
 {
   private static final Logger log = new Logger(AzureCloudBlobIterator.class);
   private final AzureStorage storage;
-  private final ListBlobItemHolderFactory blobItemDruidFactory;
   private final Iterator<URI> prefixesIterator;
   private final int maxListingLength;
-
-  private ResultSegment<ListBlobItem> result;
   private String currentContainer;
   private String currentPrefix;
-  private ResultContinuation continuationToken;
   private CloudBlobHolder currentBlobItem;
-  private Iterator<ListBlobItem> blobItemIterator;
+  private Iterator<BlobItem> blobItemIterator;
   private final AzureAccountConfig config;
 
   @AssistedInject
   AzureCloudBlobIterator(
       AzureStorage storage,
-      ListBlobItemHolderFactory blobItemDruidFactory,
       AzureAccountConfig config,
       @Assisted final Iterable<URI> prefixes,
       @Assisted final int maxListingLength
   )
   {
     this.storage = storage;
-    this.blobItemDruidFactory = blobItemDruidFactory;
     this.config = config;
     this.prefixesIterator = prefixes.iterator();
     this.maxListingLength = maxListingLength;
-    this.result = null;
     this.currentContainer = null;
     this.currentPrefix = null;
-    this.continuationToken = null;
     this.currentBlobItem = null;
     this.blobItemIterator = null;
 
@@ -108,8 +96,6 @@ private void prepareNextRequest()
     log.debug("currentUri: %s\ncurrentContainer: %s\ncurrentPrefix: %s",
               currentUri, currentContainer, currentPrefix
     );
-    result = null;
-    continuationToken = null;
   }
 
   private void fetchNextBatch()
@@ -121,14 +107,12 @@ private void fetchNextBatch()
           currentContainer,
           currentPrefix
       );
-      result = AzureUtils.retryAzureOperation(() -> storage.listBlobsWithPrefixInContainerSegmented(
+      // We don't need to use iterableByPage because the client handles this, it will fetch the next page when necessary.
+      blobItemIterator = storage.listBlobsWithPrefixInContainerSegmented(
           currentContainer,
           currentPrefix,
-          continuationToken,
           maxListingLength
-      ), config.getMaxTries());
-      continuationToken = result.getContinuationToken();
-      blobItemIterator = result.getResults().iterator();
+      ).stream().iterator();
     }
     catch (Exception e) {
       throw new RE(
@@ -146,19 +130,15 @@ private void fetchNextBatch()
    */
   private void advanceBlobItem()
   {
-    while (blobItemIterator.hasNext() || continuationToken != null || prefixesIterator.hasNext()) {
+    while (prefixesIterator.hasNext() || blobItemIterator.hasNext()) {
       while (blobItemIterator.hasNext()) {
-        ListBlobItemHolder blobItem = blobItemDruidFactory.create(blobItemIterator.next());
-        /* skip directory objects */
-        if (blobItem.isCloudBlob() && blobItem.getCloudBlob().getBlobLength() > 0) {
-          currentBlobItem = blobItem.getCloudBlob();
+        BlobItem blobItem = blobItemIterator.next();
+        if (!blobItem.isPrefix() && blobItem.getProperties().getContentLength() > 0) {
+          currentBlobItem = new CloudBlobHolder(blobItem, currentContainer);
           return;
         }
       }
-
-      if (continuationToken != null) {
-        fetchNextBatch();
-      } else if (prefixesIterator.hasNext()) {
+      if (prefixesIterator.hasNext()) {
         prepareNextRequest();
         fetchNextBatch();
       }