com.google.cloud.bigquery.TableResult Java Examples

@Test
public void PutBigQueryStreamingFullError() throws Exception {
    String tableName = Thread.currentThread().getStackTrace()[1].getMethodName();
    createTable(tableName);

    runner.setProperty(BigQueryAttributes.DATASET_ATTR, dataset.getDatasetId().getDataset());
    runner.setProperty(BigQueryAttributes.TABLE_NAME_ATTR, tableName);

    final JsonTreeReader jsonReader = new JsonTreeReader();
    runner.addControllerService("reader", jsonReader);
    runner.enableControllerService(jsonReader);

    runner.setProperty(BigQueryAttributes.RECORD_READER_ATTR, "reader");

    runner.enqueue(Paths.get("src/test/resources/bigquery/streaming-bad-data.json"));

    runner.run();
    runner.assertAllFlowFilesTransferred(PutBigQueryStreaming.REL_FAILURE, 1);
    runner.getFlowFilesForRelationship(PutBigQueryStreaming.REL_FAILURE).get(0).assertAttributeEquals(BigQueryAttributes.JOB_NB_RECORDS_ATTR, "0");

    TableResult result = bigquery.listTableData(dataset.getDatasetId().getDataset(), tableName, schema);
    assertFalse(result.getValues().iterator().hasNext());

    deleteTable(tableName);
}
 

@Test
public void runQuery() throws Exception {
  QueryJobConfiguration queryConfig = QueryJobConfiguration.newBuilder("SELECT * FROM [sample:table] LIMIT 1000")
      .setUseQueryCache(true)
      .setUseLegacySql(useLegacySql)
      .build();

  TableResult mockQueryResponse = mock(TableResult.class);
  Job mockJob = mock(Job.class);
  JobStatus mockJobStatus = mock(JobStatus.class);

  // First pretend we haven't finished running the query, second time around its completed.
  when(mockJob.isDone()).thenReturn(false).thenReturn(true);
  when(mockJob.getJobId()).thenReturn(jobId);
  when(mockJobStatus.getError()).thenReturn(null);
  when(mockJob.getStatus()).thenReturn(mockJobStatus);

  when(mockBigquery.create((JobInfo)any())).thenReturn(mockJob);
  when(mockBigquery.cancel(jobId)).thenReturn(true);
  when(mockJob.getQueryResults()).thenReturn(mockQueryResponse);

  BigQueryDelegate delegate = new BigQueryDelegate(mockBigquery, useLegacySql);
  delegate.runQuery(queryConfig, 1000, 1000);
}
 

@Test
public void testWriteAndListTableData()
    throws IOException, InterruptedException, TimeoutException, URISyntaxException {
  // Create table
  String tableName = "test_write_and_list_table_data";
  String fieldName = "string_field";
  assertNotNull(bigquerySnippets.createTable(DATASET, tableName, fieldName));
  // Add rows from file
  Path csvPath =
      Paths.get(Resources.getResource("bigquery/test_write_and_list_table_data.csv").toURI());
  long outputRows = bigquerySnippets.writeFileToTable(DATASET, tableName, csvPath, "us");
  assertEquals(2L, outputRows);
  // List all rows
  TableResult tableData = bigquerySnippets.listTableData(DATASET, tableName);
  String tableDataString = tableData.toString();
  assertTrue(tableDataString.contains("StringValue3"));
  assertTrue(tableDataString.contains("StringValue4"));
  assertTrue(bigquerySnippets.deleteTable(DATASET, tableName));
}
 

