How to use

queryWithRetries

method

in

org.apache.beam.sdk.io.gcp.testing.BigqueryClient

@Nonnull
public QueryResponse queryWithRetries(String query, String projectId)
  throws IOException, InterruptedException {
 return queryWithRetries(query, projectId, false);
}

@Override
protected boolean matchesSafely(PipelineResult pipelineResult) {
 LOG.info("Verifying Bigquery data");
 // execute query
 LOG.debug("Executing query: {}", query);
 try {
  response = bigqueryClient.queryWithRetries(query, this.projectId);
 } 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);
 }
}

@Test
public void testBigqueryMatcherThatSucceeds() throws Exception {
 BigqueryMatcher matcher =
   spy(
     new BigqueryMatcher(
       appName, projectId, query, "9bb47f5c90d2a99cad526453dff5ed5ec74650dc"));
 when(mockBigqueryClient.queryWithRetries(anyString(), anyString()))
   .thenReturn(createResponseContainingTestData());
 assertThat(mockResult, matcher);
}

private void verifyLegacyQueryRes() throws Exception {
 LOG.info("Starting verifyLegacyQueryRes in outputTable {}", outputTable);
 List<String> legacyQueryExpectedRes = ImmutableList.of("apple", "orange");
 QueryResponse response =
   bqClient.queryWithRetries(String.format("SELECT fruit from [%s];", outputTable), project);
 LOG.info("Finished to query result table {}", this.outputTable);
 List<String> tableResult =
   response
     .getRows()
     .stream()
     .flatMap(row -> row.getF().stream().map(cell -> cell.getV().toString()))
     .sorted()
     .collect(Collectors.toList());
 assertEquals(legacyQueryExpectedRes, tableResult);
}

@Test
public void testQueryWithRetriesWhenServiceFails() throws Exception {
 when(mockQuery.execute()).thenThrow(new IOException());
 thrown.expect(RuntimeException.class);
 thrown.expectMessage("Unable to get BigQuery response after retrying");
 try {
  bqClient.queryWithRetries(query, projectId);
 } finally {
  verify(mockJobs, atLeast(BigqueryClient.MAX_QUERY_RETRIES))
    .query(eq(projectId), any(QueryRequest.class));
 }
}

@Test
public void testBigqueryMatcherFailsForChecksumMismatch() throws Exception {
 BigqueryMatcher matcher =
   spy(new BigqueryMatcher(appName, projectId, query, "incorrect-checksum"));
 when(mockBigqueryClient.queryWithRetries(anyString(), anyString()))
   .thenReturn(createResponseContainingTestData());
 thrown.expect(AssertionError.class);
 thrown.expectMessage("Total number of rows are: 1");
 thrown.expectMessage("abc");
 assertThat(mockResult, matcher);
}

 @Test
 public void testQueryWithRetriesWhenQueryResponseNull() throws Exception {
  when(mockQuery.execute()).thenReturn(null);

  thrown.expect(RuntimeException.class);
  thrown.expectMessage("Unable to get BigQuery response after retrying");
  try {
   bqClient.queryWithRetries(query, projectId);
  } finally {
   verify(mockJobs, atLeast(BigqueryClient.MAX_QUERY_RETRIES))
     .query(eq(projectId), any(QueryRequest.class));
  }
 }
}

@Test
public void testBigqueryMatcherFailsWhenQueryJobNotComplete() throws Exception {
 BigqueryMatcher matcher = spy(new BigqueryMatcher(appName, projectId, query, "some-checksum"));
 when(mockBigqueryClient.queryWithRetries(anyString(), anyString()))
   .thenReturn(new QueryResponse().setJobComplete(false));
 thrown.expect(AssertionError.class);
 thrown.expectMessage("The query job hasn't completed.");
 thrown.expectMessage("jobComplete=false");
 assertThat(mockResult, matcher);
}

    "xyw=,2011-01-01,23:59:59.999999");
QueryResponse response =
  bqClient.queryWithRetries(
    String.format("SELECT bytes, date, time FROM [%s];", this.outputTable), this.project);
LOG.info("Finished to query result table {}", this.outputTable);

 @Test
 public void testE2ETrafficRoutes() throws Exception {
  this.options.setBigQuerySchema(FormatStatsFn.getSchema());
  this.options.setProject(this.projectId);
  this.options.setBigQueryDataset(this.outputDatasetId);
  this.options.setBigQueryTable(this.outputTable);
  TrafficRoutes.runTrafficRoutes(options);

  QueryResponse response =
    this.bqClient.queryWithRetries(
      String.format(
        "SELECT count(*) as total FROM [%s:%s.%s]",
        this.projectId, this.outputDatasetId, this.outputTable),
      this.projectId);
  String res = response.getRows().get(0).getF().get(0).getV().toString();
  assertEquals("27", res);
 }
}

 @Test
 public void testE2ETrafficMaxLaneFlow() throws Exception {
  this.options.setBigQuerySchema(FormatMaxesFn.getSchema());
  this.options.setProject(this.projectId);
  this.options.setBigQueryDataset(this.outputDatasetId);
  this.options.setBigQueryTable(this.outputTable);
  TrafficMaxLaneFlow.runTrafficMaxLaneFlow(this.options);

  QueryResponse response =
    this.bqClient.queryWithRetries(
      String.format(
        "SELECT count(*) as total FROM [%s:%s.%s]",
        this.projectId, this.outputDatasetId, this.outputTable),
      this.projectId);
  String res = response.getRows().get(0).getF().get(0).getV().toString();
  assertEquals("9763", res);
 }
}

.queryWithRetries(options.getInput(), bigQueryOptions.getProject(), true)
.getRows()
.get(0);