Examples with QueryRequest com.google.api.services.bigquery.model.QueryRequest used on opensource projects

Example 1 with QueryRequest

use of com.google.api.services.bigquery.model.QueryRequest in project beam by apache.

the class BigqueryMatcherTest method testQueryWithRetriesWhenServiceFails.

@Test
public void testQueryWithRetriesWhenServiceFails() throws Exception {
    BigqueryMatcher matcher = spy(new BigqueryMatcher(appName, projectId, query, "some-checksum"));
    when(mockQuery.execute()).thenThrow(new IOException());
    thrown.expect(RuntimeException.class);
    thrown.expectMessage("Unable to get BigQuery response after retrying");
    try {
        matcher.queryWithRetries(mockBigqueryClient, new QueryRequest(), fastClock, BackOffAdapter.toGcpBackOff(BigqueryMatcher.BACKOFF_FACTORY.backoff()));
    } finally {
        verify(mockJobs, atLeast(BigqueryMatcher.MAX_QUERY_RETRIES)).query(eq(projectId), eq(new QueryRequest()));
    }
}

Example 2 with QueryRequest

use of com.google.api.services.bigquery.model.QueryRequest in project beam by apache.

the class BigqueryMatcherTest method testBigqueryMatcherFailsForChecksumMismatch.

@Test
public void testBigqueryMatcherFailsForChecksumMismatch() throws IOException {
    BigqueryMatcher matcher = spy(new BigqueryMatcher(appName, projectId, query, "incorrect-checksum"));
    doReturn(mockBigqueryClient).when(matcher).newBigqueryClient(anyString());
    when(mockQuery.execute()).thenReturn(createResponseContainingTestData());
    thrown.expect(AssertionError.class);
    thrown.expectMessage("Total number of rows are: 1");
    thrown.expectMessage("abc");
    try {
        assertThat(mockResult, matcher);
    } finally {
        verify(matcher).newBigqueryClient(eq(appName));
        verify(mockJobs).query(eq(projectId), eq(new QueryRequest().setQuery(query)));
    }
}

Example 3 with QueryRequest

use of com.google.api.services.bigquery.model.QueryRequest in project beam by apache.

the class BigqueryMatcherTest method testQueryWithRetriesWhenQueryResponseNull.

@Test
public void testQueryWithRetriesWhenQueryResponseNull() throws Exception {
    BigqueryMatcher matcher = spy(new BigqueryMatcher(appName, projectId, query, "some-checksum"));
    when(mockQuery.execute()).thenReturn(null);
    thrown.expect(RuntimeException.class);
    thrown.expectMessage("Unable to get BigQuery response after retrying");
    try {
        matcher.queryWithRetries(mockBigqueryClient, new QueryRequest(), fastClock, BackOffAdapter.toGcpBackOff(BigqueryMatcher.BACKOFF_FACTORY.backoff()));
    } finally {
        verify(mockJobs, atLeast(BigqueryMatcher.MAX_QUERY_RETRIES)).query(eq(projectId), eq(new QueryRequest()));
    }
}

Example 4 with QueryRequest

use of com.google.api.services.bigquery.model.QueryRequest in project beam by apache.

the class BigqueryMatcherTest method testBigqueryMatcherThatSucceeds.

@Test
public void testBigqueryMatcherThatSucceeds() throws Exception {
    BigqueryMatcher matcher = spy(new BigqueryMatcher(appName, projectId, query, "9bb47f5c90d2a99cad526453dff5ed5ec74650dc"));
    doReturn(mockBigqueryClient).when(matcher).newBigqueryClient(anyString());
    when(mockQuery.execute()).thenReturn(createResponseContainingTestData());
    assertThat(mockResult, matcher);
    verify(matcher).newBigqueryClient(eq(appName));
    verify(mockJobs).query(eq(projectId), eq(new QueryRequest().setQuery(query)));
}

Example 5 with QueryRequest

use of com.google.api.services.bigquery.model.QueryRequest in project zeppelin by apache.

the class BigQueryInterpreter method run.

//Function to run the SQL on bigQuery service
public static Iterator<GetQueryResultsResponse> run(final String queryString, final String projId, final long wTime, final long maxRows) throws IOException {
    try {
        QueryResponse query = service.jobs().query(projId, new QueryRequest().setTimeoutMs(wTime).setQuery(queryString).setMaxResults(maxRows)).execute();
        jobId = query.getJobReference().getJobId();
        projectId = query.getJobReference().getProjectId();
        GetQueryResults getRequest = service.jobs().getQueryResults(projectId, jobId);
        return getPages(getRequest);
    } catch (IOException ex) {
        throw ex;
    }
}