use of org.apache.drill.exec.oauth.PersistentTokenTable in project drill by apache.
the class StorageResources method updateAuthToken.
@GET
@Path("/storage/{name}/update_oath2_authtoken")
@Produces(MediaType.TEXT_HTML)
public Response updateAuthToken(@PathParam("name") String name, @QueryParam("code") String code) {
try {
if (storage.getPlugin(name).getConfig() instanceof AbstractSecuredStoragePluginConfig) {
AbstractSecuredStoragePluginConfig securedStoragePluginConfig = (AbstractSecuredStoragePluginConfig) storage.getPlugin(name).getConfig();
CredentialsProvider credentialsProvider = securedStoragePluginConfig.getCredentialsProvider();
String callbackURL = this.request.getRequestURL().toString();
// Now exchange the authorization token for an access token
Builder builder = new OkHttpClient.Builder();
OkHttpClient client = builder.build();
Request accessTokenRequest = OAuthUtils.getAccessTokenRequest(credentialsProvider, code, callbackURL);
Map<String, String> updatedTokens = OAuthUtils.getOAuthTokens(client, accessTokenRequest);
// Add to token registry
TokenRegistry tokenRegistry = ((AbstractStoragePlugin) storage.getPlugin(name)).getContext().getoAuthTokenProvider().getOauthTokenRegistry();
// Add a token registry table if none exists
tokenRegistry.createTokenTable(name);
PersistentTokenTable tokenTable = tokenRegistry.getTokenTable(name);
// Add tokens to persistent storage
tokenTable.setAccessToken(updatedTokens.get(OAuthTokenCredentials.ACCESS_TOKEN));
tokenTable.setRefreshToken(updatedTokens.get(OAuthTokenCredentials.REFRESH_TOKEN));
// Get success page
String successPage = null;
try (InputStream inputStream = Resource.newClassPathResource(OAUTH_SUCCESS_PAGE).getInputStream()) {
InputStreamReader reader = new InputStreamReader(inputStream, StandardCharsets.UTF_8);
BufferedReader bufferedReader = new BufferedReader(reader);
successPage = bufferedReader.lines().collect(Collectors.joining("\n"));
bufferedReader.close();
reader.close();
} catch (IOException e) {
Response.status(Status.OK).entity("You may close this window.").build();
}
return Response.status(Status.OK).entity(successPage).build();
} else {
logger.error("{} is not a HTTP plugin. You can only add auth code to HTTP plugins.", name);
return Response.status(Status.INTERNAL_SERVER_ERROR).entity(message("Unable to add authorization code: %s", name)).build();
}
} catch (PluginException e) {
logger.error("Error when adding auth token to {}", name);
return Response.status(Status.INTERNAL_SERVER_ERROR).entity(message("Unable to add authorization code: %s", e.getMessage())).build();
}
}
use of org.apache.drill.exec.oauth.PersistentTokenTable in project drill by apache.
the class TestOAuthProcess method testGetDataWithAuthentication.
@Test
public void testGetDataWithAuthentication() {
String url = hostname + "/update_oath2_authtoken?code=ABCDEF";
Request request = new Request.Builder().url(url).build();
try (MockWebServer server = startServer()) {
server.enqueue(new MockResponse().setResponseCode(200).setBody(ACCESS_TOKEN_RESPONSE));
Response response = httpClient.newCall(request).execute();
// Verify that the request succeeded w/o error
assertEquals(200, response.code());
// Verify that the access and refresh tokens were saved
PersistentTokenTable tokenTable = ((HttpStoragePlugin) cluster.storageRegistry().getPlugin("localOauth")).getTokenRegistry().getTokenTable("localOauth");
assertEquals("you_have_access", tokenTable.getAccessToken());
assertEquals("refresh_me", tokenTable.getRefreshToken());
// Now execute a query and get query results.
server.enqueue(new MockResponse().setResponseCode(200).setBody(TEST_JSON_RESPONSE_WITH_DATATYPES));
String sql = "SELECT * FROM localOauth.test";
DirectRowSet results = queryBuilder().sql(sql).rowSet();
TupleMetadata expectedSchema = new SchemaBuilder().add("col_1", MinorType.FLOAT8, DataMode.OPTIONAL).add("col_2", MinorType.BIGINT, DataMode.OPTIONAL).add("col_3", MinorType.VARCHAR, DataMode.OPTIONAL).build();
RowSet expected = new RowSetBuilder(client.allocator(), expectedSchema).addRow(1.0, 2, "3.0").addRow(4.0, 5, "6.0").build();
RowSetUtilities.verify(expected, results);
} catch (Exception e) {
logger.debug(e.getMessage());
fail();
}
}
use of org.apache.drill.exec.oauth.PersistentTokenTable in project drill by apache.
the class TestOAuthProcess method testGetDataWithTokenRefresh.
@Test
public void testGetDataWithTokenRefresh() {
String url = hostname + "/update_oath2_authtoken?code=ABCDEF";
Request request = new Request.Builder().url(url).build();
try (MockWebServer server = startServer()) {
server.enqueue(new MockResponse().setResponseCode(200).setBody(ACCESS_TOKEN_RESPONSE));
Response response = httpClient.newCall(request).execute();
// Verify that the request succeeded w/o error
assertEquals(200, response.code());
// Verify that the access and refresh tokens were saved
PersistentTokenTable tokenTable = ((HttpStoragePlugin) cluster.storageRegistry().getPlugin("localOauth")).getTokenRegistry().getTokenTable("localOauth");
assertEquals("you_have_access", tokenTable.getAccessToken());
assertEquals("refresh_me", tokenTable.getRefreshToken());
// Now execute a query and get a refresh token
// The API should return a 401 error. This should trigger Drill to automatically
// fire off a second call with the refresh token and then a third request with the
// new access token to obtain the actual data.
server.enqueue(new MockResponse().setResponseCode(401).setBody("Access Denied"));
server.enqueue(new MockResponse().setResponseCode(200).setBody(REFRESH_TOKEN_RESPONSE));
server.enqueue(new MockResponse().setResponseCode(200).setBody(TEST_JSON_RESPONSE_WITH_DATATYPES));
String sql = "SELECT * FROM localOauth.test";
DirectRowSet results = queryBuilder().sql(sql).rowSet();
// Verify that the access and refresh tokens were saved
assertEquals("token 2.0", tokenTable.getAccessToken());
assertEquals("refresh 2.0", tokenTable.getRefreshToken());
TupleMetadata expectedSchema = new SchemaBuilder().add("col_1", MinorType.FLOAT8, DataMode.OPTIONAL).add("col_2", MinorType.BIGINT, DataMode.OPTIONAL).add("col_3", MinorType.VARCHAR, DataMode.OPTIONAL).build();
RowSet expected = new RowSetBuilder(client.allocator(), expectedSchema).addRow(1.0, 2, "3.0").addRow(4.0, 5, "6.0").build();
RowSetUtilities.verify(expected, results);
} catch (Exception e) {
logger.debug(e.getMessage());
fail();
}
}
use of org.apache.drill.exec.oauth.PersistentTokenTable in project drill by apache.
the class TestOAuthProcess method testAccessToken.
@Test
public void testAccessToken() {
String url = hostname + "/update_oath2_authtoken?code=ABCDEF";
Request request = new Request.Builder().url(url).build();
try (MockWebServer server = startServer()) {
server.enqueue(new MockResponse().setResponseCode(200).setBody(ACCESS_TOKEN_RESPONSE));
Response response = httpClient.newCall(request).execute();
// Verify that the request succeeded w/o error
assertEquals(200, response.code());
// Verify that the access and refresh tokens were saved
PersistentTokenTable tokenTable = ((HttpStoragePlugin) cluster.storageRegistry().getPlugin("localOauth")).getTokenTable();
assertEquals("you_have_access", tokenTable.getAccessToken());
assertEquals("refresh_me", tokenTable.getRefreshToken());
} catch (Exception e) {
logger.debug(e.getMessage());
fail();
}
}
Aggregations