/** Example of listing table rows with schema. */
// [TARGET listTableData(TableId, Schema, TableDataListOption...)]
public FieldValueList listTableDataSchemaId() {
  // [START ]
  Schema schema =
      Schema.of(
          Field.of("word", LegacySQLTypeName.STRING),
          Field.of("word_count", LegacySQLTypeName.STRING),
          Field.of("corpus", LegacySQLTypeName.STRING),
          Field.of("corpus_date", LegacySQLTypeName.STRING));
  TableResult tableData =
      bigquery.listTableData(
          TableId.of("bigquery-public-data", "samples", "shakespeare"), schema);
  FieldValueList row = tableData.getValues().iterator().next();
  System.out.println(row.get("word").getStringValue());
  // [END ]
  return row;
}
 

@Test
public void testFileUpload() throws InterruptedException, IOException {
	LinkedMultiValueMap<String, Object> map = new LinkedMultiValueMap<>();
	map.add("file", csvFile);
	map.add("tableName", TABLE_NAME);

	HttpHeaders headers = new HttpHeaders();
	headers.setContentType(MediaType.MULTIPART_FORM_DATA);

	HttpEntity<LinkedMultiValueMap<String, Object>> request = new HttpEntity<>(map, headers);
	ResponseEntity<String> response =
			this.restTemplate.postForEntity("/uploadFile", request, String.class);
	assertThat(response.getStatusCode().is2xxSuccessful()).isTrue();

	QueryJobConfiguration queryJobConfiguration = QueryJobConfiguration
			.newBuilder("SELECT * FROM " + DATASET_NAME + "." + TABLE_NAME)
			.build();

	TableResult queryResult = this.bigQuery.query(queryJobConfiguration);
	assertThat(queryResult.getTotalRows()).isEqualTo(3);

	List<String> names = StreamSupport.stream(queryResult.getValues().spliterator(), false)
			.map(valueList -> valueList.get(0).getStringValue())
			.collect(Collectors.toList());
	assertThat(names).containsExactlyInAnyOrder("Nathaniel", "Diaz", "Johnson");
}
 

@Test
public void testCsvDataUpload() throws InterruptedException {
	LinkedMultiValueMap<String, Object> map = new LinkedMultiValueMap<>();
	map.add("csvText", "name,age,location\nBob,24,Wyoming");
	map.add("tableName", TABLE_NAME);

	HttpHeaders headers = new HttpHeaders();
	HttpEntity<LinkedMultiValueMap<String, Object>> request = new HttpEntity<>(map, headers);
	ResponseEntity<String> response =
			this.restTemplate.postForEntity("/uploadCsvText", request, String.class);
	assertThat(response.getStatusCode().is2xxSuccessful()).isTrue();

	QueryJobConfiguration queryJobConfiguration = QueryJobConfiguration
			.newBuilder("SELECT * FROM " + DATASET_NAME + "." + TABLE_NAME)
			.build();

	TableResult queryResult = this.bigQuery.query(queryJobConfiguration);
	assertThat(queryResult.getTotalRows()).isEqualTo(1);

	FieldValueList row = queryResult.getValues().iterator().next();
	assertThat(row.get(0).getStringValue()).isEqualTo("Bob");
	assertThat(row.get(1).getLongValue()).isEqualTo(24);
	assertThat(row.get(2).getStringValue()).isEqualTo("Wyoming");
}
 

/** Example of listing table rows with schema. */
// [TARGET listTableData(String, String, Schema, TableDataListOption...)]
// [VARIABLE "my_dataset_name"]
// [VARIABLE "my_table_name"]
// [VARIABLE ...]
// [VARIABLE "field"]
public TableResult listTableDataSchema(
    String datasetName, String tableName, Schema schema, String field) {
  // [START ]
  TableResult tableData = bigquery.listTableData(datasetName, tableName, schema);
  for (FieldValueList row : tableData.iterateAll()) {
    row.get(field);
  }
  // [END ]
  return tableData;
}
 

