@Test public void testCustomBucketWriterRoleAssignment() { when(configuration.isDefaultS3WriterRoleEnabled()).thenReturn(true); ContainerInfo.Builder containerInfoBuilder = ContainerInfo.newBuilder(); Job<BatchJobExt> job = JobGenerator.oneBatchJob(); job = job.toBuilder() .withJobDescriptor(job.getJobDescriptor().toBuilder() .withContainer(job.getJobDescriptor().getContainer().toBuilder() .withAttributes(Collections.singletonMap( JobAttributes.JOB_CONTAINER_ATTRIBUTE_S3_BUCKET_NAME, "myOwnBucket" )).build() ) .build() ) .build(); converter.appendS3WriterRole(containerInfoBuilder, job, JobGenerator.oneBatchTask()); assertThat(containerInfoBuilder.getPassthroughAttributesMap()).containsEntry( S3_WRITER_ROLE, job.getJobDescriptor().getContainer().getSecurityProfile().getIamRole() ); }
@VisibleForTesting void appendS3WriterRole(TitanProtos.ContainerInfo.Builder containerInfoBuilder, Job<?> job, Task task) { if (!configuration.isDefaultS3WriterRoleEnabled()) { return; } if (LogStorageInfos.findCustomS3Bucket(job).isPresent()) { containerInfoBuilder.putPassthroughAttributes( S3_WRITER_ROLE, job.getJobDescriptor().getContainer().getSecurityProfile().getIamRole() ); } else { Evaluators.applyNotNull( configuration.getDefaultS3WriterRole(), role -> containerInfoBuilder.putPassthroughAttributes(S3_WRITER_ROLE, role) ); } logStorageInfo.getS3LogLocation(task, false).ifPresent(s3LogLocation -> Evaluators.applyNotNull( s3LogLocation.getBucket(), bucket -> containerInfoBuilder.putPassthroughAttributes(S3_BUCKET_NAME, bucket) ) ); }
DefaultTaskToPodConverter implements TaskToPodConverter { @VisibleForTesting void appendS3WriterRole(TitanProtos.ContainerInfo.Builder containerInfoBuilder, Job<?> job, Task task) { if (!configuration.isDefaultS3WriterRoleEnabled()) { return; } if (LogStorageInfos.findCustomS3Bucket(job).isPresent()) { containerInfoBuilder.putPassthroughAttributes( S3_WRITER_ROLE, job.getJobDescriptor().getContainer().getSecurityProfile().getIamRole() ); } else { Evaluators.applyNotNull( configuration.getDefaultS3WriterRole(), role -> containerInfoBuilder.putPassthroughAttributes(S3_WRITER_ROLE, role) ); } logStorageInfo.getS3LogLocation(task, false).ifPresent(s3LogLocation -> Evaluators.applyNotNull( s3LogLocation.getBucket(), bucket -> containerInfoBuilder.putPassthroughAttributes(S3_BUCKET_NAME, bucket) ) ); } }
DefaultTaskToPodConverter implements TaskToPodConverter { @VisibleForTesting void appendS3WriterRole(TitanProtos.ContainerInfo.Builder containerInfoBuilder, Job<?> job, Task task) { if (!configuration.isDefaultS3WriterRoleEnabled()) { return; } if (LogStorageInfos.findCustomS3Bucket(job).isPresent()) { containerInfoBuilder.putPassthroughAttributes( S3_WRITER_ROLE, job.getJobDescriptor().getContainer().getSecurityProfile().getIamRole() ); } else { Evaluators.applyNotNull( configuration.getDefaultS3WriterRole(), role -> containerInfoBuilder.putPassthroughAttributes(S3_WRITER_ROLE, role) ); } logStorageInfo.getS3LogLocation(task, false).ifPresent(s3LogLocation -> Evaluators.applyNotNull( s3LogLocation.getBucket(), bucket -> containerInfoBuilder.putPassthroughAttributes(S3_BUCKET_NAME, bucket) ) ); } @Inject  DefaultTaskToPodConverter(DirectKubeConfiguration configuration,
                                     PodAffinityFactory podAffinityFactory,
                                     TaintTolerationFactory taintTolerationFactory,
                                     LogStorageInfo<Task> logStorageInfo); }
DefaultTaskToPodConverter implements TaskToPodConverter { @VisibleForTesting void appendS3WriterRole(TitanProtos.ContainerInfo.Builder containerInfoBuilder, Job<?> job, Task task) { if (!configuration.isDefaultS3WriterRoleEnabled()) { return; } if (LogStorageInfos.findCustomS3Bucket(job).isPresent()) { containerInfoBuilder.putPassthroughAttributes( S3_WRITER_ROLE, job.getJobDescriptor().getContainer().getSecurityProfile().getIamRole() ); } else { Evaluators.applyNotNull( configuration.getDefaultS3WriterRole(), role -> containerInfoBuilder.putPassthroughAttributes(S3_WRITER_ROLE, role) ); } logStorageInfo.getS3LogLocation(task, false).ifPresent(s3LogLocation -> Evaluators.applyNotNull( s3LogLocation.getBucket(), bucket -> containerInfoBuilder.putPassthroughAttributes(S3_BUCKET_NAME, bucket) ) ); } @Inject  DefaultTaskToPodConverter(DirectKubeConfiguration configuration,
                                     PodAffinityFactory podAffinityFactory,
                                     TaintTolerationFactory taintTolerationFactory,
                                     LogStorageInfo<Task> logStorageInfo); @Override V1Pod apply(Job<?> job, Task task); }
DefaultTaskToPodConverter implements TaskToPodConverter { @VisibleForTesting void appendS3WriterRole(TitanProtos.ContainerInfo.Builder containerInfoBuilder, Job<?> job, Task task) { if (!configuration.isDefaultS3WriterRoleEnabled()) { return; } if (LogStorageInfos.findCustomS3Bucket(job).isPresent()) { containerInfoBuilder.putPassthroughAttributes( S3_WRITER_ROLE, job.getJobDescriptor().getContainer().getSecurityProfile().getIamRole() ); } else { Evaluators.applyNotNull( configuration.getDefaultS3WriterRole(), role -> containerInfoBuilder.putPassthroughAttributes(S3_WRITER_ROLE, role) ); } logStorageInfo.getS3LogLocation(task, false).ifPresent(s3LogLocation -> Evaluators.applyNotNull( s3LogLocation.getBucket(), bucket -> containerInfoBuilder.putPassthroughAttributes(S3_BUCKET_NAME, bucket) ) ); } @Inject  DefaultTaskToPodConverter(DirectKubeConfiguration configuration,
                                     PodAffinityFactory podAffinityFactory,
                                     TaintTolerationFactory taintTolerationFactory,
                                     LogStorageInfo<Task> logStorageInfo); @Override V1Pod apply(Job<?> job, Task task);  }