add upstream security information pass thru mechanism

genie-client/src/main/java/com/netflix/genie/client/JobClient.java

@@ -43,6 +43,7 @@
 import okhttp3.ResponseBody;
 import okio.BufferedSink;
 import org.apache.commons.lang3.StringUtils;
+import retrofit2.Call;
 import retrofit2.Retrofit;
 
 import javax.annotation.Nullable;
@@ -52,9 +53,12 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.BiFunction;
+import java.util.function.Function;
 
 /**
  * Client library for the Job Service.
@@ -119,16 +123,38 @@ public JobClient(
      */
     public String submitJob(
         final JobRequest jobRequest
-    ) throws IOException, GenieClientException {
+    ) throws IOException {
+        return submitJob(jobRequest, jobService::submitJob);
+    }
+
+    /**
+     * Submit a job to genie using the jobRequest and upstream security token information.
+     *
+     * @param jobRequest A job request containing all the details for running a job.
+     * @param upstreamSecurityTokenName the security token name provided by upstream.
+     * @param upstreamSecurityTokenValue the security token value provided by upstream.
+     * @return jobId The id of the job submitted.
+     * @throws GenieClientException If the response recieved is not 2xx.
+     * @throws IOException          For Network and other IO issues.
+     */
+    public String submitJob(
+        final JobRequest jobRequest,
+        final String upstreamSecurityTokenName,
+        final String upstreamSecurityTokenValue
+    ) throws IOException {
+        final Map<String, String> headers =
+            Collections.singletonMap(upstreamSecurityTokenName, upstreamSecurityTokenValue);
+        return submitJob(jobRequest, jr -> jobService.submitJob(headers, jr));
+    }
+
+    private String submitJob(
+        final JobRequest jobRequest,
+        final Function<JobRequest, Call<Void>> submitFn) throws IOException {
         if (jobRequest == null) {
             throw new IllegalArgumentException("Job Request cannot be null.");
         }
-        final String locationHeader = this.jobService
-            .submitJob(jobRequest)
-            .execute()
-            .headers()
-            .get(GenieClientUtils.LOCATION_HEADER);
-
+        final String locationHeader =
+            submitFn.apply(jobRequest).execute().headers().get(GenieClientUtils.LOCATION_HEADER);
         if (StringUtils.isBlank(locationHeader)) {
             throw new GenieClientException("No location header. Unable to get ID");
         }
@@ -147,7 +173,37 @@ public String submitJob(
     public String submitJobWithAttachments(
         final JobRequest jobRequest,
         final Map<String, InputStream> attachments
-    ) throws IOException, GenieClientException {
+    ) throws IOException {
+        return submitJobWithAttachments(jobRequest, attachments, jobService::submitJobWithAttachments);
+    }
+
+    /**
+     * Submit a job to genie using the jobRequest and attachments provided with the upstream security information.
+     *
+     * @param jobRequest  A job request containing all the details for running a job.
+     * @param attachments A map of filenames/input-streams needed to be sent to the server as attachments.
+     * @param upstreamSecurityTokenName the security token name provided by upstream.
+     * @param upstreamSecurityTokenValue the security token value provided by upstream.
+     * @return jobId The id of the job submitted.
+     * @throws GenieClientException If the response recieved is not 2xx.
+     * @throws IOException          For Network and other IO issues.
+     */
+    public String submitJobWithAttachments(
+        final JobRequest jobRequest,
+        final Map<String, InputStream> attachments,
+        final String upstreamSecurityTokenName,
+        final String upstreamSecurityTokenValue
+    ) throws IOException {
+        final Map<String, String> headers =
+            Collections.singletonMap(upstreamSecurityTokenName, upstreamSecurityTokenValue);
+        return submitJobWithAttachments(
+            jobRequest, attachments, (jr, at) -> jobService.submitJobWithAttachments(headers, jr, at));
+    }
+
+    private String submitJobWithAttachments(
+        final JobRequest jobRequest,
+        final Map<String, InputStream> attachments,
+        final BiFunction<JobRequest, List<MultipartBody.Part>, Call<Void>> submitFn) throws IOException {
         if (jobRequest == null) {
             throw new IllegalArgumentException("Job Request cannot be null.");
         }
@@ -177,8 +233,7 @@ public void writeTo(final BufferedSink sink) throws IOException {
 
             attachmentFiles.add(part);
         }
-        final String locationHeader = this.jobService
-            .submitJobWithAttachments(jobRequest, attachmentFiles)
+        final String locationHeader = submitFn.apply(jobRequest, attachmentFiles)
             .execute()
             .headers()
             .get(GenieClientUtils.LOCATION_HEADER);

genie-client/src/main/java/com/netflix/genie/client/apis/JobService.java

@@ -32,6 +32,7 @@
 import retrofit2.http.DELETE;
 import retrofit2.http.GET;
 import retrofit2.http.Header;
+import retrofit2.http.HeaderMap;
 import retrofit2.http.Multipart;
 import retrofit2.http.POST;
 import retrofit2.http.Part;
@@ -40,6 +41,7 @@
 import retrofit2.http.Streaming;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -64,6 +66,17 @@ public interface JobService {
     @POST(JOBS_URL_SUFFIX)
     Call<Void> submitJob(@Body JobRequest request);
 
+    /**
+     * Method to submit a job to Genie with user defined custom headers.
+     *
+     * @param customHeaders client defined custom headers.
+     * @param request The request object containing all the
+     * @return A callable object.
+     */
+    @POST(JOBS_URL_SUFFIX)
+    Call<Void> submitJob(@HeaderMap Map<String, String> customHeaders,
+                         @Body JobRequest request);
+
     /**
      * Submit a job with attachments.
      *
@@ -78,6 +91,22 @@ Call<Void> submitJobWithAttachments(
         @Part List<MultipartBody.Part> attachments
     );
 
+    /**
+     * Submit a job with attachments and custom headers.
+     *
+     * @param customHeaders client defined custom headers.
+     * @param request     A JobRequest object containing all the details needed to run the job.
+     * @param attachments A list of all the attachment files to be sent to the server.
+     * @return A callable object.
+     */
+    @Multipart
+    @POST(JOBS_URL_SUFFIX)
+    Call<Void> submitJobWithAttachments(
+        @HeaderMap Map<String, String> customHeaders,
+        @Part("request") JobRequest request,
+        @Part List<MultipartBody.Part> attachments
+    );
+
     /**
      * Method to get all jobs from Genie.
      *

genie-client/src/test/groovy/com/netflix/genie/client/JobClientSpec.groovy

@@ -17,6 +17,7 @@
  */
 package com.netflix.genie.client
 
+import com.netflix.genie.common.dto.JobRequest
 import com.netflix.genie.common.external.util.GenieObjectMapper
 import okhttp3.OkHttpClient
 import okhttp3.mockwebserver.MockResponse
@@ -67,4 +68,75 @@ class JobClientSpec extends Specification {
         "{\"_embedded\": {\"jobSearchResultList\": {}}}" | _
         "{\"_embedded\": {\"jobSearchResultList\": []}}" | _
     }
+
+    @Unroll
+    def "submit job with upstream security information formed correct request"() {
+        setup:
+        def server = new MockWebServer()
+        server.enqueue(new MockResponse().setHeader("location", genieId))
+        server.start()
+        def url = server.url("")
+        def okHttpClient = new OkHttpClient.Builder().build()
+        def retrofit = new Retrofit.Builder()
+            .baseUrl(url)
+            .client(okHttpClient)
+            .addConverterFactory(JacksonConverterFactory.create(GenieObjectMapper.getMapper()))
+            .build()
+        def jobClient = new JobClient(retrofit, 5)
+
+        when:
+        def jobId = jobClient.submitJob(Mock(JobRequest), securityName, securityValue)
+        def request = server.takeRequest()
+
+        then: "No exception is thrown and empty list is returned"
+        noExceptionThrown()
+
+        expect:
+        request.headers.get(securityName) == securityValue
+        jobId == genieId
+
+        cleanup:
+        server.shutdown()
+
+        where:
+        genieId          | securityName                | securityValue
+        "a-genie-id"     | "x-forwarded-authorization" | "some-security-value"
+        "b-genie-id"     | "Authorization"             | "some-auth-value"
+    }
+
+    @Unroll
+    def "submit job with attachments and upstream security information formed correct request"() {
+        setup:
+        def server = new MockWebServer()
+        server.enqueue(new MockResponse().setHeader("location", genieId))
+        server.start()
+        def url = server.url("")
+        def okHttpClient = new OkHttpClient.Builder().build()
+        def retrofit = new Retrofit.Builder()
+            .baseUrl(url)
+            .client(okHttpClient)
+            .addConverterFactory(JacksonConverterFactory.create(GenieObjectMapper.getMapper()))
+            .build()
+        def jobClient = new JobClient(retrofit, 5)
+
+        when:
+        def jobId = jobClient.submitJobWithAttachments(
+            Mock(JobRequest), new HashMap<String, InputStream>(), securityName, securityValue)
+        def request = server.takeRequest()
+
+        then: "No exception is thrown and empty list is returned"
+        noExceptionThrown()
+
+        expect:
+        request.headers.get(securityName) == securityValue
+        jobId == genieId
+
+        cleanup:
+        server.shutdown()
+
+        where:
+        genieId          | securityName                | securityValue
+        "a-genie-id"     | "x-forwarded-authorization" | "some-security-value"
+        "b-genie-id"     | "Authorization"             | "some-auth-value"
+    }
 }