@Test
public void testLoadFile_sync() throws InterruptedException {
	this.messageHandler.setSync(true);

	HashMap<String, Object> messageHeaders = new HashMap<>();
	messageHeaders.put(BigQuerySpringMessageHeaders.TABLE_NAME, TABLE_NAME);
	messageHeaders.put(BigQuerySpringMessageHeaders.FORMAT_OPTIONS, FormatOptions.csv());

	Message<File> message = MessageBuilder.createMessage(
			new File("src/test/resources/data.csv"),
			new MessageHeaders(messageHeaders));

	Job job = (Job) this.messageHandler.handleRequestMessage(message);
	assertThat(job).isNotNull();

	QueryJobConfiguration queryJobConfiguration = QueryJobConfiguration
			.newBuilder("SELECT * FROM test_dataset.test_table").build();
	TableResult result = this.bigquery.query(queryJobConfiguration);
	assertThat(result.getTotalRows()).isEqualTo(1);
}
 

@Test
public void testLoadFile() throws IOException, ExecutionException, InterruptedException {
	ListenableFuture<Job> bigQueryJobFuture =
			bigQueryTemplate.writeDataToTable(TABLE_NAME, dataFile.getInputStream(), FormatOptions.csv());

	Job job = bigQueryJobFuture.get();
	assertThat(job.getStatus().getState()).isEqualTo(JobStatus.State.DONE);

	QueryJobConfiguration queryJobConfiguration = QueryJobConfiguration
			.newBuilder("SELECT * FROM test_dataset.template_test_table").build();
	TableResult result = this.bigQuery.query(queryJobConfiguration);

	assertThat(result.getTotalRows()).isEqualTo(1);
	assertThat(
			result.getValues().iterator().next().get("State").getStringValue()).isEqualTo("Alabama");
}
 

@Test
public void testLoadBytes() throws ExecutionException, InterruptedException {
	byte[] byteArray =
			"CountyId,State,County\n1001,Alabama,Autauga County\n".getBytes();
	ByteArrayInputStream byteStream = new ByteArrayInputStream(byteArray);

	ListenableFuture<Job> bigQueryJobFuture =
			bigQueryTemplate.writeDataToTable(TABLE_NAME, byteStream, FormatOptions.csv());

	Job job = bigQueryJobFuture.get();
	assertThat(job.getStatus().getState()).isEqualTo(JobStatus.State.DONE);

	QueryJobConfiguration queryJobConfiguration = QueryJobConfiguration
			.newBuilder("SELECT * FROM test_dataset.template_test_table").build();
	TableResult result = this.bigQuery.query(queryJobConfiguration);

	assertThat(result.getTotalRows()).isEqualTo(1);
	assertThat(
			result.getValues().iterator().next().get("State").getStringValue()).isEqualTo("Alabama");
}
 

/** Example of listing table rows, specifying the page size. */
// [TARGET listTableData(TableId, TableDataListOption...)]
// [VARIABLE "my_dataset_name"]
// [VARIABLE "my_table_name"]
public TableResult listTableDataFromId(String datasetName, String tableName) {
  // [START bigquery_browse_table]
  TableId tableIdObject = TableId.of(datasetName, tableName);
  // This example reads the result 100 rows per RPC call. If there's no need to limit the number,
  // simply omit the option.
  TableResult tableData =
      bigquery.listTableData(tableIdObject, TableDataListOption.pageSize(100));
  for (FieldValueList row : tableData.iterateAll()) {
    // do something with the row
  }
  // [END bigquery_browse_table]
  return tableData;
}
 

/** Example of listing table rows, specifying the page size. */
// [TARGET listTableData(String, String, TableDataListOption...)]
// [VARIABLE "my_dataset_name"]
// [VARIABLE "my_table_name"]
public TableResult listTableData(String datasetName, String tableName) {
  // [START ]
  // This example reads the result 100 rows per RPC call. If there's no need to limit the number,
  // simply omit the option.
  TableResult tableData =
      bigquery.listTableData(datasetName, tableName, TableDataListOption.pageSize(100));
  for (FieldValueList row : tableData.iterateAll()) {
    // do something with the row
  }
  // [END ]
  return tableData;
}
 

