private void setupBigQueryTable(
      String projectId, String datasetId, String tableId, TableSchema schema) throws IOException {
    if (bigQueryClient == null) {
      bigQueryClient = Transport.newBigQueryClient(options.as(BigQueryOptions.class)).build();
    }

    Datasets datasetService = bigQueryClient.datasets();
    if (executeNullIfNotFound(datasetService.get(projectId, datasetId)) == null) {
      Dataset newDataset =
          new Dataset()
              .setDatasetReference(
                  new DatasetReference().setProjectId(projectId).setDatasetId(datasetId));
      datasetService.insert(projectId, newDataset).execute();
    }

    Tables tableService = bigQueryClient.tables();
    Table table = executeNullIfNotFound(tableService.get(projectId, datasetId, tableId));
    if (table == null) {
      Table newTable =
          new Table()
              .setSchema(schema)
              .setTableReference(
                  new TableReference()
                      .setProjectId(projectId)
                      .setDatasetId(datasetId)
                      .setTableId(tableId));
      tableService.insert(projectId, datasetId, newTable).execute();
    } else if (!table.getSchema().equals(schema)) {
      throw new RuntimeException(
          "Table exists and schemas do not match, expecting: "
              + schema.toPrettyString()
              + ", actual: "
              + table.getSchema().toPrettyString());
    }
  }
  /**
   * Retrieves or creates the table.
   *
   * <p>The table is checked to conform to insertion requirements as specified by WriteDisposition
   * and CreateDisposition.
   *
   * <p>If table truncation is requested (WriteDisposition.WRITE_TRUNCATE), then this will re-create
   * the table if necessary to ensure it is empty.
   *
   * <p>If an empty table is required (WriteDisposition.WRITE_EMPTY), then this will fail if the
   * table exists and is not empty.
   *
   * <p>When constructing a table, a {@code TableSchema} must be available. If a schema is provided,
   * then it will be used. If no schema is provided, but an existing table is being cleared
   * (WRITE_TRUNCATE option above), then the existing schema will be re-used. If no schema is
   * available, then an {@code IOException} is thrown.
   */
  public Table getOrCreateTable(
      WriteDisposition writeDisposition,
      CreateDisposition createDisposition,
      @Nullable TableSchema schema)
      throws IOException {
    // Check if table already exists.
    Bigquery.Tables.Get get =
        client.tables().get(ref.getProjectId(), ref.getDatasetId(), ref.getTableId());
    Table table = null;
    try {
      table = get.execute();
    } catch (IOException e) {
      ApiErrorExtractor errorExtractor = new ApiErrorExtractor();
      if (!errorExtractor.itemNotFound(e)
          || createDisposition != CreateDisposition.CREATE_IF_NEEDED) {
        // Rethrow.
        throw e;
      }
    }

    // If we want an empty table, and it isn't, then delete it first.
    if (table != null) {
      if (writeDisposition == WriteDisposition.WRITE_APPEND) {
        return table;
      }

      boolean empty = isEmpty();
      if (empty) {
        if (writeDisposition == WriteDisposition.WRITE_TRUNCATE) {
          LOG.info("Empty table found, not removing {}", BigQueryIO.toTableSpec(ref));
        }
        return table;

      } else if (writeDisposition == WriteDisposition.WRITE_EMPTY) {
        throw new IOException("WriteDisposition is WRITE_EMPTY, " + "but table is not empty");
      }

      // Reuse the existing schema if none was provided.
      if (schema == null) {
        schema = table.getSchema();
      }

      // Delete table and fall through to re-creating it below.
      LOG.info("Deleting table {}", BigQueryIO.toTableSpec(ref));
      Bigquery.Tables.Delete delete =
          client.tables().delete(ref.getProjectId(), ref.getDatasetId(), ref.getTableId());
      delete.execute();
    }

    if (schema == null) {
      throw new IllegalArgumentException("Table schema required for new table.");
    }

    // Create the table.
    return tryCreateTable(schema);
  }
  private void onInsertAll(List<List<Long>> errorIndicesSequence) throws Exception {
    when(mockClient.tabledata()).thenReturn(mockTabledata);

    List<TableDataInsertAllResponse> responses = new ArrayList<>();
    for (List<Long> errorIndices : errorIndicesSequence) {
      List<TableDataInsertAllResponse.InsertErrors> errors = new ArrayList<>();
      for (long i : errorIndices) {
        TableDataInsertAllResponse.InsertErrors error =
            new TableDataInsertAllResponse.InsertErrors();
        error.setIndex(i);
      }
      TableDataInsertAllResponse response = new TableDataInsertAllResponse();
      response.setInsertErrors(errors);
      responses.add(response);
    }

    when(mockTabledata.insertAll(
            anyString(), anyString(), anyString(), any(TableDataInsertAllRequest.class)))
        .thenReturn(mockInsertAll);
    when(mockInsertAll.execute())
        .thenReturn(
            responses.get(0),
            responses
                .subList(1, responses.size())
                .toArray(new TableDataInsertAllResponse[responses.size() - 1]));
  }
  @Test
  public void testReadMultiPage() throws IOException {
    onTableGet(basicTableSchema());

    TableDataList page1 = rawDataList(rawRow("Row1", 1)).setPageToken("page2");
    TableDataList page2 = rawDataList(rawRow("Row2", 2)).setTotalRows(2L);

    when(mockClient.tabledata()).thenReturn(mockTabledata);
    when(mockTabledata.list(anyString(), anyString(), anyString())).thenReturn(mockTabledataList);
    when(mockTabledataList.execute()).thenReturn(page1).thenReturn(page2);

    try (BigQueryTableRowIterator iterator =
        BigQueryTableRowIterator.of(
            mockClient, BigQueryIO.parseTableSpec("project:dataset.table"))) {

      List<String> names = new LinkedList<>();
      Iterators.addAll(
          names,
          Iterators.transform(
              iterator,
              new Function<TableRow, String>() {
                @Override
                public String apply(TableRow input) {
                  return (String) input.get("name");
                }
              }));

      Assert.assertThat(names, Matchers.hasItems("Row1", "Row2"));

      verifyTableGet();
      verifyTabledataList();
      // The second call should have used a page token.
      verify(mockTabledataList).setPageToken("page2");
    }
  }
  /** Checks if a table is empty. */
  public boolean isEmpty() throws IOException {
    Bigquery.Tabledata.List list =
        client.tabledata().list(ref.getProjectId(), ref.getDatasetId(), ref.getTableId());
    list.setMaxResults(1L);
    TableDataList dataList = list.execute();

    return dataList.getRows() == null || dataList.getRows().isEmpty();
  }
  @Test
  public void testReadOpenFailure() throws IOException {
    thrown.expect(RuntimeException.class);

    when(mockClient.tables()).thenReturn(mockTables);
    when(mockTables.get(anyString(), anyString(), anyString())).thenReturn(mockTablesGet);
    when(mockTablesGet.execute()).thenThrow(new IOException("No such table"));

    try (BigQueryTableRowIterator iterator =
        BigQueryTableRowIterator.of(
            mockClient, BigQueryIO.parseTableSpec("project:dataset.table"))) {
      try {
        Assert.assertFalse(iterator.hasNext()); // throws.
      } finally {
        verifyTableGet();
      }
    }
  }
  /**
   * Tries to create the BigQuery table. If a table with the same name already exists in the
   * dataset, the table creation fails, and the function returns null. In such a case, the existing
   * table doesn't necessarily have the same schema as specified by the parameter.
   *
   * @param schema Schema of the new BigQuery table.
   * @return The newly created BigQuery table information, or null if the table with the same name
   *     already exists.
   * @throws IOException if other error than already existing table occurs.
   */
  @Nullable
  public Table tryCreateTable(TableSchema schema) throws IOException {
    LOG.info("Trying to create BigQuery table: {}", BigQueryIO.toTableSpec(ref));

    Table content = new Table();
    content.setTableReference(ref);
    content.setSchema(schema);

    try {
      return client.tables().insert(ref.getProjectId(), ref.getDatasetId(), content).execute();
    } catch (IOException e) {
      if (new ApiErrorExtractor().itemAlreadyExists(e)) {
        LOG.info("The BigQuery table already exists.");
        return null;
      }
      throw e;
    }
  }
 @VisibleForTesting
 QueryResponse queryWithRetries(
     Bigquery bigqueryClient, QueryRequest queryContent, Sleeper sleeper, BackOff backOff)
     throws IOException, InterruptedException {
   IOException lastException = null;
   do {
     try {
       return bigqueryClient.jobs().query(projectId, queryContent).execute();
     } catch (IOException e) {
       // ignore and retry
       LOG.warn("Ignore the error and retry the query.");
       lastException = e;
     }
   } while (BackOffUtils.next(sleeper, backOff));
   throw new IOException(
       String.format("Unable to get BigQuery response after retrying %d times", MAX_QUERY_RETRIES),
       lastException);
 }
 private void onTableList(TableDataList result) throws IOException {
   when(mockClient.tabledata()).thenReturn(mockTabledata);
   when(mockTabledata.list(anyString(), anyString(), anyString())).thenReturn(mockTabledataList);
   when(mockTabledataList.execute()).thenReturn(result);
 }
 private void onTableGet(Table table) throws IOException {
   when(mockClient.tables()).thenReturn(mockTables);
   when(mockTables.get(anyString(), anyString(), anyString())).thenReturn(mockTablesGet);
   when(mockTablesGet.execute()).thenReturn(table);
 }
  /** Insert all rows from the given list using specified insertIds if not null. */
  public void insertAll(List<TableRow> rowList, @Nullable List<String> insertIdList)
      throws IOException {
    if (insertIdList != null && rowList.size() != insertIdList.size()) {
      throw new AssertionError(
          "If insertIdList is not null it needs to have at least " + "as many elements as rowList");
    }

    AttemptBoundedExponentialBackOff backoff =
        new AttemptBoundedExponentialBackOff(
            MAX_INSERT_ATTEMPTS, INITIAL_INSERT_BACKOFF_INTERVAL_MS);

    final List<TableDataInsertAllResponse.InsertErrors> allErrors = new ArrayList<>();
    // These lists contain the rows to publish. Initially the contain the entire list. If there are
    // failures, they will contain only the failed rows to be retried.
    List<TableRow> rowsToPublish = rowList;
    List<String> idsToPublish = insertIdList;
    while (true) {
      final List<TableRow> retryRows = new ArrayList<>();
      final List<String> retryIds = (idsToPublish != null) ? new ArrayList<String>() : null;

      int strideIndex = 0;
      // Upload in batches.
      List<TableDataInsertAllRequest.Rows> rows = new LinkedList<>();
      int dataSize = 0;

      List<Future<?>> futures = new ArrayList<>();

      for (int i = 0; i < rowsToPublish.size(); ++i) {
        TableRow row = rowsToPublish.get(i);
        TableDataInsertAllRequest.Rows out = new TableDataInsertAllRequest.Rows();
        if (idsToPublish != null) {
          out.setInsertId(idsToPublish.get(i));
        }
        out.setJson(row.getUnknownKeys());
        rows.add(out);

        dataSize += row.toString().length();
        if (dataSize >= UPLOAD_BATCH_SIZE_BYTES
            || rows.size() >= maxRowsPerBatch
            || i == rowsToPublish.size() - 1) {
          TableDataInsertAllRequest content = new TableDataInsertAllRequest();
          content.setRows(rows);

          final Bigquery.Tabledata.InsertAll insert =
              client
                  .tabledata()
                  .insertAll(ref.getProjectId(), ref.getDatasetId(), ref.getTableId(), content);

          final int finalStrideIndex = strideIndex;
          final List<TableRow> finalRowsToPublish = rowsToPublish;
          final List<String> finalIdsToPublish = idsToPublish;

          futures.add(
              executor.submit(
                  new Runnable() {
                    @Override
                    public void run() {
                      try {
                        TableDataInsertAllResponse response = insert.execute();

                        List<TableDataInsertAllResponse.InsertErrors> errors =
                            response.getInsertErrors();
                        if (errors != null) {
                          synchronized (this) {
                            allErrors.addAll(errors);
                            for (TableDataInsertAllResponse.InsertErrors error : errors) {
                              if (error.getIndex() == null) {
                                throw new IOException("Insert failed: " + allErrors);
                              }

                              int errorIndex = error.getIndex().intValue() + finalStrideIndex;
                              retryRows.add(finalRowsToPublish.get(errorIndex));
                              if (retryIds != null) {
                                retryIds.add(finalIdsToPublish.get(errorIndex));
                              }
                            }
                          }
                        }
                      } catch (IOException e) {
                        throw new RuntimeException(e);
                      }
                    }
                  }));

          dataSize = 0;
          strideIndex = i + 1;
          rows = new LinkedList<>();
        }
      }

      try {
        for (Future<?> future : futures) {
          future.get();
        }
      } catch (InterruptedException e) {
      } catch (ExecutionException e) {
        Throwables.propagate(e.getCause());
      }

      if (!allErrors.isEmpty() && !backoff.atMaxAttempts()) {
        try {
          Thread.sleep(backoff.nextBackOffMillis());
        } catch (InterruptedException e) {
          // ignore.
        }
        LOG.info("Retrying failed inserts to BigQuery");
        rowsToPublish = retryRows;
        idsToPublish = retryIds;
        allErrors.clear();
      } else {
        break;
      }
    }
    if (!allErrors.isEmpty()) {
      throw new IOException("Insert failed: " + allErrors);
    }
  }