Set env var BD_USER

Netflix · Jun 8, 2024 · a2cb84b · a2cb84b
1 parent 6bad017
commit a2cb84b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/...-web/src/main/java/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImpl.java b/...-web/src/main/java/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImpl.java
@@ -81,6 +81,7 @@
 public class TitusAgentLauncherImpl implements AgentLauncher {
 
     static final int MEGABYTE_TO_MEGABIT = 8;
+    private static final String BD_USER_ENV_VAR = "BD_USER";
     private static final String GENIE_USER_ATTR = "genie.user";
     private static final String GENIE_SOURCE_HOST_ATTR = "genie.sourceHost";
     private static final String GENIE_ENDPOINT_ATTR = "genie.endpoint";
@@ -246,6 +247,8 @@ public Health health() {
     private TitusBatchJobRequest createJobRequest(final ResolvedJob resolvedJob) throws AgentLaunchException {
         final String jobId = resolvedJob.getJobSpecification().getJob().getId();
 
+        final String jobUser = resolvedJob.getJobMetadata().getUser();
+
         // Map placeholders in entry point template to their values
         final Map<String, String> placeholdersMap = Map.of(
             TitusAgentLauncherProperties.JOB_ID_PLACEHOLDER,
@@ -267,7 +270,7 @@ private TitusBatchJobRequest createJobRequest(final ResolvedJob resolvedJob) thr
         );
         final Duration runtimeLimit = this.titusAgentLauncherProperties.getRuntimeLimit();
 
-        final Map<String, String> jobAttributes = this.createJobAttributes(jobId, resolvedJob);
+        final Map<String, String> jobAttributes = this.createJobAttributes(jobId, jobUser);
 
         final TitusBatchJobRequest.TitusBatchJobRequestBuilder requestBuilder = TitusBatchJobRequest.builder()
             .owner(
@@ -299,7 +302,7 @@ private TitusBatchJobRequest createJobRequest(final ResolvedJob resolvedJob) thr
                     .image(this.getTitusImage(resolvedJob))
                     .entryPoint(entryPoint)
                     .command(command)
-                    .env(this.createJobEnvironment())
+                    .env(this.createJobEnvironment(jobUser))
                     .attributes(
                         this.binder
                             .bind(
@@ -408,9 +411,9 @@ private DataSize getDataSizeProperty(final String propertyKey, final DataSize de
         }
     }
 
-    private Map<String, String> createJobAttributes(final String jobId, final ResolvedJob resolvedJob) {
+    private Map<String, String> createJobAttributes(final String jobId, final String jobUser) {
         final Map<String, String> jobAttributes = new HashMap<>();
-        jobAttributes.put(GENIE_USER_ATTR, resolvedJob.getJobMetadata().getUser());
+        jobAttributes.put(GENIE_USER_ATTR, jobUser);
         jobAttributes.put(GENIE_SOURCE_HOST_ATTR, this.genieHostInfo.getHostname());
         jobAttributes.put(GENIE_ENDPOINT_ATTR, this.titusAgentLauncherProperties.getGenieServerHost());
         jobAttributes.put(GENIE_JOB_ID_ATTR, jobId);
@@ -425,7 +428,7 @@ private Map<String, String> createJobAttributes(final String jobId, final Resolv
         return jobAttributes;
     }
 
-    private Map<String, String> createJobEnvironment() {
+    private Map<String, String> createJobEnvironment(final String jobUser) {
         final Map<String, String> jobEnvironment = this.binder
             .bind(
                 TitusAgentLauncherProperties.ADDITIONAL_ENVIRONMENT_PROPERTY,
@@ -438,6 +441,8 @@ private Map<String, String> createJobEnvironment() {
             jobEnvironment.putAll(this.tracePropagator.injectForAgent(currentSpan.context()));
         }
 
+        jobEnvironment.putIfAbsent(BD_USER_ENV_VAR, jobUser);
+
         return jobEnvironment;
     }
 

diff --git a/.../test/groovy/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImplSpec.groovy b/.../test/groovy/com/netflix/genie/web/agent/launchers/impl/TitusAgentLauncherImplSpec.groovy
@@ -91,6 +91,7 @@ class TitusAgentLauncherImplSpec extends Specification {
         this.requestedMemory = 1024L
         this.launcherProperties = new TitusAgentLauncherProperties()
         this.launcherProperties.setEndpoint(URI.create(TITUS_ENDPOINT_PREFIX))
+        this.launcherProperties.setAdditionalEnvironment(["BD_USER": USER])
 
         this.job = Mock(JobSpecification.ExecutionResource) {
             getId() >> JOB_ID
@@ -508,7 +509,7 @@ class TitusAgentLauncherImplSpec extends Specification {
         1 * this.cache.put(JOB_ID, TITUS_JOB_ID)
         launcherExt.isPresent()
         requestCapture != null
-        requestCapture.getContainer().getEnv().isEmpty()
+        requestCapture.getContainer().getEnv() == ["BD_USER": USER]
 
         when:
         def prop1Key = "${UUID.randomUUID()}.${UUID.randomUUID()}.${UUID.randomUUID()}".toString()
@@ -529,7 +530,7 @@ class TitusAgentLauncherImplSpec extends Specification {
         1 * this.cache.put(JOB_ID, TITUS_JOB_ID)
         launcherExt.isPresent()
         requestCapture != null
-        requestCapture.getContainer().getEnv().size() == 1
+        requestCapture.getContainer().getEnv().size() == 2
         requestCapture.getContainer().getEnv().get(prop1Key) == prop1Value
 
         when:
@@ -551,7 +552,7 @@ class TitusAgentLauncherImplSpec extends Specification {
         1 * this.cache.put(JOB_ID, TITUS_JOB_ID)
         launcherExt.isPresent()
         requestCapture != null
-        requestCapture.getContainer().getEnv().size() == 2
+        requestCapture.getContainer().getEnv().size() == 3
         requestCapture.getContainer().getEnv().get(prop1Key) == prop1Value
         requestCapture.getContainer().getEnv().get(prop2Key) == prop2Value
     }