private TableResult queryWithLarge(BigQuery bigquery, QueryJobConfiguration queryRequest, String projectId,
                                   BigQuery.JobOption... options) {
    String tempDataset = genTempName("dataset");
    String tempTable = genTempName("table");
    bigquery.create(DatasetInfo.of(tempDataset));
    TableId tableId = TableId.of(projectId, tempDataset, tempTable);
    QueryJobConfiguration jobConfiguration = QueryJobConfiguration
            .newBuilder(queryRequest.getQuery())
            .setAllowLargeResults(true)
            .setUseLegacySql(queryRequest.useLegacySql())
            .setDestinationTable(tableId)
            .build();
    try {
        return query(bigquery, jobConfiguration, projectId, options);
    } finally {
        bigquery.delete(tableId);
    }
}
 

/**
 * Get the schema by table name or query. This method also needed for read and write, because we can not get schema
 * from the ongoing data. BigQueryIO.Read return TableRow, which do not include schema in itself. So use BigQuery
 * client to get it before read and write.
 * 
 * @return
 */
@Override
public Schema getSchema() {
    BigQuery bigquery = BigQueryConnection.createClient(properties.getDatastoreProperties());
    com.google.cloud.bigquery.Schema bqRowSchema = null;
    switch (properties.sourceType.getValue()) {
    case TABLE_NAME: {
        TableId tableId = TableId.of(properties.getDatastoreProperties().projectName.getValue(),
                properties.bqDataset.getValue(), properties.tableName.getValue());
        Table table = bigquery.getTable(tableId);
        if (table == null) {
            ComponentException.build(CommonErrorCodes.UNEXPECTED_EXCEPTION).setAndThrow(
                    "Table not found:" + tableId.toString());
        }
        bqRowSchema = table.getDefinition().getSchema();
        break;
    }
    case QUERY: {
        QueryJobConfiguration queryRequest = QueryJobConfiguration
                .newBuilder(properties.query.getValue())
                .setUseLegacySql(properties.useLegacySql.getValue())
                .build();
        TableResult queryResponse =
                query(bigquery, queryRequest, properties.getDatastoreProperties().projectName.getValue());
        bqRowSchema = queryResponse.getSchema();
        break;
    }
    default:
        throw new RuntimeException("To be implemented: " + properties.sourceType.getValue());
    }
    return BigQueryAvroRegistry.get().inferSchema(bqRowSchema);
}
 

@Test
public void PutBigQueryStreamingPartialError() throws Exception {
    String tableName = Thread.currentThread().getStackTrace()[1].getMethodName();
    createTable(tableName);

    runner.setProperty(BigQueryAttributes.DATASET_ATTR, dataset.getDatasetId().getDataset());
    runner.setProperty(BigQueryAttributes.TABLE_NAME_ATTR, tableName);

    final JsonTreeReader jsonReader = new JsonTreeReader();
    runner.addControllerService("reader", jsonReader);
    runner.enableControllerService(jsonReader);

    runner.setProperty(BigQueryAttributes.RECORD_READER_ATTR, "reader");
    runner.setProperty(BigQueryAttributes.SKIP_INVALID_ROWS_ATTR, "true");

    runner.enqueue(Paths.get("src/test/resources/bigquery/streaming-bad-data.json"));

    runner.run();
    runner.assertAllFlowFilesTransferred(PutBigQueryStreaming.REL_FAILURE, 1);
    runner.getFlowFilesForRelationship(PutBigQueryStreaming.REL_FAILURE).get(0).assertAttributeEquals(BigQueryAttributes.JOB_NB_RECORDS_ATTR, "1");

    TableResult result = bigquery.listTableData(dataset.getDatasetId().getDataset(), tableName, schema);
    Iterator<FieldValueList> iterator = result.getValues().iterator();

    FieldValueList firstElt = iterator.next();
    assertFalse(iterator.hasNext());
    assertEquals(firstElt.get("name").getStringValue(), "Jane Doe");

    deleteTable(tableName);
}
 

