Search in sources :

Example 1 with PersistentTokenTable

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();
    }
}
Also used : OkHttpClient(okhttp3.OkHttpClient) InputStreamReader(java.io.InputStreamReader) TokenRegistry(org.apache.drill.exec.oauth.TokenRegistry) InputStream(java.io.InputStream) Builder(okhttp3.OkHttpClient.Builder) PluginException(org.apache.drill.exec.store.StoragePluginRegistry.PluginException) Request(okhttp3.Request) HttpServletRequest(javax.servlet.http.HttpServletRequest) CredentialsProvider(org.apache.drill.common.logical.security.CredentialsProvider) IOException(java.io.IOException) BufferedReader(java.io.BufferedReader) PersistentTokenTable(org.apache.drill.exec.oauth.PersistentTokenTable) AbstractSecuredStoragePluginConfig(org.apache.drill.common.logical.AbstractSecuredStoragePluginConfig) AbstractStoragePlugin(org.apache.drill.exec.store.AbstractStoragePlugin) Path(javax.ws.rs.Path) Produces(javax.ws.rs.Produces) GET(javax.ws.rs.GET)

Example 2 with PersistentTokenTable

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();
    }
}
Also used : MockResponse(okhttp3.mockwebserver.MockResponse) SchemaBuilder(org.apache.drill.exec.record.metadata.SchemaBuilder) RowSetBuilder(org.apache.drill.exec.physical.rowSet.RowSetBuilder) ClusterFixtureBuilder(org.apache.drill.test.ClusterFixtureBuilder) Request(okhttp3.Request) DirectRowSet(org.apache.drill.exec.physical.rowSet.DirectRowSet) DirectRowSet(org.apache.drill.exec.physical.rowSet.DirectRowSet) RowSet(org.apache.drill.exec.physical.rowSet.RowSet) IOException(java.io.IOException) Response(okhttp3.Response) MockResponse(okhttp3.mockwebserver.MockResponse) RowSetBuilder(org.apache.drill.exec.physical.rowSet.RowSetBuilder) TupleMetadata(org.apache.drill.exec.record.metadata.TupleMetadata) MockWebServer(okhttp3.mockwebserver.MockWebServer) SchemaBuilder(org.apache.drill.exec.record.metadata.SchemaBuilder) PersistentTokenTable(org.apache.drill.exec.oauth.PersistentTokenTable) ClusterTest(org.apache.drill.test.ClusterTest) Test(org.junit.Test)

Example 3 with PersistentTokenTable

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();
    }
}
Also used : MockResponse(okhttp3.mockwebserver.MockResponse) SchemaBuilder(org.apache.drill.exec.record.metadata.SchemaBuilder) RowSetBuilder(org.apache.drill.exec.physical.rowSet.RowSetBuilder) ClusterFixtureBuilder(org.apache.drill.test.ClusterFixtureBuilder) Request(okhttp3.Request) DirectRowSet(org.apache.drill.exec.physical.rowSet.DirectRowSet) DirectRowSet(org.apache.drill.exec.physical.rowSet.DirectRowSet) RowSet(org.apache.drill.exec.physical.rowSet.RowSet) IOException(java.io.IOException) Response(okhttp3.Response) MockResponse(okhttp3.mockwebserver.MockResponse) RowSetBuilder(org.apache.drill.exec.physical.rowSet.RowSetBuilder) TupleMetadata(org.apache.drill.exec.record.metadata.TupleMetadata) MockWebServer(okhttp3.mockwebserver.MockWebServer) SchemaBuilder(org.apache.drill.exec.record.metadata.SchemaBuilder) PersistentTokenTable(org.apache.drill.exec.oauth.PersistentTokenTable) ClusterTest(org.apache.drill.test.ClusterTest) Test(org.junit.Test)

Example 4 with PersistentTokenTable

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();
    }
}
Also used : Response(okhttp3.Response) MockResponse(okhttp3.mockwebserver.MockResponse) MockResponse(okhttp3.mockwebserver.MockResponse) SchemaBuilder(org.apache.drill.exec.record.metadata.SchemaBuilder) RowSetBuilder(org.apache.drill.exec.physical.rowSet.RowSetBuilder) ClusterFixtureBuilder(org.apache.drill.test.ClusterFixtureBuilder) Request(okhttp3.Request) MockWebServer(okhttp3.mockwebserver.MockWebServer) PersistentTokenTable(org.apache.drill.exec.oauth.PersistentTokenTable) IOException(java.io.IOException) ClusterTest(org.apache.drill.test.ClusterTest) Test(org.junit.Test)

Aggregations

IOException (java.io.IOException)4 Request (okhttp3.Request)4 PersistentTokenTable (org.apache.drill.exec.oauth.PersistentTokenTable)4 Response (okhttp3.Response)3 MockResponse (okhttp3.mockwebserver.MockResponse)3 MockWebServer (okhttp3.mockwebserver.MockWebServer)3 RowSetBuilder (org.apache.drill.exec.physical.rowSet.RowSetBuilder)3 SchemaBuilder (org.apache.drill.exec.record.metadata.SchemaBuilder)3 ClusterFixtureBuilder (org.apache.drill.test.ClusterFixtureBuilder)3 ClusterTest (org.apache.drill.test.ClusterTest)3 Test (org.junit.Test)3 DirectRowSet (org.apache.drill.exec.physical.rowSet.DirectRowSet)2 RowSet (org.apache.drill.exec.physical.rowSet.RowSet)2 TupleMetadata (org.apache.drill.exec.record.metadata.TupleMetadata)2 BufferedReader (java.io.BufferedReader)1 InputStream (java.io.InputStream)1 InputStreamReader (java.io.InputStreamReader)1 HttpServletRequest (javax.servlet.http.HttpServletRequest)1 GET (javax.ws.rs.GET)1 Path (javax.ws.rs.Path)1