com.google.cloud.bigquery.InsertAllResponse Java Examples
The following examples show how to use
com.google.cloud.bigquery.InsertAllResponse.
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: BqSink.java From beast with Apache License 2.0 | 6 votes |
@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 #2
Source File: TableSnippets.java From google-cloud-java with Apache License 2.0 | 6 votes |
/** 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; }
Example #3
Source File: BQResponseParser.java From beast with Apache License 2.0 | 6 votes |
/** * 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 #4
Source File: TableSnippets.java From google-cloud-java with Apache License 2.0 | 6 votes |
/** 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; }
Example #5
Source File: ITTableSnippets.java From google-cloud-java with Apache License 2.0 | 6 votes |
@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 #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: ITBigQuerySnippets.java From google-cloud-java with Apache License 2.0 | 5 votes |
@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)); }
Example #8
Source File: TestBigQueryTarget.java From datacollector with Apache License 2.0 | 5 votes |
@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); } }
Example #9
Source File: ITTableSnippets.java From google-cloud-java with Apache License 2.0 | 5 votes |
@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()); }
Example #10
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 #11
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 #12
Source File: BigQueryClient.java From beam with Apache License 2.0 | 5 votes |
private void handleBigQueryResponseExceptions(InsertAllResponse response) { if (response.hasErrors()) { throw new RuntimeException( format( "The following errors occurred while inserting to BigQuery: %s", response.getInsertErrors())); } }
Example #13
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 #14
Source File: BigQueryWriteTest.java From gcp-ingestion with Mozilla Public License 2.0 | 5 votes |
/** 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()); }
Example #15
Source File: BigQuery.java From gcp-ingestion with Mozilla Public License 2.0 | 5 votes |
@Override protected void checkResultFor(InsertAllResponse batchResult, int index) { Optional.ofNullable(batchResult.getErrorsFor(index)).filter(errors -> !errors.isEmpty()) .ifPresent(errors -> { throw new BigQueryErrors(errors); }); }
Example #16
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 #17
Source File: ITTableSnippets.java From google-cloud-java with Apache License 2.0 | 4 votes |
@Test public void testInsert() throws InterruptedException { InsertAllResponse response = tableSnippets.insert("row1", "row2"); assertFalse(response.hasErrors()); verifyTestRows(table); }
Example #18
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 #19
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 #20
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 #21
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 #22
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 #23
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 #24
Source File: BigQuery.java From gcp-ingestion with Mozilla Public License 2.0 | 4 votes |
@Override protected CompletableFuture<InsertAllResponse> close() { return CompletableFuture.completedFuture(bigQuery.insertAll(builder.build())); }