private TableResult query(BigQuery bigquery, QueryJobConfiguration queryRequest, String projectId,
                          BigQuery.JobOption... options) {
    TableResult queryResponse = null;
    try {
        queryResponse = bigquery.query(queryRequest, options);
    } catch (BigQueryException exception) {
        if ("responseTooLarge".equals(exception.getReason())) {
            return queryWithLarge(bigquery, queryRequest, projectId, options);
        }
    } catch (final InterruptedException e) {
        Thread.currentThread().interrupt();
    }
    return loopQueryResponse(queryResponse);
}
 

private TableResult loopQueryResponse(TableResult queryResponse) {
    if (queryResponse == null) {
        TalendRuntimeException.build(ComponentsErrorCode.IO_EXCEPTION).throwIt();
    }
    /* JobException now
    if (queryResponse.hasErrors()) {
        TalendRuntimeException.build(ComponentsErrorCode.IO_EXCEPTION).setAndThrow(
                queryResponse.getExecutionErrors().toArray(new String[] {}));
    }
    */
    return queryResponse;
}
 

@Test
public void PutBigQueryStreamingNoError() throws Exception {
    String tableName = Thread.currentThread().getStackTrace()[1].getMethodName();
    createTable(tableName);

    runner.setProperty(BigQueryAttributes.DATASET_ATTR, dataset.getDatasetId().getDataset());
    runner.setProperty(BigQueryAttributes.TABLE_NAME_ATTR, tableName);

    final JsonTreeReader jsonReader = new JsonTreeReader();
    runner.addControllerService("reader", jsonReader);
    runner.enableControllerService(jsonReader);

    runner.setProperty(BigQueryAttributes.RECORD_READER_ATTR, "reader");

    runner.enqueue(Paths.get("src/test/resources/bigquery/streaming-correct-data.json"));

    runner.run();
    runner.assertAllFlowFilesTransferred(PutBigQueryStreaming.REL_SUCCESS, 1);
    runner.getFlowFilesForRelationship(PutBigQueryStreaming.REL_SUCCESS).get(0).assertAttributeEquals(BigQueryAttributes.JOB_NB_RECORDS_ATTR, "2");

    TableResult result = bigquery.listTableData(dataset.getDatasetId().getDataset(), tableName, schema);
    Iterator<FieldValueList> iterator = result.getValues().iterator();

    FieldValueList firstElt = iterator.next();
    FieldValueList sndElt = iterator.next();
    assertTrue(firstElt.get("name").getStringValue().endsWith("Doe"));
    assertTrue(sndElt.get("name").getStringValue().endsWith("Doe"));

    FieldValueList john;
    FieldValueList jane;
    john = firstElt.get("name").getStringValue().equals("John Doe") ? firstElt : sndElt;
    jane = firstElt.get("name").getStringValue().equals("Jane Doe") ? firstElt : sndElt;

    assertEquals(jane.get("job").getRecordValue().get(0).getStringValue(), "Director");
    assertTrue(john.get("alias").getRepeatedValue().size() == 2);
    assertTrue(john.get("addresses").getRepeatedValue().get(0).getRecordValue().get(0).getStringValue().endsWith("000"));

    deleteTable(tableName);
}
 

TableResult query(String sql)
{
    try {
        return bigQuery.query(QueryJobConfiguration.of(sql));
    }
    catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        throw new BigQueryException(BaseHttpServiceException.UNKNOWN_CODE, format("Failed to run the query [%s]", sql), e);
    }
}
 

