com.google.cloud.bigquery.InsertAllResponse Java Examples

@Override
public Status push(Records records) {
    InsertAllResponse response = insertIntoBQ(records);
    InsertStatus status = new InsertStatus(!response.hasErrors(), response.getInsertErrors());
    //if bq has errors
    if (response.hasErrors()) {
        //parse the error records
        Map<Record, List<BQInsertionRecordsErrorType>> parsedRecords = responseParser.parseBQResponse(records, response);
        //sink error records
        boolean isTheSinkSuccessful = errorHandler.handleErrorRecords(parsedRecords).isSuccess();
        if (!isTheSinkSuccessful) {
            return new InsertStatus(isTheSinkSuccessful, response.getInsertErrors());
        }
        Records bqValidRecords = errorHandler.getBQValidRecords(parsedRecords);
        if (!bqValidRecords.getRecords().isEmpty()) { // there are valid records
            //insert the valid records into bq
            isTheSinkSuccessful &= !insertIntoBQ(bqValidRecords).hasErrors();
        }
        return new InsertStatus(isTheSinkSuccessful, response.getInsertErrors());
    }
    // return the original response
    return new InsertStatus(!response.hasErrors(), response.getInsertErrors());
}
 

/** Example of inserting rows into the table, ignoring invalid rows. */
// [TARGET insert(Iterable, boolean, boolean)]
// [VARIABLE "rowId1"]
// [VARIABLE "rowId2"]
public InsertAllResponse insertWithParams(String rowId1, String rowId2) {
  // [START ]
  List<RowToInsert> rows = new ArrayList<>();
  Map<String, Object> row1 = new HashMap<>();
  row1.put("stringField", 1);
  row1.put("booleanField", true);
  Map<String, Object> row2 = new HashMap<>();
  row2.put("stringField", "value2");
  row2.put("booleanField", false);
  rows.add(RowToInsert.of(rowId1, row1));
  rows.add(RowToInsert.of(rowId2, row2));
  InsertAllResponse response = table.insert(rows, true, true);
  // do something with response
  // [END ]
  return response;
}
 

/**
 * Parses the {@link InsertAllResponse} object and constructs a mapping of each record in @{@link Records} that were
 * tried to sink in BQ and the error type {@link BQInsertionRecordsErrorType}.
 *
 * @param records - records that were tried with BQ insertion
 * @param bqResponse - the status of insertion for all records as returned by BQ
 * @return - map of each record and the associated list of error types.
 */
public Map<Record, List<BQInsertionRecordsErrorType>> parseBQResponse(final Records records, final InsertAllResponse bqResponse) {
    if (!bqResponse.hasErrors()) {
        return Collections.emptyMap();
    }
    Map<Record, List<BQInsertionRecordsErrorType>> parsedRecords = new HashMap<>();
    Map<Long, List<BigQueryError>> insertErrorsMap = bqResponse.getInsertErrors();
    for (long recordIndex : insertErrorsMap.keySet()) {
        final Record message = records.getRecords().get((int) recordIndex);
        final List<BQInsertionRecordsErrorType> errorTypeList = new ArrayList<>();
        parsedRecords.put(message, errorTypeList);
        for (final BigQueryError err : insertErrorsMap.get(recordIndex)) {
            errorTypeList.add(ErrorTypeFactory.getErrorType(err.getReason(), err.getMessage()));
        } //end of for each row
    }
    return parsedRecords;
}
 

/** Example of inserting rows into the table. */
// [TARGET insert(Iterable)]
// [VARIABLE "rowId1"]
// [VARIABLE "rowId2"]
public InsertAllResponse insert(String rowId1, String rowId2) {
  // [START ]
  List<RowToInsert> rows = new ArrayList<>();
  Map<String, Object> row1 = new HashMap<>();
  row1.put("stringField", "value1");
  row1.put("booleanField", true);
  Map<String, Object> row2 = new HashMap<>();
  row2.put("stringField", "value2");
  row2.put("booleanField", false);
  rows.add(RowToInsert.of(rowId1, row1));
  rows.add(RowToInsert.of(rowId2, row2));
  InsertAllResponse response = table.insert(rows);
  // do something with response
  // [END ]
  return response;
}
 

