KAFKA-16246: Cleanups in ConsoleConsumer

tools/src/main/java/org/apache/kafka/tools/consumer/ConsoleConsumer.java

@@ -22,8 +22,6 @@
 import java.util.Iterator;
 import java.util.Map;
 import java.util.Optional;
-import java.util.OptionalInt;
-import java.util.OptionalLong;
 import java.util.concurrent.CountDownLatch;
 import java.util.regex.Pattern;
 import java.util.Collections;
@@ -68,11 +66,8 @@ public static void main(String[] args) throws Exception {
 
     public static void run(ConsoleConsumerOptions opts) {
         messageCount = 0;
-        long timeoutMs = opts.timeoutMs() >= 0 ? opts.timeoutMs() : Long.MAX_VALUE;
         Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(opts.consumerProps(), new ByteArrayDeserializer(), new ByteArrayDeserializer());
-        ConsumerWrapper consumerWrapper = opts.partitionArg().isPresent()
-            ? new ConsumerWrapper(Optional.of(opts.topicArg()), opts.partitionArg(), OptionalLong.of(opts.offsetArg()), Optional.empty(), consumer, timeoutMs)
-            : new ConsumerWrapper(Optional.of(opts.topicArg()), OptionalInt.empty(), OptionalLong.empty(), Optional.ofNullable(opts.includedTopicsArg()), consumer, timeoutMs);
+        ConsumerWrapper consumerWrapper = new ConsumerWrapper(opts, consumer);
 
         addShutdownHook(consumerWrapper, opts);
 
@@ -148,43 +143,26 @@ static boolean checkErr(PrintStream output) {
     }
 
     public static class ConsumerWrapper {
-        final Optional<String> topic;
-        final OptionalInt partitionId;
-        final OptionalLong offset;
-        final Optional<String> includedTopics;
-        final Consumer<byte[], byte[]> consumer;
-        final long timeoutMs;
         final Time time = Time.SYSTEM;
+        final long timeoutMs;
+        final Consumer<byte[], byte[]> consumer;
 
         Iterator<ConsumerRecord<byte[], byte[]>> recordIter = Collections.emptyIterator();
 
-        public ConsumerWrapper(Optional<String> topic,
-                               OptionalInt partitionId,
-                               OptionalLong offset,
-                               Optional<String> includedTopics,
-                               Consumer<byte[], byte[]> consumer,
-                               long timeoutMs) {
-            this.topic = topic;
-            this.partitionId = partitionId;
-            this.offset = offset;
-            this.includedTopics = includedTopics;
+        public ConsumerWrapper(ConsoleConsumerOptions opts, Consumer<byte[], byte[]> consumer) {
+            Optional<String> topic = Optional.ofNullable(opts.topicArg());
+            Optional<String> includedTopics = Optional.ofNullable(opts.includedTopicsArg());
             this.consumer = consumer;
-            this.timeoutMs = timeoutMs;
-
-            if (topic.isPresent() && partitionId.isPresent() && offset.isPresent() && !includedTopics.isPresent()) {
-                seek(topic.get(), partitionId.getAsInt(), offset.getAsLong());
-            } else if (topic.isPresent() && partitionId.isPresent() && !offset.isPresent() && !includedTopics.isPresent()) {
-                // default to latest if no offset is provided
-                seek(topic.get(), partitionId.getAsInt(), ListOffsetsRequest.LATEST_TIMESTAMP);
-            } else if (topic.isPresent() && !partitionId.isPresent() && !offset.isPresent() && !includedTopics.isPresent()) {
-                consumer.subscribe(Collections.singletonList(topic.get()));
-            } else if (!topic.isPresent() && !partitionId.isPresent() && !offset.isPresent() && includedTopics.isPresent()) {
-                consumer.subscribe(Pattern.compile(includedTopics.get()));
+            timeoutMs = opts.timeoutMs() >= 0 ? opts.timeoutMs() : Long.MAX_VALUE;
+
+            if (topic.isPresent()) {
+                if (opts.partitionArg().isPresent()) {
+                    seek(topic.get(), opts.partitionArg().getAsInt(), opts.offsetArg());
+                } else {
+                    consumer.subscribe(Collections.singletonList(topic.get()));
+                }
             } else {
-                throw new IllegalArgumentException("An invalid combination of arguments is provided. " +
-                        "Exactly one of 'topic' or 'include' must be provided. " +
-                        "If 'topic' is provided, an optional 'partition' may also be provided. " +
-                        "If 'partition' is provided, an optional 'offset' may also be provided, otherwise, consumption starts from the end of the partition.");
+                includedTopics.ifPresent(topics -> consumer.subscribe(Pattern.compile(topics)));
             }
         }
 

tools/src/test/java/org/apache/kafka/tools/consumer/ConsoleConsumerTest.java

@@ -24,21 +24,19 @@
 import org.apache.kafka.common.MessageFormatter;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.TimeoutException;
-import org.apache.kafka.common.requests.ListOffsetsRequest;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.server.util.MockTime;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
+import java.io.IOException;
 import java.io.PrintStream;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.Optional;
-import java.util.OptionalInt;
-import java.util.OptionalLong;
+import java.util.regex.Pattern;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -58,8 +56,7 @@ public void setup() {
     }
 
     @Test
-    public void shouldThrowTimeoutExceptionWhenTimeoutIsReached() {
-        String topic = "test";
+    public void shouldThrowTimeoutExceptionWhenTimeoutIsReached() throws IOException {
         final Time time = new MockTime();
         final int timeoutMs = 1000;
 
@@ -71,20 +68,22 @@ public void shouldThrowTimeoutExceptionWhenTimeoutIsReached() {
             return ConsumerRecords.EMPTY;
         });
 
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", "test",
+            "--timeout-ms", String.valueOf(timeoutMs)
+        };
+
         ConsoleConsumer.ConsumerWrapper consumer = new ConsoleConsumer.ConsumerWrapper(
-                Optional.of(topic),
-                OptionalInt.empty(),
-                OptionalLong.empty(),
-                Optional.empty(),
-                mockConsumer,
-                timeoutMs
+            new ConsoleConsumerOptions(args),
+            mockConsumer
         );
 
         assertThrows(TimeoutException.class, consumer::receive);
     }
 
     @Test
-    public void shouldResetUnConsumedOffsetsBeforeExit() {
+    public void shouldResetUnConsumedOffsetsBeforeExit() throws IOException {
         String topic = "test";
         int maxMessages = 123;
         int totalMessages = 700;
@@ -94,13 +93,16 @@ public void shouldResetUnConsumedOffsetsBeforeExit() {
         TopicPartition tp1 = new TopicPartition(topic, 0);
         TopicPartition tp2 = new TopicPartition(topic, 1);
 
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", topic,
+            "--timeout-ms", "1000"
+        };
+
         ConsoleConsumer.ConsumerWrapper consumer = new ConsoleConsumer.ConsumerWrapper(
-                Optional.of(topic),
-                OptionalInt.empty(),
-                OptionalLong.empty(),
-                Optional.empty(),
-                mockConsumer,
-                1000L);
+            new ConsoleConsumerOptions(args),
+            mockConsumer
+        );
 
         mockConsumer.rebalance(Arrays.asList(tp1, tp2));
         Map<TopicPartition, Long> offsets = new HashMap<>();
@@ -165,47 +167,75 @@ public void shouldStopWhenOutputCheckErrorFails() {
 
     @Test
     @SuppressWarnings("unchecked")
-    public void shouldSeekWhenOffsetIsSet() {
+    public void shouldSeekWhenOffsetIsSet() throws IOException {
         Consumer<byte[], byte[]> mockConsumer = mock(Consumer.class);
         TopicPartition tp0 = new TopicPartition("test", 0);
 
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", tp0.topic(),
+            "--partition", String.valueOf(tp0.partition()),
+            "--timeout-ms", "1000"
+        };
+
         ConsoleConsumer.ConsumerWrapper consumer = new ConsoleConsumer.ConsumerWrapper(
-                Optional.of(tp0.topic()),
-                OptionalInt.of(tp0.partition()),
-                OptionalLong.empty(),
-                Optional.empty(),
-                mockConsumer,
-                1000L);
+            new ConsoleConsumerOptions(args),
+            mockConsumer
+        );
 
         verify(mockConsumer).assign(eq(Collections.singletonList(tp0)));
         verify(mockConsumer).seekToEnd(eq(Collections.singletonList(tp0)));
         consumer.cleanup();
         reset(mockConsumer);
 
-        consumer = new ConsoleConsumer.ConsumerWrapper(
-                Optional.of(tp0.topic()),
-                OptionalInt.of(tp0.partition()),
-                OptionalLong.of(123L),
-                Optional.empty(),
-                mockConsumer,
-                1000L);
+        args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", tp0.topic(),
+            "--partition", String.valueOf(tp0.partition()),
+            "--offset", "123",
+            "--timeout-ms", "1000"
+        };
+
+        consumer = new ConsoleConsumer.ConsumerWrapper(new ConsoleConsumerOptions(args), mockConsumer);
 
         verify(mockConsumer).assign(eq(Collections.singletonList(tp0)));
         verify(mockConsumer).seek(eq(tp0), eq(123L));
         consumer.cleanup();
         reset(mockConsumer);
 
-        consumer = new ConsoleConsumer.ConsumerWrapper(
-                Optional.of(tp0.topic()),
-                OptionalInt.of(tp0.partition()),
-                OptionalLong.of(ListOffsetsRequest.EARLIEST_TIMESTAMP),
-                Optional.empty(),
-                mockConsumer,
-                1000L);
+        args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--topic", tp0.topic(),
+            "--partition", String.valueOf(tp0.partition()),
+            "--offset", "earliest",
+            "--timeout-ms", "1000"
+        };
+
+        consumer = new ConsoleConsumer.ConsumerWrapper(new ConsoleConsumerOptions(args), mockConsumer);
 
         verify(mockConsumer).assign(eq(Collections.singletonList(tp0)));
         verify(mockConsumer).seekToBeginning(eq(Collections.singletonList(tp0)));
         consumer.cleanup();
         reset(mockConsumer);
     }
+
+    @Test
+    @SuppressWarnings("unchecked")
+    public void shouldWorkWithoutTopicOption() throws IOException {
+        Consumer<byte[], byte[]> mockConsumer = mock(Consumer.class);
+
+        String[] args = new String[]{
+            "--bootstrap-server", "localhost:9092",
+            "--include", "includeTest*",
+            "--from-beginning"
+        };
+
+        ConsoleConsumer.ConsumerWrapper consumer = new ConsoleConsumer.ConsumerWrapper(
+            new ConsoleConsumerOptions(args),
+            mockConsumer
+        );
+
+        verify(mockConsumer).subscribe(any(Pattern.class));
+        consumer.cleanup();
+    }
 }