@Test(expected = StageException.class)
public void runQueryHasErrors() throws Exception {
  QueryJobConfiguration queryRequest = QueryJobConfiguration.newBuilder("SELECT * FROM [sample:table] LIMIT 1000")
      .setUseQueryCache(true)
      .setUseLegacySql(useLegacySql)
      .build();

  TableResult mockQueryResponse = mock(TableResult.class);
  Job mockJob = mock(Job.class);
  JobStatus mockJobStatus = mock(JobStatus.class);

  // First pretend we haven't finished running the query, second time around its completed.
  when(mockJob.isDone()).thenReturn(true);
  when(mockJob.getJobId()).thenReturn(jobId);

  when(mockJob.getQueryResults()).thenReturn(mockQueryResponse);
  when(mockJobStatus.getError()).thenReturn(new BigQueryError(
      "Some Error",
      "Some Location",
      "Some Error Message"
  ));
  when(mockJob.getStatus()).thenReturn(mockJobStatus);

  when(mockBigquery.create((JobInfo)any())).thenReturn(mockJob);
  when(mockBigquery.cancel(jobId)).thenReturn(true);

  BigQueryDelegate delegate = new BigQueryDelegate(mockBigquery, useLegacySql);

  ErrorCode code = null;
  try {
    delegate.runQuery(queryRequest, 1000, 1000);
  } catch (StageException e) {
    code = e.getErrorCode();
    throw e;
  } finally {
    assertEquals(Errors.BIGQUERY_02, code);
  }
}
 

@Test(expected = StageException.class)
public void runQueryTimeout() throws Exception {
  QueryJobConfiguration queryConfig = QueryJobConfiguration.newBuilder("SELECT * FROM [sample:table] LIMIT 1000")
      .setUseQueryCache(true)
      .setUseLegacySql(useLegacySql)
      .build();

  TableResult mockQueryResponse = mock(TableResult.class);
  Job mockJob = mock(Job.class);
  JobStatus mockJobStatus = mock(JobStatus.class);

  // First pretend we haven't finished running the query, second time around its completed.
  when(mockJob.isDone()).thenReturn(false).thenReturn(true);
  when(mockJob.getJobId()).thenReturn(jobId);
  when(mockJobStatus.getError()).thenReturn(null);
  when(mockJob.getStatus()).thenReturn(mockJobStatus);

  when(mockBigquery.create((JobInfo)any())).thenReturn(mockJob);
  when(mockBigquery.cancel(jobId)).thenReturn(true);
  when(mockJob.getQueryResults()).thenReturn(mockQueryResponse);

  BigQueryDelegate delegate = new BigQueryDelegate(
      mockBigquery,
      useLegacySql,
      Clock.offset(Clock.systemDefaultZone(), Duration.ofSeconds(2))
  );

  ErrorCode code = null;
  try {
    delegate.runQuery(queryConfig, 1000, 1000);
  } catch (StageException e) {
    code = e.getErrorCode();
    throw e;
  } finally {
    assertEquals(Errors.BIGQUERY_00, code);
  }
}
 

@Before
public void setUp() throws Exception {
  mockBigquery = mock(BigQuery.class);
  jobId = JobId.of("test-project", "datacollector");
  mockResult = mock(TableResult.class);

  mockStatic(ServiceAccountCredentials.class);
  mockStatic(GoogleCredentials.class);
  ServiceAccountCredentials serviceAccountCredentials = mock(ServiceAccountCredentials.class);
  when(ServiceAccountCredentials.fromStream(any())).thenReturn(serviceAccountCredentials);
}
 

private static String convertRunToHtmlTable(TableResult result) {
  if (result == null) {
    return "";
  }

  StringBuilder sb = new StringBuilder();
  for (FieldValueList row : result.iterateAll()) {
    sb.append("<tr>");
    String url = row.get("url").getStringValue();
    addColumn(sb, String.format("<a href=\"%s\">%s</a>", url, url));
    addColumn(sb, row.get("view_count").getLongValue());
    sb.append("</tr>");
  }
  return sb.toString();
}
 