@Test
public void testInsertParams() throws InterruptedException {
  InsertAllResponse response = tableSnippets.insertWithParams("row1", "row2");
  assertFalse(response.hasErrors());
  List<FieldValueList> rows = ImmutableList.copyOf(tableSnippets.list().getValues());
  while (rows.isEmpty()) {
    Thread.sleep(500);
    rows = ImmutableList.copyOf(tableSnippets.list().getValues());
  }
  Set<List<?>> values =
      FluentIterable.from(rows)
          .transform(
              new Function<FieldValueList, List<?>>() {
                @Override
                public List<?> apply(FieldValueList row) {
                  return ImmutableList.of(
                      row.get(0).getStringValue(), row.get(1).getBooleanValue());
                }
              })
          .toSet();
  assertEquals(ImmutableSet.of(ROW2), values);
}
 

/** Example of inserting rows into a table without running a load job. */
// [TARGET insertAll(InsertAllRequest)]
// [VARIABLE "my_dataset_name"]
// [VARIABLE "my_table_name"]
public InsertAllResponse insertAll(String datasetName, String tableName) {
  // [START bigquery_table_insert_rows]
  TableId tableId = TableId.of(datasetName, tableName);
  // Values of the row to insert
  Map<String, Object> rowContent = new HashMap<>();
  rowContent.put("booleanField", true);
  // Bytes are passed in base64
  rowContent.put("bytesField", "Cg0NDg0="); // 0xA, 0xD, 0xD, 0xE, 0xD in base64
  // Records are passed as a map
  Map<String, Object> recordsContent = new HashMap<>();
  recordsContent.put("stringField", "Hello, World!");
  rowContent.put("recordField", recordsContent);
  InsertAllResponse response =
      bigquery.insertAll(
          InsertAllRequest.newBuilder(tableId)
              .addRow(rowContent)
              // More rows can be added in the same RPC by invoking .addRow() on the builder.
              // You can also supply optional unique row keys to support de-duplication scenarios.
              .build());
  if (response.hasErrors()) {
    // If any of the insertions failed, this lets you inspect the errors
    for (Entry<Long, List<BigQueryError>> entry : response.getInsertErrors().entrySet()) {
      // inspect row error
    }
  }
  // [END bigquery_table_insert_rows]
  return response;
}
 

@Test
public void testInsertAllAndListTableData() throws IOException, InterruptedException {
  String tableName = "test_insert_all_and_list_table_data";
  String fieldName1 = "booleanField";
  String fieldName2 = "bytesField";
  String fieldName3 = "recordField";
  String fieldName4 = "stringField";
  TableId tableId = TableId.of(DATASET, tableName);
  Schema schema =
      Schema.of(
          Field.of(fieldName1, LegacySQLTypeName.BOOLEAN),
          Field.of(fieldName2, LegacySQLTypeName.BYTES),
          Field.of(
              fieldName3,
              LegacySQLTypeName.RECORD,
              Field.of(fieldName4, LegacySQLTypeName.STRING)));
  TableInfo table = TableInfo.of(tableId, StandardTableDefinition.of(schema));
  assertNotNull(bigquery.create(table));
  InsertAllResponse response = bigquerySnippets.insertAll(DATASET, tableName);
  assertFalse(response.hasErrors());
  assertTrue(response.getInsertErrors().isEmpty());
  Page<FieldValueList> listPage = bigquerySnippets.listTableDataFromId(DATASET, tableName);
  while (Iterators.size(listPage.iterateAll().iterator()) < 1) {
    Thread.sleep(500);
    listPage = bigquerySnippets.listTableDataFromId(DATASET, tableName);
  }
  FieldValueList row = listPage.getValues().iterator().next();
  assertEquals(true, row.get(0).getBooleanValue());
  assertArrayEquals(new byte[] {0xA, 0xD, 0xD, 0xE, 0xD}, row.get(1).getBytesValue());
  assertEquals("Hello, World!", row.get(2).getRecordValue().get(0).getStringValue());

  listPage = bigquerySnippets.listTableDataSchema(DATASET, tableName, schema, fieldName1);
  row = listPage.getValues().iterator().next();
  assertNotNull(row.get(fieldName1));
  assertArrayEquals(new byte[] {0xA, 0xD, 0xD, 0xE, 0xD}, row.get(fieldName2).getBytesValue());

  bigquerySnippets.listTableDataSchemaId();

  assertTrue(bigquerySnippets.deleteTable(DATASET, tableName));
}
 

