@Override public GetRecordsResult answer(InvocationOnMock invocation) throws Throwable { final String shardIterator = ((GetRecordsRequest) invocation.getArguments()[0]).getShardIterator(); // note that HashMap returns null when there is no entry in the map. // A null 'nextShardIterator' indicates that the shard has finished // and we should move onto the next shard. String nextShardIterator = shardIterators.get(shardIterator); Matcher m = Pattern.compile("shard_iterator_d_0*(\\d+)").matcher(shardIterator); Collection<Record> ans = answers.get(shardIterator); if (nextShardIterator == null && m.matches()) { // last shard iterates forever. Integer num = Integer.parseInt(m.group(1)); nextShardIterator = "shard_iterator_d_" + pad(Integer.toString(num + 1), 3); } if (null == ans) { // default to an empty list of records. ans = createRecords(); } return new GetRecordsResult() .withRecords(ans) .withNextShardIterator(nextShardIterator); }
@Override public List<SourceRecord> poll() throws InterruptedException { // TODO rate limiting? if (assignedShards.isEmpty()) { throw new ConnectException("No remaining source shards"); } final String shardId = assignedShards.get(currentShardIdx); final GetRecordsRequest req = new GetRecordsRequest(); req.setShardIterator(shardIterator(shardId)); req.setLimit(100); // TODO configurable final GetRecordsResult rsp = streamsClient.getRecords(req); if (rsp.getNextShardIterator() == null) { log.info("Shard ID `{}` for table `{}` has been closed, it will no longer be polled", shardId, config.tableForShard(shardId)); shardIterators.remove(shardId); assignedShards.remove(shardId); } else { log.debug("Retrieved {} records from shard ID `{}`", rsp.getRecords().size(), shardId); shardIterators.put(shardId, rsp.getNextShardIterator()); } currentShardIdx = (currentShardIdx + 1) % assignedShards.size(); final String tableName = config.tableForShard(shardId); final String topic = config.topicFormat.replace("${table}", tableName); final Map<String, String> sourcePartition = sourcePartition(shardId); return rsp.getRecords().stream() .map(dynamoRecord -> toSourceRecord(sourcePartition, topic, dynamoRecord.getDynamodb())) .collect(Collectors.toList()); }