@Test
public void testLoadFile() throws InterruptedException, ExecutionException {
	HashMap<String, Object> messageHeaders = new HashMap<>();
	this.messageHandler.setTableName(TABLE_NAME);
	this.messageHandler.setFormatOptions(FormatOptions.csv());

	Message<File> message = MessageBuilder.createMessage(
			new File("src/test/resources/data.csv"),
			new MessageHeaders(messageHeaders));

	ListenableFuture<Job> jobFuture =
			(ListenableFuture<Job>) this.messageHandler.handleRequestMessage(message);

	// Assert that a BigQuery polling task is scheduled successfully.
	await().atMost(Duration.FIVE_SECONDS)
			.untilAsserted(
					() -> assertThat(
							this.taskScheduler.getScheduledThreadPoolExecutor().getQueue()).hasSize(1));
	jobFuture.get();

	QueryJobConfiguration queryJobConfiguration = QueryJobConfiguration
			.newBuilder("SELECT * FROM test_dataset.test_table").build();
	TableResult result = this.bigquery.query(queryJobConfiguration);

	assertThat(result.getTotalRows()).isEqualTo(1);
	assertThat(
			result.getValues().iterator().next().get("State").getStringValue()).isEqualTo("Alabama");

	// This asserts that the BigQuery job polling task is no longer in the scheduler.
	assertThat(this.taskScheduler.getScheduledThreadPoolExecutor().getQueue()).hasSize(0);
}
 

private Job issueQueryToBQ(String statement) throws InterruptedException {
  QueryJobConfiguration jobConfiguration = QueryJobConfiguration.newBuilder(statement)
      .build();

  String jobId = makeJobId(jobIdPrefix, statement);

  LOG.info("Triggering job {} for statement |{}|", jobId, statement);

  TableResult result = bigQueryClient.query(jobConfiguration, JobId.of(jobId));
  return bigQueryClient.getJob(JobId.of(jobId));
}
 

private List<BigQuerySplit> createEmptyProjection(TableId tableId, int actualParallelism, Optional<String> filter)
{
    log.debug("createEmptyProjection(tableId=%s, actualParallelism=%s, filter=[%s])", tableId, actualParallelism, filter);
    try {
        long numberOfRows;
        if (filter.isPresent()) {
            // count the rows based on the filter
            String sql = bigQueryClient.selectSql(tableId, "COUNT(*)", new String[] {filter.get()});
            TableResult result = bigQueryClient.query(sql);
            numberOfRows = result.iterateAll().iterator().next().get(0).getLongValue();
        }
        else {
            // no filters, so we can take the value from the table info
            numberOfRows = bigQueryClient.getTable(tableId).getNumRows().longValue();
        }

        long rowsPerSplit = numberOfRows / actualParallelism;
        long remainingRows = numberOfRows - (rowsPerSplit * actualParallelism); // need to be added to one fo the split due to integer division
        List<BigQuerySplit> splits = range(0, actualParallelism)
                .mapToObj(ignored -> BigQuerySplit.emptyProjection(rowsPerSplit))
                .collect(toList());
        splits.set(0, BigQuerySplit.emptyProjection(rowsPerSplit + remainingRows));
        return splits;
    }
    catch (BigQueryException e) {
        throw new PrestoException(BIGQUERY_FAILED_TO_EXECUTE_QUERY, "Failed to compute empty projection", e);
    }
}
 