@Test
public void testDataSetOrTableExistence() throws Exception {
  List<Record> records = new ArrayList<>();

  Record record1 = createRecord(ImmutableMap.of("a", 1));
  Record record2 = createRecord(ImmutableMap.of("a",  2));
  Record record3 = createRecord(ImmutableMap.of("a",  3));

  record1.getHeader().setAttribute("dataSet", "correctDataset");
  record1.getHeader().setAttribute("table", "wrongTable");

  record2.getHeader().setAttribute("dataSet", "wrongDataset");
  record2.getHeader().setAttribute("table", "correctTable");

  record3.getHeader().setAttribute("dataSet", "correctDataset");
  record3.getHeader().setAttribute("table", "correctTable");

  records.add(record1);
  records.add(record2);
  records.add(record3);

  mockBigQueryInsertAllRequest(invocationOnMock -> {
    InsertAllResponse response = PowerMockito.mock(InsertAllResponse.class);
    Mockito.doReturn(Collections.emptyMap()).when(response).getInsertErrors();
    Mockito.doReturn(false).when(response).hasErrors();
    return response;
  });

  BigQueryTargetConfigBuilder configBuilder = new BigQueryTargetConfigBuilder();
  configBuilder.datasetEL("${record:attribute('dataSet')}");
  configBuilder.tableNameEL("${record:attribute('table')}");
  TargetRunner targetRunner = createAndRunner(configBuilder.build(), records);

  Assert.assertEquals(2, targetRunner.getErrorRecords().size());
  for (Record errorRecord : targetRunner.getErrorRecords()) {
    String errorCode = errorRecord.getHeader().getErrorCode();
    Assert.assertEquals(Errors.BIGQUERY_17.getCode(), errorCode);
  }
}
 

@Test
public void testList() throws InterruptedException {
  List<FieldValueList> rows = ImmutableList.copyOf(tableSnippets.list().getValues());
  assertEquals(0, rows.size());

  InsertAllResponse response = tableSnippets.insert("row1", "row2");
  assertFalse(response.hasErrors());
  rows = ImmutableList.copyOf(tableSnippets.list().getValues());
  while (rows.isEmpty()) {
    Thread.sleep(500);
    rows = ImmutableList.copyOf(tableSnippets.list().getValues());
  }
  assertEquals(2, rows.size());
}
 

@Test
public void testCorrectImplicitMapping() throws Exception {
  List<Record> records = new ArrayList<>();

  records.add(createRecord(ImmutableMap.of("a", "1", "b", 1, "c", 1.0)));
  records.add(createRecord(ImmutableMap.of("a", "2", "b", 2, "c", 2.0)));
  records.add(createRecord(ImmutableMap.of("a", "3", "b", 3, "c", 3.0)));

  mockBigQueryInsertAllRequest(invocationOnMock -> {
    InsertAllRequest insertAllRequest = (InsertAllRequest)invocationOnMock.getArguments()[0];

    List<InsertAllRequest.RowToInsert> rows = insertAllRequest.getRows();

    final AtomicInteger recordIdx = new AtomicInteger(0);
    rows.forEach(row -> {
      int idx = recordIdx.getAndIncrement();
      Record record = records.get(idx);
      Map<String, ?> rowContent = row.getContent();
      record.get().getValueAsMap().forEach((k, v) -> Assert.assertEquals(v.getValue(), rowContent.get(k)));
    });

    InsertAllResponse response = PowerMockito.mock(InsertAllResponse.class);
    Mockito.doReturn(Collections.emptyMap()).when(response).getInsertErrors();
    Mockito.doReturn(false).when(response).hasErrors();
    return response;
  });

  BigQueryTargetConfigBuilder configBuilder = new BigQueryTargetConfigBuilder();

  createAndRunner(configBuilder.build(), records);
}
 

private void mockBigQueryInsertAllRequest(Answer<InsertAllResponse> insertAllResponseAnswer) {
  PowerMockito.doAnswer(insertAllResponseAnswer).when(bigQuery).insertAll(Mockito.any(InsertAllRequest.class));
  PowerMockito.doAnswer((Answer<Table>) invocationOnMock -> {
    TableId tableId = (TableId) invocationOnMock.getArguments()[0];
    return tableId.equals(TableId.of("correctDataset", "correctTable")) ?
        Mockito.mock(Table.class) : null;
  }).when(bigQuery).getTable(Mockito.any(TableId.class));
}
 

