Admin: Cleanup alter_partition_reassignments (#3002)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: dpkp

kafka/admin/_partitions.py

@@ -268,45 +268,63 @@ def _process_alter_partition_reassignments_input(reassignments):
         _Partition = _Topic.ReassignablePartition
         topic2partitions = defaultdict(list)
         for tp, replicas in reassignments.items():
+            if replicas is not None:
+                replicas = list(replicas)
+                if not replicas:
+                    raise ValueError(
+                        "Replica list for %s must be non-empty; "
+                        "use None to cancel a reassignment." % (tp,))
+                elif not all(isinstance(item, int) for item in replicas):
+                    raise ValueError(
+                        "Replica list for %s must be int broker_ids." % (tp,))
             topic2partitions[tp.topic].append(_Partition(
                 partition_index=tp.partition,
-                replicas=list(replicas) if replicas is not None else None,
+                replicas=replicas,
             ))
         return [_Topic(name=topic, partitions=parts) for topic, parts in topic2partitions.items()]
 
-    def alter_partition_reassignments(self, reassignments, timeout_ms=None, raise_errors=True):
+    def alter_partition_reassignments(self, reassignments, timeout_ms=None):
         """Alter the replica sets for the given partitions.
 
         Arguments:
             reassignments (dict): A dict mapping
-                :class:`~kafka.TopicPartition` to a list of broker IDs for
-                the new replica set, or ``None`` to cancel a pending
-                reassignment for that partition.
+                :class:`~kafka.TopicPartition` to a list of broker IDs
+                for the new replica set, or ``None`` to cancel a
+                pending reassignment for that partition.
 
         Keyword Arguments:
-            timeout_ms (numeric, optional): The time in ms to wait for the
-                request to complete.
-            raise_errors (bool, optional): Whether to raise errors as
-                exceptions. Default True.
+            timeout_ms (numeric, optional): The time in ms to wait for
+                the request to complete.
+
+        Raises: top-level failures that prevents processing request.
+            Does not raise partition-specific errors.
 
         Returns:
-            Decoded AlterPartitionReassignmentsResponse (as a dict).
+            dict: A dict mapping each :class:`~kafka.TopicPartition`
+            that the broker acknowledged to the error class for that
+            partition, or ``None`` if the reassignment was accepted.
+            Partitions the broker did not report on are absent from the
+            dict.
         """
         timeout_ms = self._validate_timeout(timeout_ms)
 
-        def response_errors(r):
-            yield Errors.for_code(r.error_code)
-            for topic in r.responses:
-                for partition in topic.partitions:
-                    yield Errors.for_code(partition.error_code)
-
         request = AlterPartitionReassignmentsRequest(
             timeout_ms=timeout_ms,
             topics=self._process_alter_partition_reassignments_input(reassignments),
         )
+
+        def top_level_error(r):
+            yield Errors.for_code(r.error_code)
         response = self._manager.run(
-            self._send_request_to_controller, request, response_errors, raise_errors)
-        return response.to_dict()
+            self._send_request_to_controller, request, top_level_error)
+
+        results = {}
+        for topic in response.responses:
+            for partition in topic.partitions:
+                tp = TopicPartition(topic.name, partition.partition_index)
+                err = Errors.for_code(partition.error_code)
+                results[tp] = err if err is not Errors.NoError else None
+        return results
 
     async def _async_list_partition_reassignments(self, topic_partitions=None, timeout_ms=None):
         timeout_ms = self._validate_timeout(timeout_ms)
@@ -333,12 +351,10 @@ async def _async_list_partition_reassignments(self, topic_partitions=None, timeo
             timeout_ms=timeout_ms,
             topics=topics_field,
         )
-        response = await self._manager.send(request)
 
-        top_level_error = Errors.for_code(response.error_code)
-        if top_level_error is not Errors.NoError:
-            raise top_level_error(
-                "ListPartitionReassignmentsRequest failed: %s" % response.error_message)
+        def top_level_error(r):
+            yield Errors.for_code(r.error_code)
+        response = await self._send_request_to_controller(request, top_level_error)
 
         ret = {}
         for topic in response.topics:

kafka/cli/admin/partitions/alter_reassignments.py

@@ -16,9 +16,6 @@ def add_arguments(cls, parser):
         parser.add_argument(
             '--timeout-ms', type=int, default=None,
             help='Request timeout in milliseconds')
-        parser.add_argument(
-            '--no-raise-errors', dest='raise_errors', action='store_false',
-            help='Do not raise on partition-level errors; return the response instead')
 
     @classmethod
     def command(cls, client, args):