@Override
public void getSample(int limit, Consumer<IndexedRecord> consumer) {
    BigQuery bigquery = BigQueryConnection.createClient(properties.getDatastoreProperties());
    com.google.cloud.bigquery.Schema bqRowSchema = null;
    String query = null;
    boolean useLegacySql = true;
    switch (properties.sourceType.getValue()) {
    case TABLE_NAME: {
        query = String.format("select * from `%s.%s.%s` LIMIT %d",
                properties.getDatastoreProperties().projectName.getValue(), properties.bqDataset.getValue(),
                properties.tableName.getValue(), limit);
        useLegacySql = false;
        break;
    }
    case QUERY: {
        query = properties.query.getValue();
        useLegacySql = properties.useLegacySql.getValue();
        break;
    }
    default:
        throw new RuntimeException("To be implemented: " + properties.sourceType.getValue());
    }
    QueryJobConfiguration queryRequest = QueryJobConfiguration
            .newBuilder(query)
            .setUseLegacySql(useLegacySql)
            .build();
    // todo: proper pagination, not critical for getSample yet
    TableResult queryResponse =
            query(bigquery, queryRequest, properties.getDatastoreProperties().projectName.getValue());
    bqRowSchema = queryResponse.getSchema();
    Schema schema = BigQueryAvroRegistry.get().inferSchema(bqRowSchema);
    Iterator<FieldValueList> iterator = queryResponse.getValues().iterator();
    IndexedRecordConverter<Map<String, Object>, IndexedRecord> converter =
            new BigQueryFieldValueListIndexedRecordConverter();
    converter.setSchema(schema);
    int count = 0; // need this only for legacy sql with large result
    while (iterator.hasNext() && count < limit) {
        List<FieldValue> values = iterator.next();
        consumer.accept(converter.convertToAvro(BigQueryAvroRegistry.get().convertFileds(values, schema)));
        count++;
    }
}
 

@VisibleForTesting
TableResult runQuery(QueryJobConfiguration queryRequest, long pageSize) throws StageException {
  return delegate.runQuery(queryRequest, conf.timeout * 1000, pageSize);
}
 

public void runQuery() throws InterruptedException {
  QueryJobConfiguration queryConfig =
      QueryJobConfiguration.newBuilder(
              "SELECT "
                  + "CONCAT('https://stackoverflow.com/questions/', CAST(id as STRING)) as url, "
                  + "view_count "
                  + "FROM `bigquery-public-data.stackoverflow.posts_questions` "
                  + "WHERE tags like '%google-bigquery%' "
                  + "ORDER BY favorite_count DESC LIMIT 10")
          // Use standard SQL syntax for queries.
          // See: https://cloud.google.com/bigquery/sql-reference/
          .setUseLegacySql(false)
          .build();

  List<TimeSeries> timeSeriesList = new ArrayList<>();

  long queryStartTime = System.currentTimeMillis();

  // Create a job ID so that we can safely retry.
  JobId jobId = JobId.of(UUID.randomUUID().toString());
  Job queryJob = bigquery.create(JobInfo.newBuilder(queryConfig).setJobId(jobId).build());

  // Wait for the query to complete.
  queryJob = queryJob.waitFor();

  // Check for errors
  if (queryJob == null) {
    throw new RuntimeException("Job no longer exists");
  } else if (queryJob.getStatus().getError() != null) {
    // You can also look at queryJob.getStatus().getExecutionErrors() for all
    // errors, not just the latest one.
    throw new RuntimeException(queryJob.getStatus().getError().toString());
  }

  // Log the result metrics.
  TableResult result = queryJob.getQueryResults();

  long queryEndTime = System.currentTimeMillis();
  // Add query duration metric.
  timeSeriesList.add(prepareMetric(QUERY_DURATION_METRIC, queryEndTime - queryStartTime));

  // Add rows returned metric.
  timeSeriesList.add(prepareMetric(ROWS_RETURNED_METRIC, result.getTotalRows()));

  // Prepares the time series request
  CreateTimeSeriesRequest request =
      CreateTimeSeriesRequest.newBuilder()
          .setName(projectName)
          .addAllTimeSeries(timeSeriesList)
          .build();

  createMetricsIfNeeded();
  client.createTimeSeries(request);
  os.println("Done writing metrics.");

  mostRecentRunResult = result;
}
 

public static TableResult getMostRecentRunResult() {
  return mostRecentRunResult;
}