Examples with DirectRowSet org.apache.drill.exec.physical.rowSet.DirectRowSet used on opensource projects

Example 41 with DirectRowSet

use of org.apache.drill.exec.physical.rowSet.DirectRowSet 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();
    }
}

Example 42 with DirectRowSet

use of org.apache.drill.exec.physical.rowSet.DirectRowSet 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();
    }
}