Make command and cluster criteria tags available to job as environmen…

…t variables

genie-core/src/main/java/com/netflix/genie/core/jobs/JobConstants.java

@@ -245,6 +245,16 @@ public final class JobConstants {
      */
     public static final String GENIE_JOB_MEMORY_ENV_VAR = "GENIE_JOB_MEMORY";
 
+    /**
+     * Environment variable for the Genie command tags in the job request.
+     */
+    public static final String GENIE_REQUESTED_COMMAND_TAGS_ENV_VAR = "GENIE_REQUESTED_COMMAND_TAGS";
+
+    /**
+     * Environment variable for the Genie cluster criteria tags in the job request.
+     */
+    public static final String GENIE_REQUESTED_CLUSTER_TAGS_ENV_VAR = "GENIE_REQUESTED_CLUSTER_TAGS";
+
     /**
      * Process ID.
      **/

genie-core/src/main/java/com/netflix/genie/core/jobs/workflow/impl/InitialSetupTask.java

@@ -19,7 +19,9 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.netflix.genie.common.dto.Cluster;
+import com.netflix.genie.common.dto.ClusterCriteria;
 import com.netflix.genie.common.dto.Command;
+import com.netflix.genie.common.dto.JobRequest;
 import com.netflix.genie.common.exceptions.GenieException;
 import com.netflix.genie.common.exceptions.GeniePreconditionException;
 import com.netflix.genie.common.exceptions.GenieServerException;
@@ -38,6 +40,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
@@ -106,6 +109,9 @@ public void executeTask(@NotNull final Map<String, Object> context) throws Genie
                 jobExecEnv.getMemory()
             );
 
+            // create environment variables for the job request
+            this.createJobRequestEnvironmentVariables(writer, jobExecEnv.getJobRequest());
+
             //Export the Genie Version
             writer.write(GENIE_VERSION_EXPORT);
             writer.write(LINE_SEPARATOR);
@@ -376,6 +382,60 @@ void createJobEnvironmentVariables(
         writer.write(LINE_SEPARATOR);
     }
 
+    @VisibleForTesting
+    void createJobRequestEnvironmentVariables(
+        final Writer writer,
+        final JobRequest jobRequest
+    ) throws IOException {
+
+        // create environment variable for the command tags/criteria in the job request
+        writer.write(JobConstants.EXPORT
+            + JobConstants.GENIE_REQUESTED_COMMAND_TAGS_ENV_VAR
+            + JobConstants.EQUALS_SYMBOL
+            + JobConstants.DOUBLE_QUOTE_SYMBOL
+            + tagsToString(jobRequest.getCommandCriteria())
+            + JobConstants.DOUBLE_QUOTE_SYMBOL
+            + LINE_SEPARATOR);
+
+        final List<ClusterCriteria> clusterCriterias = jobRequest.getClusterCriterias();
+
+        final List<String> clusterCriteriasStrings = new ArrayList<>(clusterCriterias.size());
+
+        for (ClusterCriteria clusterCriteria : clusterCriterias) {
+            clusterCriteriasStrings.add("[" + tagsToString(clusterCriteria.getTags()) + "]");
+        }
+
+        // Append new line
+        writer.write(System.lineSeparator());
+
+        // create environment variable for the list of cluster tags/criteria in the job request as a single
+        // value in the form "[[x,y,z],[a,b,c]]"
+        writer.write(JobConstants.EXPORT
+            + JobConstants.GENIE_REQUESTED_CLUSTER_TAGS_ENV_VAR
+            + JobConstants.EQUALS_SYMBOL
+            + JobConstants.DOUBLE_QUOTE_SYMBOL
+            + "[" + StringUtils.join(clusterCriteriasStrings, ',') + "]"
+            + JobConstants.DOUBLE_QUOTE_SYMBOL
+            + LINE_SEPARATOR);
+
+        // create environment variables for individual tags/criteria in the job request
+        for (int i = 0; i < clusterCriterias.size(); i++) {
+            final ClusterCriteria clusterCriteria = clusterCriterias.get(i);
+
+            // create environment variable for the job name
+            writer.write(JobConstants.EXPORT
+                + JobConstants.GENIE_REQUESTED_CLUSTER_TAGS_ENV_VAR + "_" + i
+                + JobConstants.EQUALS_SYMBOL
+                + JobConstants.DOUBLE_QUOTE_SYMBOL
+                + tagsToString(clusterCriteria.getTags())
+                + JobConstants.DOUBLE_QUOTE_SYMBOL
+                + LINE_SEPARATOR);
+        }
+
+        // Append new line
+        writer.write(LINE_SEPARATOR);
+    }
+
     /**
      * Helper to convert a set of tags into a string that is a suitable value for a shell environment variable.
      * Adds double quotes as necessary (i.e. in case of spaces, newlines), performs escaping of in-tag quotes.

genie-core/src/test/java/com/netflix/genie/core/jobs/workflow/impl/InitialSetupTaskUnitTest.java

@@ -17,8 +17,11 @@
  */
 package com.netflix.genie.core.jobs.workflow.impl;
 
+import com.google.common.collect.Sets;
 import com.netflix.genie.common.dto.Cluster;
+import com.netflix.genie.common.dto.ClusterCriteria;
 import com.netflix.genie.common.dto.Command;
+import com.netflix.genie.common.dto.JobRequest;
 import com.netflix.genie.core.jobs.JobConstants;
 import com.netflix.genie.test.categories.UnitTest;
 import com.netflix.spectator.api.Registry;
