com.epam.eco.commons.kafka.helpers.TopicOffsetForTimeFetcher Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of commons-kafka Show documentation
A library of utilities, helpers and higher-level APIs for the Kafka client library
There is a newer version: 3.0.5
/*
 * Copyright 2020 EPAM Systems
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License.  You may obtain a copy
 * of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package com.epam.eco.commons.kafka.helpers;

import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.apache.commons.lang3.Validate;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
import org.apache.kafka.common.TopicPartition;

import com.epam.eco.commons.kafka.KafkaUtils;
import com.epam.eco.commons.kafka.config.ConsumerConfigBuilder;

/**
 * @author Naira_Tamrazyan
 */
public class TopicOffsetForTimeFetcher {

    private final Map consumerConfig;

    private TopicOffsetForTimeFetcher(String bootstrapServers, Map consumerConfig) {
        this.consumerConfig = ConsumerConfigBuilder.
                with(consumerConfig).
                bootstrapServers(bootstrapServers).
                minRequiredConfigs().
                enableAutoCommitDisabled().
                clientIdRandom().
                build();
    }

    public static TopicOffsetForTimeFetcher with(Map consumerConfig) {
        return new TopicOffsetForTimeFetcher(null, consumerConfig);
    }

    public static TopicOffsetForTimeFetcher with(String bootstrapServers) {
        return new TopicOffsetForTimeFetcher(bootstrapServers, null);
    }

    public Map fetchForPartition(TopicPartition topicPartition, Long timestamp) {
        return fetchForPartitions(
                topicPartition != null ? Collections.singletonList(topicPartition) : null, timestamp);
    }

    public Map fetchForPartitions(Map topicPartitionTimes) {
        Validate.notEmpty(topicPartitionTimes, "Map of partition timestamps is null or empty");
        Validate.noNullElements(topicPartitionTimes.keySet(), "Collection of partitions null elements");
        Validate.noNullElements(topicPartitionTimes.values(), "Collection of timestamps contains null elements");

        try (KafkaConsumer consumer = new KafkaConsumer<>(consumerConfig)) {
            return doFetch(consumer, topicPartitionTimes);
        }
    }

    public Map fetchForPartitions(
            Collection topicPartitions,
            Long timestamp) {
        Validate.notEmpty(topicPartitions, "Collection of partitions is null or empty");
        Validate.noNullElements(topicPartitions, "Collection of partitions contains null elements");
        Validate.notNull(timestamp, "Timestamp is null");

        try (KafkaConsumer consumer = new KafkaConsumer<>(consumerConfig)) {
            Map partitionTimestamps =
                    toPartitionTimestamps(topicPartitions, timestamp);
            return doFetch(consumer, partitionTimestamps);
        }
    }

    public Map fetchForTopic(String topicName, Long timestamp) {
        return fetchForTopics(
                topicName != null ? Collections.singletonList(topicName) : null, timestamp);
    }

    public Map fetchForTopics(Collection topicNames, Long timestamp) {
        Validate.notEmpty(topicNames, "Collection of topic names is null or empty");
        Validate.noNullElements(topicNames, "Collection of topic names contains null elements");
        Validate.notNull(timestamp, "Timestamp is null");

        try (KafkaConsumer consumer = new KafkaConsumer<>(consumerConfig)) {
            List partitions = KafkaUtils.getTopicPartitionsAsList(consumer, topicNames);
            Map partitionTimestamps =
                    toPartitionTimestamps(partitions, timestamp);
            return doFetch(consumer, partitionTimestamps);
        }
    }

    private static Map toPartitionTimestamps(
            Collection partitions,
            Long timestamp) {
        return partitions.stream().
                collect(Collectors.toMap(Function.identity(), partition -> timestamp));
    }

    private static Map doFetch(Consumer consumer,
                                                     Map partitionTimestamps) {
        Map offsetsForTimes = consumer.offsetsForTimes(partitionTimestamps);

        return offsetsForTimes.entrySet().stream()
                .collect(Collectors.toMap(Map.Entry::getKey, entry ->
                        entry.getValue() == null ? null : entry.getValue().offset()));
    }

}