Fix support for GCS requester pays bucket option

Signed-off-by: Paolo Di Tommaso <paolo.ditommaso@gmail.com>
nextflow-io · Apr 3, 2024 · d9d61cf · d9d61cf
1 parent a550e52
commit d9d61cf
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 3 deletions.
diff --git a/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchScriptLauncher.groovy b/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchScriptLauncher.groovy
@@ -25,6 +25,7 @@ import com.google.cloud.batch.v1.Volume
 import com.google.cloud.storage.contrib.nio.CloudStoragePath
 import groovy.transform.CompileStatic
 import groovy.util.logging.Slf4j
+import nextflow.cloud.google.batch.client.BatchConfig
 import nextflow.executor.BashWrapperBuilder
 import nextflow.extension.FilesEx
 import nextflow.processor.TaskBean
@@ -43,6 +44,7 @@ class GoogleBatchScriptLauncher extends BashWrapperBuilder implements GoogleBatc
 
     private static final String MOUNT_ROOT = '/mnt/disks'
 
+    private BatchConfig config
     private CloudStoragePath remoteWorkDir
     private Path remoteBinDir
     private Set<String> buckets = new HashSet<>()
@@ -127,14 +129,18 @@ class GoogleBatchScriptLauncher extends BashWrapperBuilder implements GoogleBatc
     List<Volume> getVolumes() {
         final result = new ArrayList(10)
         for( String it : buckets ) {
+            final mountOptions = ['-o rw', '-implicit-dirs']
+            if( config && config.googleOpts.enableRequesterPaysBuckets )
+                mountOptions << "--billing-project ${config.googleOpts.projectId}".toString()
+
             result.add(
                 Volume.newBuilder()
                     .setGcs(
                         GCS.newBuilder()
                             .setRemotePath(it)
                     )
                     .setMountPath( "${MOUNT_ROOT}/${it}".toString() )
-                    .addAllMountOptions( ['-o rw', '-implicit-dirs'] )
+                    .addAllMountOptions( mountOptions )
                     .build()
             )
         }
@@ -160,4 +166,9 @@ class GoogleBatchScriptLauncher extends BashWrapperBuilder implements GoogleBatc
         return remoteWorkDir.resolve(TaskRun.CMD_INFILE)
     }
 
+    GoogleBatchScriptLauncher withConfig(BatchConfig config) {
+        this.config = config
+        return this
+    }
+
 }
diff --git a/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy b/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy
@@ -125,6 +125,7 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
         else {
             final taskBean = task.toTaskBean()
             return new GoogleBatchScriptLauncher(taskBean, executor.remoteBinDir)
+                .withConfig(executor.config)
         }
     }
 

diff --git a/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchScriptLauncherTest.groovy b/plugins/nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchScriptLauncherTest.groovy
@@ -20,6 +20,8 @@ package nextflow.cloud.google.batch
 import java.nio.file.Paths
 
 import com.google.cloud.storage.contrib.nio.CloudStorageFileSystem
+import nextflow.cloud.google.GoogleOpts
+import nextflow.cloud.google.batch.client.BatchConfig
 import spock.lang.Specification
 import spock.lang.Unroll
 
@@ -49,6 +51,13 @@ class GoogleBatchScriptLauncherTest extends Specification{
     def 'should compute volume mounts' () {
         given:
         def launcher = new GoogleBatchScriptLauncher()
+        launcher.config = Mock(BatchConfig) {
+            getGoogleOpts() >> Mock(GoogleOpts) {
+                getProjectId() >> 'my-project'
+                getEnableRequesterPaysBuckets() >> true
+            }
+        }
+        and:
         def PATH1 = CloudStorageFileSystem.forBucket('alpha').getPath('/data/sample1.bam')
         def PATH2 = CloudStorageFileSystem.forBucket('alpha').getPath('/data/sample2.bam')
         def PATH3 = CloudStorageFileSystem.forBucket('omega').getPath('/data/sample3.bam')
@@ -71,10 +80,10 @@ class GoogleBatchScriptLauncherTest extends Specification{
         volumes.size() == 2
         volumes[0].getGcs().getRemotePath() == 'alpha'
         volumes[0].getMountPath() == '/mnt/disks/alpha'
-        volumes[0].getMountOptionsList() == ['-o rw', '-implicit-dirs']
+        volumes[0].getMountOptionsList() == ['-o rw', '-implicit-dirs', '--billing-project my-project']
         volumes[1].getGcs().getRemotePath() == 'omega'
         volumes[1].getMountPath() == '/mnt/disks/omega'
-        volumes[1].getMountOptionsList() == ['-o rw', '-implicit-dirs']
+        volumes[1].getMountOptionsList() == ['-o rw', '-implicit-dirs', '--billing-project my-project']
     }
 
 }