com.google.cloud.bigquery.InsertAllRequest Java Examples

The following examples show how to use com.google.cloud.bigquery.InsertAllRequest. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example #1
Source File: BqSinkTest.java    From beast with Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example #2
Source File: BqSink.java    From beast with Apache License 2.0 5 votes vote down vote up
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;
}
 
Example #3
Source File: BqSinkTest.java    From beast with Apache License 2.0 5 votes vote down vote up
@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);
}
 
Example #4
Source File: BqSinkTest.java    From beast with Apache License 2.0 5 votes vote down vote up
@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());
}
 
Example #5
Source File: BqSinkTest.java    From beast with Apache License 2.0 5 votes vote down vote up
@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 #6
Source File: BigQuerySnippets.java    From google-cloud-java with Apache License 2.0 5 votes vote down vote up
/** 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;
}
 
Example #7
Source File: BigQuery.java    From gcp-ingestion with Mozilla Public License 2.0 5 votes vote down vote up
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);
}
 
Example #8
Source File: BigQueryWriteTest.java    From gcp-ingestion with Mozilla Public License 2.0 5 votes vote down vote up
@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);
}
 
Example #9
Source File: BigQueryClient.java    From beam with Apache License 2.0 5 votes vote down vote up
/** 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);
}
 
Example #10
Source File: TestBigQueryTarget.java    From datacollector with Apache License 2.0 5 votes vote down vote up
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));
}
 
Example #11
Source File: TestBigQueryTarget.java    From datacollector with Apache License 2.0 5 votes vote down vote up
@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);
}
 
Example #12
Source File: PutBigQueryStreaming.java    From nifi with Apache License 2.0 4 votes vote down vote up
@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);
    }
}
 
Example #13
Source File: BQRowWithInsertId.java    From beast with Apache License 2.0 4 votes vote down vote up
@Override
public InsertAllRequest.RowToInsert of(Record record) {
    return InsertAllRequest.RowToInsert.of(record.getId(), record.getColumns());
}
 
Example #14
Source File: InsertDataAndQueryTable.java    From google-cloud-java with Apache License 2.0 4 votes vote down vote up
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);
  }
}
 
Example #15
Source File: TestBigQueryTarget.java    From datacollector with Apache License 2.0 4 votes vote down vote up
@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"));
  }

}
 
Example #16
Source File: TestBigQueryTarget.java    From datacollector with Apache License 2.0 4 votes vote down vote up
@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));
}
 
Example #17
Source File: TestBigQueryTarget.java    From datacollector with Apache License 2.0 4 votes vote down vote up
@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;
    }
  });
}
 
Example #18
Source File: TestBigQueryTarget.java    From datacollector with Apache License 2.0 4 votes vote down vote up
@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));
}
 
Example #19
Source File: BqIntegrationTest.java    From beast with Apache License 2.0 4 votes vote down vote up
@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"));
}
 
Example #20
Source File: BQRowWithoutId.java    From beast with Apache License 2.0 4 votes vote down vote up
@Override
public InsertAllRequest.RowToInsert of(Record record) {
    return InsertAllRequest.RowToInsert.of(record.getColumns());
}
 
Example #21
Source File: BQRow.java    From beast with Apache License 2.0 votes vote down vote up
InsertAllRequest.RowToInsert of(Record record);