ATLAS-1906: Atlas client should support setting hadoop-jwt token in header.
Signed-off-by: apoorvnaik <apoorvnaik@apache.org>
diff --git a/client/src/main/java/org/apache/atlas/AtlasAdminClient.java b/client/src/main/java/org/apache/atlas/AtlasAdminClient.java
index b61b2bf..f334f6c 100644
--- a/client/src/main/java/org/apache/atlas/AtlasAdminClient.java
+++ b/client/src/main/java/org/apache/atlas/AtlasAdminClient.java
@@ -77,7 +77,7 @@
String[] basicAuthUsernamePassword = AuthenticationUtil.getBasicAuthenticationInput();
atlasClient = new AtlasClient(atlasServerUri, basicAuthUsernamePassword);
} else {
- atlasClient = new AtlasClient(atlasServerUri, null);
+ atlasClient = new AtlasClient(atlasServerUri);
}
return handleCommand(commandLine, atlasServerUri, atlasClient);
}
diff --git a/client/src/main/java/org/apache/atlas/AtlasBaseClient.java b/client/src/main/java/org/apache/atlas/AtlasBaseClient.java
index 98da51e..602831a 100644
--- a/client/src/main/java/org/apache/atlas/AtlasBaseClient.java
+++ b/client/src/main/java/org/apache/atlas/AtlasBaseClient.java
@@ -39,6 +39,7 @@
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 @@
private String basicAuthPassword;
private AtlasClientContext atlasClientContext;
private boolean retryEnabled = false;
+ private Cookie cookie = null;
protected AtlasBaseClient() {
}
@@ -106,6 +108,11 @@
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 @@
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 @@
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 @@
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());
diff --git a/client/src/main/java/org/apache/atlas/AtlasClient.java b/client/src/main/java/org/apache/atlas/AtlasClient.java
index ec482da..c50f3f5 100755
--- a/client/src/main/java/org/apache/atlas/AtlasClient.java
+++ b/client/src/main/java/org/apache/atlas/AtlasClient.java
@@ -42,6 +42,7 @@
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 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 @@
}
@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
diff --git a/client/src/main/java/org/apache/atlas/AtlasClientV2.java b/client/src/main/java/org/apache/atlas/AtlasClientV2.java
index 7e287e7..1353357 100644
--- a/client/src/main/java/org/apache/atlas/AtlasClientV2.java
+++ b/client/src/main/java/org/apache/atlas/AtlasClientV2.java
@@ -41,6 +41,7 @@
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 @@
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
diff --git a/client/src/test/java/org/apache/atlas/AtlasClientTest.java b/client/src/test/java/org/apache/atlas/AtlasClientTest.java
index 4dfd43a..c33906a 100644
--- a/client/src/test/java/org/apache/atlas/AtlasClientTest.java
+++ b/client/src/test/java/org/apache/atlas/AtlasClientTest.java
@@ -56,6 +56,8 @@
@Mock
private WebResource service;
+ @Mock
+ private WebResource.Builder resourceBuilderMock;
@Mock
private Configuration configuration;
@@ -430,11 +432,11 @@
}
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() {
