KAFKA-16455: Check partition exists before send reassignments to serv…

…er in ReassignPartitionsCommand (apache#15659) Currently, when executing kafka-reassign-partitions.sh with the --execute option, if a partition number specified in the JSON file does not exist, this check occurs only when submitting the reassignments to alterPartitionReassignments on the server-side. We can perform this check in advance before submitting the reassignments to the server side. Reviewers: Luke Chen <[email protected]>
chiacyu · Apr 8, 2024 · f895ab5 · f895ab5
1 parent 169ed60
commit f895ab5
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java b/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java
@@ -641,7 +641,7 @@ static Map<TopicPartition, List<Integer>> getReplicaAssignmentForTopics(Admin ad
      * @param adminClient     The AdminClient to use.
      * @param partitions      The partitions to get information about.
      * @return                A map from partitions to broker assignments.
-     *                        If any topic can't be found, an exception will be thrown.
+     *                        If any topic or partition can't be found, an exception will be thrown.
      */
     static Map<TopicPartition, List<Integer>> getReplicaAssignmentForPartitions(Admin adminClient,
                                                                                 Set<TopicPartition> partitions
@@ -654,6 +654,13 @@ static Map<TopicPartition, List<Integer>> getReplicaAssignmentForPartitions(Admi
                     res.put(tp, info.replicas().stream().map(Node::id).collect(Collectors.toList()));
             })
         );
+
+        if (!res.keySet().equals(partitions)) {
+            Set<TopicPartition> missingPartitions = new HashSet<>(partitions);
+            missingPartitions.removeAll(res.keySet());
+            throw new ExecutionException(new UnknownTopicOrPartitionException("Unable to find partition: " +
+                missingPartitions.stream().map(TopicPartition::toString).collect(Collectors.joining(", "))));
+        }
         return res;
     }
 

diff --git a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java
@@ -83,6 +83,7 @@
 import static org.apache.kafka.tools.reassign.ReassignPartitionsCommand.replicaMoveStatesToString;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -300,6 +301,13 @@ public void testGetReplicaAssignments() throws Exception {
 
             assertEquals(assignments,
                 getReplicaAssignmentForPartitions(adminClient, new HashSet<>(asList(new TopicPartition("foo", 0), new TopicPartition("bar", 0)))));
+
+            UnknownTopicOrPartitionException exception =
+                assertInstanceOf(UnknownTopicOrPartitionException.class,
+                    assertThrows(ExecutionException.class,
+                        () -> getReplicaAssignmentForPartitions(adminClient,
+                            new HashSet<>(asList(new TopicPartition("foo", 0), new TopicPartition("foo", 10))))).getCause());
+            assertEquals("Unable to find partition: foo-10", exception.getMessage());
         }
     }