com.google.cloud.bigquery.InsertAllRequest Java Examples

@Test
public void shouldPushMultipleMessagesToBigQuerySuccessfully() {
    Record user1 = new Record(offsetInfo, createUser("alice"));
    Record user2 = new Record(offsetInfo, createUser("bob"));
    Record user3 = new Record(offsetInfo, createUser("mary"));
    Records records = new Records(Arrays.asList(user1, user2, user3));
    InsertAllRequest request = builder
            .addRow(user1.getId(), user1.getColumns())
            .addRow(user2.getId(), user2.getColumns())
            .addRow(user3.getId(), user3.getColumns())
            .build();

    Status status = sink.push(records);

    verify(bigquery).insertAll(request);
    assertTrue(status.isSuccess());
}
 

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;
}
 

@Before
public void setUp() {
    tableId = TableId.of("test-dataset", "test-table");
    builder = InsertAllRequest.newBuilder(tableId);
    bqRow = new BQRowWithInsertId();
    BQErrorHandler errorHandlerInstance = new OOBErrorHandler(new DefaultLogWriter());
    sink = new BqSink(bigquery, tableId, new BQResponseParser(), errorHandlerInstance, bqRow);
    when(successfulResponse.hasErrors()).thenReturn(false);
    when(bigquery.insertAll(any())).thenReturn(successfulResponse);
    when(failureResponse.hasErrors()).thenReturn(true);
    insertErrors = new HashMap<>();
    List<BigQueryError> columnError = Arrays.asList(new BigQueryError("failed since type mismatched", "column location", "message"));
    insertErrors.put(0L, columnError);
    when(failureResponse.getInsertErrors()).thenReturn(insertErrors);
}
 

@Test
public void shouldPushMessageToBigQuerySuccessfully() {
    Record user = new Record(offsetInfo, createUser("alice"));
    InsertAllRequest request = builder.addRow(user.getId(), user.getColumns()).build();
    Records records = new Records(Arrays.asList(user));

    Status status = sink.push(records);

    verify(bigquery).insertAll(request);
    assertTrue(status.isSuccess());
}
 

@Test
public void shouldErrorWhenBigQueryInsertFails() {
    Record user1 = new Record(offsetInfo, createUser("alice"));
    InsertAllRequest request = builder.addRow(user1.getId(), user1.getColumns()).build();
    Records records = new Records(Arrays.asList(user1));
    when(bigquery.insertAll(request)).thenReturn(failureResponse);

    Status status = sink.push(records);

    verify(bigquery).insertAll(request);
    assertFalse("Should return false", status.isSuccess());
    assertTrue(status.getException().isPresent());
    assertEquals(insertErrors, ((BqInsertErrors) status.getException().get()).getErrors());
}
 

