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

Example 6 with QueryRequest

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

the class BigqueryMatcher method matchesSafely.

@Override
protected boolean matchesSafely(PipelineResult pipelineResult) {
    LOG.info("Verifying Bigquery data");
    Bigquery bigqueryClient = newBigqueryClient(applicationName);
    // execute query
    LOG.debug("Executing query: {}", query);
    try {
        QueryRequest queryContent = new QueryRequest();
        queryContent.setQuery(query);
        response = queryWithRetries(bigqueryClient, queryContent, Sleeper.DEFAULT, BackOffAdapter.toGcpBackOff(BACKOFF_FACTORY.backoff()));
    } catch (IOException | InterruptedException e) {
        if (e instanceof InterruptedIOException) {
            Thread.currentThread().interrupt();
        }
        throw new RuntimeException("Failed to fetch BigQuery data.", e);
    }
    if (!response.getJobComplete()) {
        // query job not complete, verification failed
        return false;
    } else {
        // compute checksum
        actualChecksum = generateHash(response.getRows());
        LOG.debug("Generated a SHA1 checksum based on queried data: {}", actualChecksum);
        return expectedChecksum.equals(actualChecksum);
    }
}

Example 7 with QueryRequest

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

the class BigqueryMatcherTest method testBigqueryMatcherFailsWhenQueryJobNotComplete.

@Test
public void testBigqueryMatcherFailsWhenQueryJobNotComplete() throws Exception {
    BigqueryMatcher matcher = spy(new BigqueryMatcher(appName, projectId, query, "some-checksum"));
    doReturn(mockBigqueryClient).when(matcher).newBigqueryClient(anyString());
    when(mockQuery.execute()).thenReturn(new QueryResponse().setJobComplete(false));
    thrown.expect(AssertionError.class);
    thrown.expectMessage("The query job hasn't completed.");
    thrown.expectMessage("jobComplete=false");
    try {
        assertThat(mockResult, matcher);
    } finally {
        verify(matcher).newBigqueryClient(eq(appName));
        verify(mockJobs).query(eq(projectId), eq(new QueryRequest().setQuery(query)));
    }
}