github.com/segmentio/kafka-go@v0.4.48-0.20240318174348-3f6244eb34fd/alterpartitionreassignments.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  	"net"
     6  	"time"
     7  
     8  	"github.com/segmentio/kafka-go/protocol/alterpartitionreassignments"
     9  )
    10  
    11  // AlterPartitionReassignmentsRequest is a request to the AlterPartitionReassignments API.
    12  type AlterPartitionReassignmentsRequest struct {
    13  	// Address of the kafka broker to send the request to.
    14  	Addr net.Addr
    15  
    16  	// Topic is the name of the topic to alter partitions in. Keep this field empty and use Topic in AlterPartitionReassignmentsRequestAssignment to
    17  	// reassign to multiple topics.
    18  	Topic string
    19  
    20  	// Assignments is the list of partition reassignments to submit to the API.
    21  	Assignments []AlterPartitionReassignmentsRequestAssignment
    22  
    23  	// Timeout is the amount of time to wait for the request to complete.
    24  	Timeout time.Duration
    25  }
    26  
    27  // AlterPartitionReassignmentsRequestAssignment contains the requested reassignments for a single
    28  // partition.
    29  type AlterPartitionReassignmentsRequestAssignment struct {
    30  	// Topic is the name of the topic to alter partitions in. If empty, the value of Topic in AlterPartitionReassignmentsRequest is used.
    31  	Topic string
    32  
    33  	// PartitionID is the ID of the partition to make the reassignments in.
    34  	PartitionID int
    35  
    36  	// BrokerIDs is a slice of brokers to set the partition replicas to, or null to cancel a pending reassignment for this partition.
    37  	BrokerIDs []int
    38  }
    39  
    40  // AlterPartitionReassignmentsResponse is a response from the AlterPartitionReassignments API.
    41  type AlterPartitionReassignmentsResponse struct {
    42  	// Error is set to a non-nil value including the code and message if a top-level
    43  	// error was encountered when doing the update.
    44  	Error error
    45  
    46  	// PartitionResults contains the specific results for each partition.
    47  	PartitionResults []AlterPartitionReassignmentsResponsePartitionResult
    48  }
    49  
    50  // AlterPartitionReassignmentsResponsePartitionResult contains the detailed result of
    51  // doing reassignments for a single partition.
    52  type AlterPartitionReassignmentsResponsePartitionResult struct {
    53  	// Topic is the topic name.
    54  	Topic string
    55  
    56  	// PartitionID is the ID of the partition that was altered.
    57  	PartitionID int
    58  
    59  	// Error is set to a non-nil value including the code and message if an error was encountered
    60  	// during the update for this partition.
    61  	Error error
    62  }
    63  
    64  func (c *Client) AlterPartitionReassignments(
    65  	ctx context.Context,
    66  	req *AlterPartitionReassignmentsRequest,
    67  ) (*AlterPartitionReassignmentsResponse, error) {
    68  	apiTopicMap := make(map[string]*alterpartitionreassignments.RequestTopic)
    69  
    70  	for _, assignment := range req.Assignments {
    71  		topic := assignment.Topic
    72  		if topic == "" {
    73  			topic = req.Topic
    74  		}
    75  
    76  		apiTopic := apiTopicMap[topic]
    77  		if apiTopic == nil {
    78  			apiTopic = &alterpartitionreassignments.RequestTopic{
    79  				Name: topic,
    80  			}
    81  			apiTopicMap[topic] = apiTopic
    82  		}
    83  
    84  		replicas := []int32{}
    85  		for _, brokerID := range assignment.BrokerIDs {
    86  			replicas = append(replicas, int32(brokerID))
    87  		}
    88  
    89  		apiTopic.Partitions = append(
    90  			apiTopic.Partitions,
    91  			alterpartitionreassignments.RequestPartition{
    92  				PartitionIndex: int32(assignment.PartitionID),
    93  				Replicas:       replicas,
    94  			},
    95  		)
    96  	}
    97  
    98  	apiReq := &alterpartitionreassignments.Request{
    99  		TimeoutMs: int32(req.Timeout.Milliseconds()),
   100  	}
   101  
   102  	for _, apiTopic := range apiTopicMap {
   103  		apiReq.Topics = append(apiReq.Topics, *apiTopic)
   104  	}
   105  
   106  	protoResp, err := c.roundTrip(
   107  		ctx,
   108  		req.Addr,
   109  		apiReq,
   110  	)
   111  	if err != nil {
   112  		return nil, err
   113  	}
   114  	apiResp := protoResp.(*alterpartitionreassignments.Response)
   115  
   116  	resp := &AlterPartitionReassignmentsResponse{
   117  		Error: makeError(apiResp.ErrorCode, apiResp.ErrorMessage),
   118  	}
   119  
   120  	for _, topicResult := range apiResp.Results {
   121  		for _, partitionResult := range topicResult.Partitions {
   122  			resp.PartitionResults = append(
   123  				resp.PartitionResults,
   124  				AlterPartitionReassignmentsResponsePartitionResult{
   125  					Topic:       topicResult.Name,
   126  					PartitionID: int(partitionResult.PartitionIndex),
   127  					Error:       makeError(partitionResult.ErrorCode, partitionResult.ErrorMessage),
   128  				},
   129  			)
   130  		}
   131  	}
   132  
   133  	return resp, nil
   134  }