ATLAS-1906: Atlas client should support setting hadoop-jwt token in header.

client/src/main/java/org/apache/atlas/AtlasAdminClient.java

@@ -77,7 +77,7 @@ public class AtlasAdminClient {
             String[] basicAuthUsernamePassword = AuthenticationUtil.getBasicAuthenticationInput();
             atlasClient = new AtlasClient(atlasServerUri, basicAuthUsernamePassword);
         } else {
-            atlasClient = new AtlasClient(atlasServerUri, null);
+            atlasClient = new AtlasClient(atlasServerUri);
         }
         return handleCommand(commandLine, atlasServerUri, atlasClient);
     }

client/src/main/java/org/apache/atlas/AtlasBaseClient.java

@@ -39,6 +39,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import javax.ws.rs.HttpMethod;
+import javax.ws.rs.core.Cookie;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.MultivaluedMap;
 import javax.ws.rs.core.Response;
@@ -77,6 +78,7 @@ public abstract class AtlasBaseClient {
     private String basicAuthPassword;
     private AtlasClientContext atlasClientContext;
     private boolean retryEnabled = false;
+    private Cookie cookie = null;
 
     protected AtlasBaseClient() {
     }
@@ -106,6 +108,11 @@ public abstract class AtlasBaseClient {
         initializeState(baseUrls, ugi, doAsUser);
     }
 
+    protected AtlasBaseClient(String[] baseUrls, Cookie cookie) {
+        this.cookie = cookie;
+        initializeState(baseUrls, null, null);
+    }
+
     @VisibleForTesting
     protected AtlasBaseClient(WebResource service, Configuration configuration) {
         this.service = service;
@@ -126,6 +133,10 @@ public abstract class AtlasBaseClient {
         initializeState(configuration, baseUrl, null, null);
     }
 
+    public void setCookie(Cookie cookie) {
+        this.cookie = cookie;
+    }
+
     protected static UserGroupInformation getCurrentUGI() throws AtlasException {
         try {
             return UserGroupInformation.getCurrentUser();
@@ -172,14 +183,14 @@ public abstract class AtlasBaseClient {
 
         final URLConnectionClientHandler handler;
 
-        if ((!AuthenticationUtil.isKerberosAuthenticationEnabled()) && basicAuthUser != null && basicAuthPassword != null) {
+        if ((AuthenticationUtil.isKerberosAuthenticationEnabled())) {
+            handler = SecureClientUtils.getClientConnectionHandler(config, configuration, doAsUser, ugi);
+        } else {
             if (configuration.getBoolean(TLS_ENABLED, false)) {
                 handler = SecureClientUtils.getUrlConnectionClientHandler();
             } else {
                 handler = new URLConnectionClientHandler();
             }
-        } else {
-            handler = SecureClientUtils.getClientConnectionHandler(config, configuration, doAsUser, ugi);
         }
         Client client = new Client(handler, config);
         client.setReadTimeout(readTimeout);
@@ -292,10 +303,20 @@ public abstract class AtlasBaseClient {
             if (LOG.isDebugEnabled()) {
                 LOG.debug("Calling API [ {} : {} ] {}", api.getMethod(), api.getPath(), requestObject != null ? "<== " + requestObject : "");
             }
-            clientResponse = resource
+
+            WebResource.Builder requestBuilder = resource.getRequestBuilder();
+
+            // Set content headers
+            requestBuilder
                     .accept(JSON_MEDIA_TYPE)
-                    .type(JSON_MEDIA_TYPE)
-                    .method(api.getMethod(), ClientResponse.class, requestObject);
+                    .type(JSON_MEDIA_TYPE);
+
+            // Set cookie if present
+            if (cookie != null) {
+                requestBuilder.cookie(cookie);
+            }
+
+            clientResponse = requestBuilder.method(api.getMethod(), ClientResponse.class, requestObject);
 
             if (LOG.isDebugEnabled()) {
                 LOG.debug("API {} returned status {}", resource.getURI(), clientResponse.getStatus());

client/src/main/java/org/apache/atlas/AtlasClient.java

@@ -42,6 +42,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import javax.ws.rs.HttpMethod;
+import javax.ws.rs.core.Cookie;
 import javax.ws.rs.core.MultivaluedMap;
 import javax.ws.rs.core.Response;
 import java.util.ArrayList;
@@ -113,9 +114,33 @@ public class AtlasClient extends AtlasBaseClient {
 
     public static final String UNKNOWN_STATUS = "Unknown status";
 
-    // New constuctor for Basic auth
-    public AtlasClient(String[] baseUrl, String[] basicAuthUserNamepassword) {
-        super(baseUrl, basicAuthUserNamepassword);
+    /**
+     * Constructor for AtlasClient with cookie params as header
+     * @param baseUrl
+     * @param cookieName
+     * @param value
+     * @param path
+     * @param domain
+     */
+
+    public AtlasClient(String[] baseUrl, String cookieName, String value, String path, String domain) {
+        super(baseUrl, new Cookie( cookieName, value, path, domain));
+    }
+
+    /**
+     * Constructor for AtlasClient with cookie as header
+     * @param baseUrl
+     * @param cookie
+     */
+
+    public AtlasClient(String[] baseUrl, Cookie cookie) {
+        super(baseUrl, cookie);
+    }
+
+
+    // New constructor for Basic auth
+    public AtlasClient(String[] baseUrl, String[] basicAuthUserNamePassword) {
+        super(baseUrl, basicAuthUserNamePassword);
     }
 
     /**
@@ -152,8 +177,8 @@ public class AtlasClient extends AtlasBaseClient {
     }
 
     @VisibleForTesting
-    public AtlasClient(Configuration configuration, String[] baseUrl, String[] basicAuthUserNamepassword) {
-        super(configuration, baseUrl, basicAuthUserNamepassword);
+    public AtlasClient(Configuration configuration, String[] baseUrl, String[] basicAuthUserNamePassword) {
+        super(configuration, baseUrl, basicAuthUserNamePassword);
     }
 
     @VisibleForTesting

client/src/main/java/org/apache/atlas/AtlasClientV2.java

@@ -41,6 +41,7 @@ import org.apache.commons.configuration.Configuration;
 import org.apache.hadoop.security.UserGroupInformation;
 
 import javax.ws.rs.HttpMethod;
+import javax.ws.rs.core.Cookie;
 import javax.ws.rs.core.MultivaluedMap;
 import javax.ws.rs.core.Response;
 import java.util.List;
@@ -122,8 +123,25 @@ public class AtlasClientV2 extends AtlasBaseClient {
         super(ugi, doAsUser, baseUrls);
     }
 
-    protected AtlasClientV2() {
-        super();
+    /**
+     * Constructor for AtlasClient with cookie params as header
+     * @param baseUrl
+     * @param cookieName
+     * @param value
+     * @param path
+     * @param domain
+     */
+    public AtlasClientV2(String[] baseUrl, String cookieName, String value, String path, String domain) {
+        super(baseUrl, new Cookie(cookieName, value, path, domain));
+    }
+
+    /**
+     * Constructor for AtlasClient with cookie as header
+     * @param baseUrl
+     * @param cookie
+     */
+    public AtlasClientV2(String[] baseUrl, Cookie cookie) {
+        super(baseUrl, cookie);
     }
 
     @VisibleForTesting

client/src/test/java/org/apache/atlas/AtlasClientTest.java

@@ -56,6 +56,8 @@ public class AtlasClientTest {
 
     @Mock
     private WebResource service;
+    @Mock
+    private WebResource.Builder resourceBuilderMock;
 
     @Mock
     private Configuration configuration;
@@ -430,11 +432,11 @@ public class AtlasClientTest {
     }
 
     private WebResource.Builder getBuilder(WebResource resourceObject) {
-        WebResource.Builder builder = mock(WebResource.Builder.class);
+        when(resourceObject.getRequestBuilder()).thenReturn(resourceBuilderMock);
         when(resourceObject.path(anyString())).thenReturn(resourceObject);
-        when(resourceObject.accept(AtlasBaseClient.JSON_MEDIA_TYPE)).thenReturn(builder);
-        when(builder.type(AtlasBaseClient.JSON_MEDIA_TYPE)).thenReturn(builder);
-        return builder;
+        when(resourceBuilderMock.accept(AtlasBaseClient.JSON_MEDIA_TYPE)).thenReturn(resourceBuilderMock);
+        when(resourceBuilderMock.type(AtlasBaseClient.JSON_MEDIA_TYPE)).thenReturn(resourceBuilderMock);
+        return resourceBuilderMock;
     }
 
     private void setupRetryParams() {