private void handleBigQueryResponseExceptions(InsertAllResponse response) {
  if (response.hasErrors()) {
    throw new RuntimeException(
        format(
            "The following errors occurred while inserting to BigQuery: %s",
            response.getInsertErrors()));
  }
}
 

/** Inserts multiple rows of the same schema to a BigQuery table. */
public void insertAll(Collection<Map<String, ?>> rows, String table) {
  TableId tableId = TableId.of(projectId, dataset, table);

  InsertAllRequest.Builder builder = InsertAllRequest.newBuilder(tableId);

  for (Map<String, ?> row : rows) {
    builder.addRow(row);
  }

  InsertAllResponse response = client.insertAll(builder.build());
  handleBigQueryResponseExceptions(response);
}
 

/** Prepare a mock BQ response. */
@Before
public void mockBigQueryResponse() {
  bigQuery = mock(com.google.cloud.bigquery.BigQuery.class);
  response = mock(InsertAllResponse.class);
  when(bigQuery.insertAll(any())).thenReturn(response);
  when(response.getErrorsFor(anyLong())).thenReturn(ImmutableList.of());
  output = new BigQuery.Write(bigQuery, MAX_BYTES, MAX_MESSAGES, MAX_DELAY, BATCH_KEY_TEMPLATE,
      ForkJoinPool.commonPool(), PubsubMessageToObjectNode.Raw.of());
}
 

@Override
protected void checkResultFor(InsertAllResponse batchResult, int index) {
  Optional.ofNullable(batchResult.getErrorsFor(index)).filter(errors -> !errors.isEmpty())
      .ifPresent(errors -> {
        throw new BigQueryErrors(errors);
      });
}
 

private InsertAllResponse insertIntoBQ(Records records) {
    Instant start = Instant.now();
    Builder builder = newBuilder(tableId);
    records.forEach((Record m) -> builder.addRow(recordInserter.of(m)));
    InsertAllRequest rows = builder.build();
    InsertAllResponse response = bigquery.insertAll(rows);
    log.info("Pushing {} records to BQ success?: {}", records.size(), !response.hasErrors());
    statsClient.count("bq.sink.push.records", records.size());
    statsClient.timeIt("bq.sink.push.time", start);
    return response;
}
 

@Test
public void testInsert() throws InterruptedException {
  InsertAllResponse response = tableSnippets.insert("row1", "row2");
  assertFalse(response.hasErrors());
  verifyTestRows(table);
}
 

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    final String projectId = context.getProperty(PROJECT_ID).evaluateAttributeExpressions().getValue();
    final String dataset = context.getProperty(DATASET).evaluateAttributeExpressions(flowFile).getValue();
    final String tableName = context.getProperty(TABLE_NAME).evaluateAttributeExpressions(flowFile).getValue();

    final TableId tableId;
    if (StringUtils.isEmpty(projectId)) {
        tableId = TableId.of(dataset, tableName);
    } else {
        tableId = TableId.of(projectId, dataset, tableName);
    }

    try {

        InsertAllRequest.Builder request = InsertAllRequest.newBuilder(tableId);
        int nbrecord = 0;

        try (final InputStream in = session.read(flowFile)) {
            final RecordReaderFactory readerFactory = context.getProperty(RECORD_READER).asControllerService(RecordReaderFactory.class);
            try (final RecordReader reader = readerFactory.createRecordReader(flowFile, in, getLogger());) {
                Record currentRecord;
                while ((currentRecord = reader.nextRecord()) != null) {
                    request.addRow(convertMapRecord(currentRecord.toMap()));
                    nbrecord++;
                }
            }
        }

        request.setIgnoreUnknownValues(context.getProperty(IGNORE_UNKNOWN).evaluateAttributeExpressions(flowFile).asBoolean());
        request.setSkipInvalidRows(context.getProperty(SKIP_INVALID_ROWS).evaluateAttributeExpressions(flowFile).asBoolean());

        InsertAllResponse response = getCloudService().insertAll(request.build());

        final Map<String, String> attributes = new HashMap<>();

        if (response.hasErrors()) {
            getLogger().log(LogLevel.WARN, "Failed to insert {} of {} records into BigQuery {} table.", new Object[] { response.getInsertErrors().size(), nbrecord, tableName });
            if (getLogger().isDebugEnabled()) {
                for (long index : response.getInsertErrors().keySet()) {
                    for (BigQueryError e : response.getInsertErrors().get(index)) {
                        getLogger().log(LogLevel.DEBUG, "Failed to insert record #{}: {}", new Object[] { index, e.getMessage() });
                    }
                }
            }

            attributes.put(BigQueryAttributes.JOB_NB_RECORDS_ATTR, Long.toString(nbrecord - response.getInsertErrors().size()));

            flowFile = session.penalize(flowFile);
            flowFile = session.putAllAttributes(flowFile, attributes);
            session.transfer(flowFile, REL_FAILURE);
        } else {
            attributes.put(BigQueryAttributes.JOB_NB_RECORDS_ATTR, Long.toString(nbrecord));
            flowFile = session.putAllAttributes(flowFile, attributes);
            session.transfer(flowFile, REL_SUCCESS);
        }

    } catch (Exception ex) {
        getLogger().log(LogLevel.ERROR, ex.getMessage(), ex);
        flowFile = session.penalize(flowFile);
        session.transfer(flowFile, REL_FAILURE);
    }
}
 