/** 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;
}
 

private Batch(TableId tableId) {
  super();
  this.tableId = tableId;
  builder = InsertAllRequest.newBuilder(tableId)
      // ignore row values for columns not present in the table
      .setIgnoreUnknownValues(true)
      // insert all valid rows when invalid rows are present in the request
      .setSkipInvalidRows(true);
}
 

@Test
public void canHandleDocumentId() {
  output = new BigQuery.Write(bigQuery, MAX_BYTES, MAX_MESSAGES, NO_DELAY, BATCH_KEY_TEMPLATE,
      ForkJoinPool.commonPool(), PubsubMessageToObjectNode.Raw.of());
  output.apply(PubsubMessage.newBuilder().putAttributes("document_id", "id").build()).join();
  List<InsertAllRequest.RowToInsert> rows = ((BigQuery.Write.Batch) output.batches
      .get(BATCH_KEY)).builder.build().getRows();

  assertNull(rows.get(0).getId());
  assertEquals(ImmutableMap.of("document_id", "id"), rows.get(0).getContent());
  assertEquals(20, output.batches.get(BATCH_KEY).byteSize);
}
 

/** 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);
}
 

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));
}
 

@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);
}
 

@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);
    }
}
 

@Override
public InsertAllRequest.RowToInsert of(Record record) {
    return InsertAllRequest.RowToInsert.of(record.getId(), record.getColumns());
}
 

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));
}
 

@Test
public void shouldParseAndPushMessagesToBq() throws Exception {
    TableId tableId = TableId.of("bqsinktest", "test_messages");
    BqSink bqSink = new BqSink(bigQueryMock, tableId, new BQResponseParser(), gcsSinkHandler, bqRow);
    String orderNumber = "order-1";
    String orderUrl = "order_url";
    String orderDetails = "order_details";
    Instant now = Instant.now();
    long second = now.getEpochSecond();
    int nano = now.getNano();
    ColumnMapping mapping = new ColumnMapping();
    mapping.put("1", "order_number");
    mapping.put("2", "order_url");
    mapping.put("3", "order_details");
    mapping.put("4", "created_at");
    mapping.put("5", "status");
    mapping.put("6", "discounted_value");
    mapping.put("7", "success");
    mapping.put("8", "order_price");
    mapping.put("12", "aliases");

    ColumnMapping currStateMapping = new ColumnMapping();
    currStateMapping.put("record_name", "current_state");
    currStateMapping.put("1", "key");
    currStateMapping.put("2", "value");
    mapping.put("9", currStateMapping);

    converter = new ConsumerRecordConverter(new RowMapper(mapping), new ProtoParser(StencilClientFactory.getClient(), TestMessage.class.getName()), clock);
    Timestamp createdAt = Timestamp.newBuilder().setSeconds(second).setNanos(nano).build();
    TestKey key = TestKey.newBuilder().setOrderNumber(orderNumber).setOrderUrl(orderUrl).build();
    com.gojek.beast.Status completed = com.gojek.beast.Status.COMPLETED;
    long discount = 1234;
    float price = 1234.5678f;
    TestMessage message = TestMessage.newBuilder()
            .setOrderNumber(orderNumber)
            .setOrderUrl(orderUrl)
            .setOrderDetails(orderDetails)
            .setCreatedAt(createdAt)
            .setStatus(completed)
            .setDiscount(discount)
            .setPrice(price)
            .setSuccess(true)
            .addAliases("alias1").addAliases("alias2")
            .putCurrentState("state_key_1", "state_value_1")
            .putCurrentState("state_key_2", "state_value_2")
            .build();
    String topic = "topic";
    int partition = 1, offset = 1;
    long recordTimestamp = Instant.now().toEpochMilli();
    ConsumerRecord<byte[], byte[]> consumerRecord = new ConsumerRecord<>(topic, partition, offset, recordTimestamp, TimestampType.CREATE_TIME,
            0, 0, 1, key.toByteArray(), message.toByteArray());

    List<ConsumerRecord<byte[], byte[]>> messages = Arrays.asList(consumerRecord);
    when(successfulResponse.hasErrors()).thenReturn(false);
    when(bigQueryMock.insertAll(insertRequestCaptor.capture())).thenReturn(successfulResponse);

    List<Record> records = converter.convert(messages);
    Status status = bqSink.push(new Records(records));
    assertTrue(status.isSuccess());

    List<InsertAllRequest.RowToInsert> bqRows = insertRequestCaptor.getValue().getRows();
    assertEquals(1, bqRows.size());
    Map<String, Object> contents = bqRows.get(0).getContent();
    assertEquals("should have same number of columns as mappings, with metadata columns", mapping.size() + 5, contents.size());
    assertEquals(orderUrl, contents.get("order_url"));
    assertEquals(orderNumber, contents.get("order_number"));
    assertEquals(orderDetails, contents.get("order_details"));
    assertEquals(new DateTime(Instant.ofEpochSecond(second, nano).toEpochMilli()), contents.get("created_at"));
    assertEquals(completed.toString(), contents.get("status"));
    assertEquals(discount, contents.get("discounted_value"));
    assertEquals(price, contents.get("order_price"));
    assertEquals(Arrays.asList("alias1", "alias2"), contents.get("aliases"));
    assertTrue(Boolean.valueOf(contents.get("success").toString()));
    containsMetadata(contents, new OffsetInfo(topic, partition, offset, recordTimestamp));
    List repeatedStateMap = (List) contents.get("current_state");
    assertEquals("state_key_1", ((Map) repeatedStateMap.get(0)).get("key"));
    assertEquals("state_value_1", ((Map) repeatedStateMap.get(0)).get("value"));
    assertEquals("state_key_2", ((Map) repeatedStateMap.get(1)).get("key"));
    assertEquals("state_value_2", ((Map) repeatedStateMap.get(1)).get("value"));
}
 

@Override
public InsertAllRequest.RowToInsert of(Record record) {
    return InsertAllRequest.RowToInsert.of(record.getColumns());
}
 

InsertAllRequest.RowToInsert of(Record record);