Examples with Context org.apache.beam.sdk.io.gcp.bigquery.InsertRetryPolicy.Context used on opensource projects

Example 1 with Context

use of org.apache.beam.sdk.io.gcp.bigquery.InsertRetryPolicy.Context in project beam by apache.

the class FakeDatasetService method insertAll.

@Override
public long insertAll(TableReference ref, List<ValueInSingleWindow<TableRow>> rowList, @Nullable List<String> insertIdList, InsertRetryPolicy retryPolicy, List<ValueInSingleWindow<TableRow>> failedInserts) throws IOException, InterruptedException {
    Map<TableRow, List<TableDataInsertAllResponse.InsertErrors>> insertErrors = getInsertErrors();
    synchronized (BigQueryIOTest.tables) {
        if (insertIdList != null) {
            assertEquals(rowList.size(), insertIdList.size());
        } else {
            insertIdList = Lists.newArrayListWithExpectedSize(rowList.size());
            for (int i = 0; i < rowList.size(); ++i) {
                insertIdList.add(Integer.toString(ThreadLocalRandom.current().nextInt()));
            }
        }
        long dataSize = 0;
        TableContainer tableContainer = getTableContainer(ref.getProjectId(), ref.getDatasetId(), ref.getTableId());
        for (int i = 0; i < rowList.size(); ++i) {
            TableRow row = rowList.get(i).getValue();
            List<TableDataInsertAllResponse.InsertErrors> allErrors = insertErrors.get(row);
            boolean shouldInsert = true;
            if (allErrors != null) {
                for (TableDataInsertAllResponse.InsertErrors errors : allErrors) {
                    if (!retryPolicy.shouldRetry(new Context(errors))) {
                        shouldInsert = false;
                    }
                }
            }
            if (shouldInsert) {
                dataSize += tableContainer.addRow(row, insertIdList.get(i));
            } else {
                failedInserts.add(rowList.get(i));
            }
        }
        return dataSize;
    }
}