Java Code Examples for com.google.cloud.bigquery.InsertAllRequest#RowToInsert

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

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

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

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

InsertAllRequest.RowToInsert of(Record record);