@Override public void deleteAll(List<K> keys) { try { deleteAllAsync(keys).get(); } catch (Exception e) { throw new SamzaException(e); } }
@Override public CompletableFuture<Void> putAllAsync(List<Entry<K, V>> records) { Preconditions.checkNotNull(writeFn, "null write function"); Preconditions.checkNotNull(records, "null records"); if (records.isEmpty()) { return CompletableFuture.completedFuture(null); } List<K> deleteKeys = records.stream() .filter(e -> e.getValue() == null).map(Entry::getKey).collect(Collectors.toList()); List<Entry<K, V>> putRecords = records.stream() .filter(e -> e.getValue() != null).collect(Collectors.toList()); CompletableFuture<Void> deleteFuture = deleteKeys.isEmpty() ? CompletableFuture.completedFuture(null) : deleteAllAsync(deleteKeys); // Return the combined future return CompletableFuture.allOf( deleteFuture, instrument(() -> asyncTable.putAllAsync(putRecords), metrics.numPutAlls, metrics.putAllNs)) .exceptionally(e -> { String strKeys = records.stream().map(r -> r.getKey().toString()).collect(Collectors.joining(",")); throw new SamzaException(String.format("Failed to put records with keys=" + strKeys), e); }); }
public void doTestDeleteAll(boolean sync, boolean error) throws Exception { TableWriteFunction<String, String> writeFn = mock(TableWriteFunction.class); RemoteTable<String, String> table = getTable("testDeleteAll-" + sync + error, mock(TableReadFunction.class), writeFn, false); CompletableFuture<Void> future; if (error) { future = new CompletableFuture(); future.completeExceptionally(new RuntimeException("Test exception")); } else { future = CompletableFuture.completedFuture(null); } // Sync is backed by async so needs to mock the async method doReturn(future).when(writeFn).deleteAllAsync(any()); List<String> keys = Arrays.asList("foo1", "foo2"); ArgumentCaptor<List> argCaptor = ArgumentCaptor.forClass(List.class); if (sync) { table.deleteAll(keys); } else { table.deleteAllAsync(keys).get(); } verify(writeFn, times(1)).deleteAllAsync(argCaptor.capture()); Assert.assertEquals(keys, argCaptor.getValue()); verify(table.writeRateLimiter, times(1)).throttle(anyCollection()); }