public static void main(String... args) throws InterruptedException {
  // Create a service instance
  BigQuery bigquery = BigQueryOptions.getDefaultInstance().getService();

  // Create a dataset
  String datasetId = "my_dataset_id";
  bigquery.create(DatasetInfo.newBuilder(datasetId).build());

  TableId tableId = TableId.of(datasetId, "my_table_id");
  // Table field definition
  Field stringField = Field.of("StringField", LegacySQLTypeName.STRING);
  // Table schema definition
  Schema schema = Schema.of(stringField);
  // Create a table
  StandardTableDefinition tableDefinition = StandardTableDefinition.of(schema);
  bigquery.create(TableInfo.of(tableId, tableDefinition));

  // Define rows to insert
  Map<String, Object> firstRow = new HashMap<>();
  Map<String, Object> secondRow = new HashMap<>();
  firstRow.put("StringField", "value1");
  secondRow.put("StringField", "value2");
  // Create an insert request
  InsertAllRequest insertRequest =
      InsertAllRequest.newBuilder(tableId).addRow(firstRow).addRow(secondRow).build();
  // Insert rows
  InsertAllResponse insertResponse = bigquery.insertAll(insertRequest);
  // Check if errors occurred
  if (insertResponse.hasErrors()) {
    System.out.println("Errors occurred while inserting rows");
  }

  // Create a query request
  QueryJobConfiguration queryConfig =
      QueryJobConfiguration.newBuilder("SELECT * FROM my_dataset_id.my_table_id").build();
  // Read rows
  System.out.println("Table rows:");
  for (FieldValueList row : bigquery.query(queryConfig).iterateAll()) {
    System.out.println(row);
  }
}
 

@Test
public void testErrorInIngestingMultipleTables() throws Exception {
  List<Record> records = new ArrayList<>();

  Record record1 = createRecord(ImmutableMap.of("a", 1));
  Record record2 = createRecord(ImmutableMap.of("a",  2));
  Record record3 = createRecord(ImmutableMap.of("a",  3));

  record1.getHeader().setAttribute("table", "table1");
  record2.getHeader().setAttribute("table", "table2");
  record3.getHeader().setAttribute("table", "table3");

  records.add(record1);
  records.add(record2);
  records.add(record3);

  Answer<InsertAllResponse> insertAllResponseAnswer = invocationOnMock -> {
    InsertAllRequest request = (InsertAllRequest) (invocationOnMock.getArguments()[0]);
    InsertAllResponse response = PowerMockito.mock(InsertAllResponse.class);
    if (request.getTable().getTable().equals("table2")) {
      BigQueryError bigQueryError = PowerMockito.mock(BigQueryError.class);
      Mockito.doReturn("Error in bigquery").when(bigQueryError).getMessage();
      Mockito.doReturn("Error in bigquery").when(bigQueryError).getReason();
      Mockito.doReturn(
          ImmutableMap.of(0L, Collections.singletonList(bigQueryError))
      ).when(response).getInsertErrors();
      Mockito.doReturn(true).when(response).hasErrors();
    } else {
      Mockito.doReturn(Collections.emptyMap()).when(response).getInsertErrors();
      Mockito.doReturn(false).when(response).hasErrors();
    }
    return response;
  };

  PowerMockito.doAnswer(insertAllResponseAnswer).when(bigQuery).insertAll(Mockito.any(InsertAllRequest.class));
  PowerMockito.doAnswer((Answer<Table>) invocationOnMock -> Mockito.mock(Table.class))
      .when(bigQuery).getTable(Mockito.any(TableId.class));

  BigQueryTargetConfigBuilder configBuilder = new BigQueryTargetConfigBuilder();
  configBuilder.datasetEL("sample");
  configBuilder.tableNameEL("${record:attribute('table')}");
  TargetRunner targetRunner = createAndRunner(configBuilder.build(), records);

  Assert.assertEquals(1, targetRunner.getErrorRecords().size());
  for (Record errorRecord : targetRunner.getErrorRecords()) {
    String errorCode = errorRecord.getHeader().getErrorCode();
    Assert.assertNotNull(errorCode);
    Assert.assertEquals("table2", errorRecord.getHeader().getAttribute("table"));
  }

}
 