@@ -31,7 +28,10 @@ def command(cls, client, args):
                 reassignments[tp] = None
             else:
                 reassignments[tp] = [int(b) for b in replicas_str.split(',') if b]
-        return client.alter_partition_reassignments(
+        results = client.alter_partition_reassignments(
             reassignments,
-            timeout_ms=args.timeout_ms,
-            raise_errors=args.raise_errors)
+            timeout_ms=args.timeout_ms)
+        return {
+            '%s:%d' % (tp.topic, tp.partition): (err.__name__ if err else None)
+            for tp, err in results.items()
+        }

test/admin/test_admin_partitions.py

@@ -4,6 +4,7 @@
 
 from kafka.admin import OffsetSpec, KafkaAdminClient
 from kafka.errors import (
+    InvalidPartitionsError,
     NotLeaderForPartitionError,
     UnknownTopicOrPartitionError,
     IncompatibleBrokerVersion,
@@ -25,7 +26,7 @@
 
 
 class TestAlterPartitionReassignmentsMockBroker:
-    def test_success_returns_dict(self, broker, admin):
+    def test_success_returns_per_partition_none(self, broker, admin):
         Topic = AlterPartitionReassignmentsResponse.ReassignableTopicResponse
         Partition = Topic.ReassignablePartitionResponse
         broker.respond(
@@ -39,6 +40,7 @@ def test_success_returns_dict(self, broker, admin):
                         name='topic-a',
                         partitions=[
                             Partition(partition_index=0, error_code=0, error_message=None),
+                            Partition(partition_index=1, error_code=0, error_message=None),
                         ],
                     ),
                 ],
@@ -47,11 +49,13 @@ def test_success_returns_dict(self, broker, admin):
 
         result = admin.alter_partition_reassignments({
             TopicPartition('topic-a', 0): [1, 2, 3],
+            TopicPartition('topic-a', 1): [4, 5, 6],
         })
 
-        assert result['error_code'] == 0
-        assert result['responses'][0]['name'] == 'topic-a'
-        assert result['responses'][0]['partitions'][0]['error_code'] == 0
+        assert result == {
+            TopicPartition('topic-a', 0): None,
+            TopicPartition('topic-a', 1): None,
+        }
 
     def test_cancel_reassignment_sends_null_replicas(self, broker, admin):
         captured = {}
@@ -66,7 +70,7 @@ def handler(api_key, api_version, correlation_id, request_bytes):
         broker.respond_fn(AlterPartitionReassignmentsRequest, handler)
 
         admin.alter_partition_reassignments({
-            TopicPartition('topic-a', 0): None,  # cancel
+            TopicPartition('topic-a', 0): None,
             TopicPartition('topic-a', 1): [4, 5],
         })
 
@@ -77,7 +81,7 @@ def handler(api_key, api_version, correlation_id, request_bytes):
         assert by_index[0].replicas is None
         assert list(by_index[1].replicas) == [4, 5]
 
-    def test_partition_level_error_raises(self, broker, admin):
+    def test_partition_level_error_returned_in_dict(self, broker, admin):
         Topic = AlterPartitionReassignmentsResponse.ReassignableTopicResponse
         Partition = Topic.ReassignablePartitionResponse
         broker.respond(
@@ -90,45 +94,69 @@ def test_partition_level_error_raises(self, broker, admin):
                     Topic(
                         name='topic-a',
                         partitions=[
-                            Partition(partition_index=0, error_code=37,  # InvalidPartitionsError
+                            Partition(partition_index=0, error_code=0, error_message=None),
+                            Partition(partition_index=1, error_code=37,
                                       error_message='bad partition'),
                         ],
                     ),
                 ],
             ),
         )
 
+        result = admin.alter_partition_reassignments({
+            TopicPartition('topic-a', 0): [1, 2, 3],
+            TopicPartition('topic-a', 1): [4, 5, 6],
+        })
+
+        assert result == {
+            TopicPartition('topic-a', 0): None,
+            TopicPartition('topic-a', 1): InvalidPartitionsError,
+        }
+
+    def test_top_level_error_raises(self, broker, admin):
+        broker.respond(
+            AlterPartitionReassignmentsRequest,
+            AlterPartitionReassignmentsResponse(
+                throttle_time_ms=0,
+                error_code=38,  # ClusterAuthorizationFailedError
+                error_message='not authorized',
+                responses=[],
+            ),
+        )
+
         with pytest.raises(Exception):
             admin.alter_partition_reassignments({
                 TopicPartition('topic-a', 0): [1, 2, 3],
             })
 