@@ -117,6 +120,12 @@ public void testCreateEnvironmentVariables() throws Exception {
         final String commandTag2 = "cmd-foo";
         final String jobName = "The Job";
         final int memory = 1000;
+        final String cmdCritTag1 = "tagX";
+        final String cmdCritTag2 = "tagY";
+        final String cmdCritTag3 = "tagZ";
+        final String cltCritTag1 = "bar";
+        final String cltCritTag2 = "baz";
+        final String cltCritTag3 = "foo";
 
         final Cluster mockCluster = Mockito.mock(Cluster.class);
         Mockito.when(mockCluster.getId()).thenReturn(java.util.Optional.of(clusterId));
@@ -129,12 +138,29 @@ public void testCreateEnvironmentVariables() throws Exception {
         Mockito.when(mockCommand.getId()).thenReturn(java.util.Optional.of(commandName));
         Mockito.when(mockCommand.getTags()).thenReturn(new HashSet<>(Arrays.asList(commandTag2, commandTag1)));
 
+        final JobRequest mockJobRequest = Mockito.mock(JobRequest.class);
+        Mockito
+            .when(mockJobRequest.getCommandCriteria())
+            .thenReturn(
+                Sets.newHashSet(Arrays.asList(cmdCritTag3, cmdCritTag2, cmdCritTag1))
+            );
+        Mockito
+            .when(mockJobRequest.getClusterCriterias())
+            .thenReturn(
+                Arrays.asList(
+                    new ClusterCriteria(Sets.newHashSet(cltCritTag3, cltCritTag1, cltCritTag2)),
+                    new ClusterCriteria(Sets.newHashSet(cltCritTag3, cltCritTag1)),
+                    new ClusterCriteria(Sets.newHashSet(cltCritTag3))
+                )
+            );
+
         final StringWriter mockWriter = new StringWriter();
         this.initialSetupTask.createJobDirEnvironmentVariables(mockWriter, tempDir.getRoot().getAbsolutePath());
         this.initialSetupTask.createApplicationEnvironmentVariables(mockWriter);
         this.initialSetupTask.createCommandEnvironmentVariables(mockWriter, mockCommand);
         this.initialSetupTask.createClusterEnvironmentVariables(mockWriter, mockCluster);
         this.initialSetupTask.createJobEnvironmentVariables(mockWriter, jobId, jobName, memory);
+        this.initialSetupTask.createJobRequestEnvironmentVariables(mockWriter, mockJobRequest);
 
         final String expextedOutput = ""
             + "export GENIE_JOB_DIR=\"" + tempDir.getRoot().getAbsolutePath() + "\"\n"
@@ -162,6 +188,17 @@ public void testCreateEnvironmentVariables() throws Exception {
             + "export GENIE_JOB_NAME=\"" + jobName + "\"\n"
             + "\n"
             + "export GENIE_JOB_MEMORY=" + memory + "\n"
+            + "\n"
+            + "export GENIE_REQUESTED_COMMAND_TAGS=\"" + cmdCritTag1 + "," + cmdCritTag2 + "," + cmdCritTag3 + "\"\n"
+            + "\n"
+            + "export GENIE_REQUESTED_CLUSTER_TAGS=\"["
+                + "[" + cltCritTag1 + "," + cltCritTag2 + "," + cltCritTag3 + "],"
+                + "[" + cltCritTag1 + "," + cltCritTag3 + "],"
+                + "[" + cltCritTag3 + "]"
+            + "]\"" + "\n"
+            + "export GENIE_REQUESTED_CLUSTER_TAGS_0=\"" + cltCritTag1 + "," + cltCritTag2 + "," + cltCritTag3 + "\"\n"
+            + "export GENIE_REQUESTED_CLUSTER_TAGS_1=\"" + cltCritTag1 + "," + cltCritTag3 + "\"\n"
+            + "export GENIE_REQUESTED_CLUSTER_TAGS_2=\"" + cltCritTag3 + "\"\n"
             + "\n";
 
         Assert.assertEquals(expextedOutput, mockWriter.getString());

genie-web/src/test/resources/com/netflix/genie/web/controllers/JobRestControllerIntegrationTests/linux-runsh.txt

@@ -89,6 +89,11 @@ export GENIE_JOB_NAME="List * ... Directories bash job"
 
 export GENIE_JOB_MEMORY=1024
 
+export GENIE_REQUESTED_COMMAND_TAGS="bash"
+
+export GENIE_REQUESTED_CLUSTER_TAGS="[[localhost]]"
+export GENIE_REQUESTED_CLUSTER_TAGS_0="localhost"
+
 export GENIE_VERSION=3
 
 # Sourcing setup file from Cluster: cluster1

genie-web/src/test/resources/com/netflix/genie/web/controllers/JobRestControllerIntegrationTests/non-linux-runsh.txt

@@ -89,6 +89,11 @@ export GENIE_JOB_NAME="List * ... Directories bash job"
 
 export GENIE_JOB_MEMORY=1024
 
+export GENIE_REQUESTED_COMMAND_TAGS="bash"
+
+export GENIE_REQUESTED_CLUSTER_TAGS="[[localhost]]"
+export GENIE_REQUESTED_CLUSTER_TAGS_0="localhost"
+
 export GENIE_VERSION=3
 
 # Sourcing setup file from Cluster: cluster1