@Test
public void testNestedAndRepeatedRecords() throws Exception {
  Map<String, Object> expectedRowContent = new LinkedHashMap<>();

  //one level map
  Map<String, Object> innerLevelMap = new LinkedHashMap<>();
  innerLevelMap.put("aMap1", 1);
  innerLevelMap.put("aMap2", 2.0);
  innerLevelMap.put("aMap3", true);

  expectedRowContent.put("a", innerLevelMap);

  //single list
  List<String> innerStringList = new ArrayList<>();
  innerStringList.add("bList1");
  innerStringList.add("bList2");
  innerStringList.add("bList3");

  expectedRowContent.put("b", innerStringList);

  //Nested Repeated map
  List<Map<String, Object>> innerLevelMapList = new ArrayList<>();

  Map<String, Object> innerLevelMap1 = new LinkedHashMap<>();
  innerLevelMap1.put("cList1a", "a");
  innerLevelMap1.put("cList1b", 1);
  innerLevelMap1.put("cList1c", 2.0);
  innerLevelMap1.put("cList1d", true);
  innerLevelMap1.put("cList1e", ImmutableMap.of("e", 1));


  Map<String, Object> innerLevelMap2 = new LinkedHashMap<>();
  innerLevelMap2.put("cList2a", "b");
  innerLevelMap2.put("cList2b", 2);
  innerLevelMap2.put("cList2c", 3.0);
  innerLevelMap2.put("cList2d", false);
  innerLevelMap1.put("cList2e", ImmutableMap.of("e", 2));


  Map<String, Object> innerLevelMap3 = new LinkedHashMap<>();
  innerLevelMap3.put("cList3a", "c");
  innerLevelMap3.put("cList3b", 3);
  innerLevelMap3.put("cList3c", 4.0);
  innerLevelMap3.put("cList3d", true);
  innerLevelMap1.put("cList3e", ImmutableMap.of("e", 3));

  innerLevelMapList.add(innerLevelMap1);
  innerLevelMapList.add(innerLevelMap2);
  innerLevelMapList.add(innerLevelMap3);

  expectedRowContent.put("c", innerLevelMapList);


  Record record = RecordCreator.create();
  record.set(createField(expectedRowContent));

  mockBigQueryInsertAllRequest(invocationOnMock -> {
    InsertAllRequest insertAllRequest = (InsertAllRequest)invocationOnMock.getArguments()[0];

    List<InsertAllRequest.RowToInsert> rows = insertAllRequest.getRows();

    Assert.assertEquals(1, rows.size());

    InsertAllRequest.RowToInsert row = rows.get(0);

    row.getContent().forEach((k, v) -> checkValues(expectedRowContent.get(k), v));

    InsertAllResponse response = PowerMockito.mock(InsertAllResponse.class);
    Mockito.doReturn(Collections.emptyMap()).when(response).getInsertErrors();
    Mockito.doReturn(false).when(response).hasErrors();
    return response;
  });

  BigQueryTargetConfigBuilder configBuilder = new BigQueryTargetConfigBuilder();
  createAndRunner(configBuilder.build(), Collections.singletonList(record));
}
 