-    def test_partition_error_suppressed_with_raise_errors_false(self, broker, admin):
-        Topic = AlterPartitionReassignmentsResponse.ReassignableTopicResponse
-        Partition = Topic.ReassignablePartitionResponse
+    def test_empty_replica_list_rejected(self, broker, admin):
+        with pytest.raises(ValueError, match='non-empty'):
+            admin.alter_partition_reassignments({
+                TopicPartition('topic-a', 0): [],
+            })
+
+    def test_non_int_replica_rejected(self, broker, admin):
+        with pytest.raises(ValueError, match='int broker_ids'):
+            admin.alter_partition_reassignments({
+                TopicPartition('topic-a', 0): ['1', '2', '3'],
+            })
+
+    def test_missing_partition_response_is_absent(self, broker, admin):
         broker.respond(
             AlterPartitionReassignmentsRequest,
             AlterPartitionReassignmentsResponse(
                 throttle_time_ms=0,
                 error_code=0,
                 error_message=None,
-                responses=[
-                    Topic(
-                        name='topic-a',
-                        partitions=[
-                            Partition(partition_index=0, error_code=37, error_message='bad'),
-                        ],
-                    ),
-                ],
+                responses=[],
             ),
         )
 
-        result = admin.alter_partition_reassignments(
-            {TopicPartition('topic-a', 0): [1, 2, 3]},
-            raise_errors=False,
-        )
+        result = admin.alter_partition_reassignments({
+            TopicPartition('topic-a', 0): [1, 2, 3],
+        })
 
-        assert result['responses'][0]['partitions'][0]['error_code'] == 37
+        assert result == {}
 
 
 # ---------------------------------------------------------------------------
@@ -242,15 +270,59 @@ def test_top_level_error_raises(self, broker, admin):
             ListPartitionReassignmentsRequest,
             ListPartitionReassignmentsResponse(
                 throttle_time_ms=0,
-                error_code=41,  # NotControllerError
-                error_message='not controller',
+                error_code=29,  # TopicAuthorizationFailedError
+                error_message='not authorized',
                 topics=[],
             ),
         )
 
         with pytest.raises(Exception) as exc_info:
             admin.list_partition_reassignments()
-        assert 'not controller' in str(exc_info.value)
+        assert 'not authorized' in str(exc_info.value)
+
+    def test_not_controller_retries(self, broker, admin):
+        Topic = ListPartitionReassignmentsResponse.OngoingTopicReassignment
+        Partition = Topic.OngoingPartitionReassignment
+        # First response: NotControllerError. After controller refresh, success.
+        broker.respond(
+            ListPartitionReassignmentsRequest,
+            ListPartitionReassignmentsResponse(
+                throttle_time_ms=0,
+                error_code=41,  # NotControllerError
+                error_message='not controller',
+                topics=[],
+            ),
+        )
+        broker.respond(
+            ListPartitionReassignmentsRequest,
+            ListPartitionReassignmentsResponse(
+                throttle_time_ms=0,
+                error_code=0,
+                error_message=None,
+                topics=[
+                    Topic(
+                        name='topic-a',
+                        partitions=[
+                            Partition(
+                                partition_index=0,
+                                replicas=[1, 2, 3],
+                                adding_replicas=[4],
+                                removing_replicas=[1],
+                            ),
+                        ],
+                    ),
+                ],
+            ),
+        )
+
+        result = admin.list_partition_reassignments()
+        assert result == {
+            TopicPartition('topic-a', 0): {
+                'replicas': [1, 2, 3],
+                'adding_replicas': [4],
+                'removing_replicas': [1],
+            },
+        }
 
 
 # ---------------------------------------------------------------------------

test/integration/test_admin_integration.py

@@ -525,15 +525,12 @@ def test_describe_metadata_quorum(kafka_admin_client):
 
 @pytest.mark.skipif(env_kafka_version() < (2, 4), reason="AlterPartitionReassignments requires broker >=2.4")
 def test_alter_partition_reassignments(kafka_admin_client, topic):
-    topic_metadata = kafka_admin_client.describe_topics([topic])[0]
     brokers = [b.node_id for b in kafka_admin_client._manager.cluster.brokers()]
     # Single-broker cluster: only valid reassignment target is [broker]
     tp = TopicPartition(topic, 0)
 
     result = kafka_admin_client.alter_partition_reassignments({tp: brokers})
-    assert result['error_code'] == 0
-    assert len(result['responses']) == 1
-    assert result['responses'][0]['name'] == topic
+    assert result == {tp: None}
 
 
 @pytest.mark.skipif(env_kafka_version() < (2, 4), reason="ListPartitionReassignments requires broker >=2.4")