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 |
@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 |
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 |
@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 |
@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 |
@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 |
/** 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 |
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 |
@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 |
/** 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 |
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 |
@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 |
@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 |
@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 |
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 |
@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 |
@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 |
@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 |
@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 |
@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 |
@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 |
InsertAllRequest.RowToInsert of(Record record);