Fix support for Azure managed identity clientId

Signed-off-by: Paolo Di Tommaso <paolo.ditommaso@gmail.com>
nextflow-io · Jun 14, 2024 · 306814e · 306814e
1 parent c08dc49
commit 306814e
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 9 deletions.
diff --git a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy
@@ -297,11 +297,16 @@ class AzBatchService implements Closeable {
     }
 
     protected TokenCredential createBatchCredentialsWithManagedIdentity() {
-        log.debug '[AZURE BATCH] Creating Azure Batch client using Managed Identity credentials'
-
-        return new ManagedIdentityCredentialBuilder()
-                .clientId(config.managedIdentity().clientId)
-                .build()
+        final clientId = config.managedIdentity().clientId
+        final credential = new ManagedIdentityCredentialBuilder()
+        if (clientId) {
+            log.debug "[AZURE BATCH] Creating Azure Batch client using Managed Identity credentials - clientId: ${clientId}"
+            credential.clientId(clientId)
+        }
+        else {
+            log.debug '[AZURE BATCH] Creating Azure Batch client using Managed Identity credentials - not clientId provided'
+        }
+        return credential.build()
     }
 
     protected BatchClient createBatchClient() {

diff --git a/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchServiceTest.groovy b/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchServiceTest.groovy
@@ -1,17 +1,17 @@
 package nextflow.cloud.azure.batch
 
-import com.azure.compute.batch.models.BatchPool
-import com.azure.compute.batch.models.ElevationLevel
-
-import java.nio.file.Path
 import java.time.Instant
 import java.time.temporal.ChronoUnit
 import java.util.function.Predicate
 
+import com.azure.compute.batch.models.BatchPool
+import com.azure.compute.batch.models.ElevationLevel
+import com.azure.identity.ManagedIdentityCredential
 import com.google.common.hash.HashCode
 import nextflow.Global
 import nextflow.Session
 import nextflow.cloud.azure.config.AzConfig
+import nextflow.cloud.azure.config.AzManagedIdentityOpts
 import nextflow.cloud.azure.config.AzPoolOpts
 import nextflow.cloud.azure.config.AzStartTaskOpts
 import nextflow.file.FileSystemPathFactory
@@ -708,4 +708,25 @@ class AzBatchServiceTest extends Specification {
         result.containerSettings.imageName == 'ubuntu:latest'
         result.containerSettings.containerRunOptions == '-v /etc/ssl/certs:/etc/ssl/certs:ro -v /etc/pki:/etc/pki:ro --privileged -e FUSION_WORK=/fusion/az/foo/work/dir -e FUSION_TAGS=[.command.*|.exitcode|.fusion.*](nextflow.io/metadata=true),[*](nextflow.io/temporary=true) -e AZURE_STORAGE_ACCOUNT=my-account -e AZURE_STORAGE_SAS_TOKEN=1234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890 '
     }
+
+    @Unroll
+    def 'should create manage credentials token' () {
+        given:
+        def config = Mock(AzConfig)
+        def exec = Mock(AzBatchExecutor) {getConfig() >> new AzConfig(CONFIG) }
+        AzBatchService service = Spy(new AzBatchService(exec))
+
+        when:
+        def token = service.createBatchCredentialsWithManagedIdentity()
+        then:
+        config.managedIdentity() >> { Mock(AzManagedIdentityOpts) }
+        then:
+        token instanceof ManagedIdentityCredential
+        (token as ManagedIdentityCredential).clientId == EXPECTED
+
+        where:
+        CONFIG                                          | EXPECTED
+        [:]                                             | null
+        [managedIdentity: [clientId: 'client-123']]     | 'client-123'
+    }
 }