@Test
public void testRowId() throws Exception {
  List<Record> records = new ArrayList<>();

  records.add(
      createRecord(
          ImmutableMap.of("a", 1, "b", 11, "c", 111)
      )
  );
  records.add(
      createRecord(
          ImmutableMap.of("a", 2, "b", 22, "c", 222)
      )
  );
  records.add(
      createRecord(
          ImmutableMap.of("a", 1, "b", 33, "c", 333)
      )
  );


  final Map<String, Map<String, Object>> rowIdToRow = new LinkedHashMap<>();

  mockBigQueryInsertAllRequest(invocationOnMock -> {
    InsertAllResponse response = PowerMockito.mock(InsertAllResponse.class);
    Mockito.doReturn(Collections.emptyMap()).when(response).getInsertErrors();
    Mockito.doReturn(false).when(response).hasErrors();

    InsertAllRequest request = (InsertAllRequest)invocationOnMock.getArguments()[0];

    request.getRows().forEach(row ->
        rowIdToRow.computeIfAbsent(row.getId(), rowId -> new LinkedHashMap<>()).putAll(row.getContent())
    );
    return response;
  });


  BigQueryTargetConfigBuilder configBuilder = new BigQueryTargetConfigBuilder();
  configBuilder.ignoreInvalidColumns(true);
  //Set value of a has row id
  configBuilder.rowIdExpression("${record:value('/a')}");
  createAndRunner(configBuilder.build(), records);

  Assert.assertEquals(2, rowIdToRow.size());

  rowIdToRow.forEach((rowId, row) ->{
    switch (rowId) {
      case "1":
        Assert.assertEquals(1, row.get("a"));
        Assert.assertEquals(33, row.get("b"));
        Assert.assertEquals(333, row.get("c"));
        break;
      case "2":
        Assert.assertEquals(2, row.get("a"));
        Assert.assertEquals(22, row.get("b"));
        Assert.assertEquals(222, row.get("c"));
        break;
      default:
        Assert.fail("Unexpected row id: " + rowId);
        break;
    }
  });
}
 

@Test
public void testDateTimeAndByteFields() throws Exception {
  Record record = RecordCreator.create();
  Map<String, Field> rootField = new LinkedHashMap<>();
  Map<String, Object> expectedContentMap = new LinkedHashMap<>();

  Date currentDate = new Date();

  String sampleBytes = "sample";

  rootField.put("dateField", Field.create(Field.Type.DATE, currentDate));
  expectedContentMap.put("dateField", DATE_FORMAT.format(currentDate));

  rootField.put("timeField", Field.create(Field.Type.TIME, currentDate));
  expectedContentMap.put("timeField", TIME_FORMAT.format(currentDate));

  rootField.put("datetimeField", Field.create(Field.Type.DATETIME, currentDate));
  expectedContentMap.put("datetimeField", DATE_TIME_FORMAT.format(currentDate));

  rootField.put("bytesField", Field.create(Field.Type.BYTE_ARRAY, sampleBytes.getBytes()));
  expectedContentMap.put("bytesField", Base64.getEncoder().encodeToString(sampleBytes.getBytes()));

  record.set(Field.create(rootField));

  mockBigQueryInsertAllRequest(invocationOnMock -> {
    InsertAllRequest request = (InsertAllRequest) invocationOnMock.getArguments()[0];
    InsertAllRequest.RowToInsert rowToInsert = request.getRows().get(0);
    Map<String,Object> actualContentMap =  rowToInsert.getContent();

    Assert.assertEquals(expectedContentMap.keySet(), actualContentMap.keySet());

    expectedContentMap.forEach((ek, ev) -> {
      Object actualContent = actualContentMap.get(ek);
      Assert.assertEquals(ev, actualContent);
    });

    InsertAllResponse response = PowerMockito.mock(InsertAllResponse.class);
    Mockito.doReturn(Collections.emptyMap()).when(response).getInsertErrors();
    Mockito.doReturn(false).when(response).hasErrors();
    return response;
  });

  BigQueryTargetConfigBuilder configBuilder = new BigQueryTargetConfigBuilder();
  configBuilder.ignoreInvalidColumns(true);
  createAndRunner(configBuilder.build(), Collections.singletonList(record));
}
 

@Override
protected CompletableFuture<InsertAllResponse> close() {
  return CompletableFuture.completedFuture(bigQuery.insertAll(builder.build()));
}