github.com/streamdal/segmentio-kafka-go@v0.4.47-streamdal/listpartitionreassignments.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  	"net"
     6  	"time"
     7  
     8  	"github.com/segmentio/kafka-go/protocol/listpartitionreassignments"
     9  )
    10  
    11  // ListPartitionReassignmentsRequest is a request to the ListPartitionReassignments API.
    12  type ListPartitionReassignmentsRequest struct {
    13  	// Address of the kafka broker to send the request to.
    14  	Addr net.Addr
    15  
    16  	// Topics we want reassignments for, mapped by their name, or nil to list everything.
    17  	Topics map[string]ListPartitionReassignmentsRequestTopic
    18  
    19  	// Timeout is the amount of time to wait for the request to complete.
    20  	Timeout time.Duration
    21  }
    22  
    23  // ListPartitionReassignmentsRequestTopic contains the requested partitions for a single
    24  // topic.
    25  type ListPartitionReassignmentsRequestTopic struct {
    26  	// The partitions to list partition reassignments for.
    27  	PartitionIndexes []int
    28  }
    29  
    30  // ListPartitionReassignmentsResponse is a response from the ListPartitionReassignments API.
    31  type ListPartitionReassignmentsResponse struct {
    32  	// Error is set to a non-nil value including the code and message if a top-level
    33  	// error was encountered.
    34  	Error error
    35  
    36  	// Topics contains results for each topic, mapped by their name.
    37  	Topics map[string]ListPartitionReassignmentsResponseTopic
    38  }
    39  
    40  // ListPartitionReassignmentsResponseTopic contains the detailed result of
    41  // ongoing reassignments for a topic.
    42  type ListPartitionReassignmentsResponseTopic struct {
    43  	// Partitions contains result for topic partitions.
    44  	Partitions []ListPartitionReassignmentsResponsePartition
    45  }
    46  
    47  // ListPartitionReassignmentsResponsePartition contains the detailed result of
    48  // ongoing reassignments for a single partition.
    49  type ListPartitionReassignmentsResponsePartition struct {
    50  	// PartitionIndex contains index of the partition.
    51  	PartitionIndex int
    52  
    53  	// Replicas contains the current replica set.
    54  	Replicas []int
    55  
    56  	// AddingReplicas contains the set of replicas we are currently adding.
    57  	AddingReplicas []int
    58  
    59  	// RemovingReplicas contains the set of replicas we are currently removing.
    60  	RemovingReplicas []int
    61  }
    62  
    63  func (c *Client) ListPartitionReassignments(
    64  	ctx context.Context,
    65  	req *ListPartitionReassignmentsRequest,
    66  ) (*ListPartitionReassignmentsResponse, error) {
    67  	apiReq := &listpartitionreassignments.Request{
    68  		TimeoutMs: int32(req.Timeout.Milliseconds()),
    69  	}
    70  
    71  	for topicName, topicReq := range req.Topics {
    72  		apiReq.Topics = append(
    73  			apiReq.Topics,
    74  			listpartitionreassignments.RequestTopic{
    75  				Name:             topicName,
    76  				PartitionIndexes: intToInt32Array(topicReq.PartitionIndexes),
    77  			},
    78  		)
    79  	}
    80  
    81  	protoResp, err := c.roundTrip(
    82  		ctx,
    83  		req.Addr,
    84  		apiReq,
    85  	)
    86  	if err != nil {
    87  		return nil, err
    88  	}
    89  	apiResp := protoResp.(*listpartitionreassignments.Response)
    90  
    91  	resp := &ListPartitionReassignmentsResponse{
    92  		Error:  makeError(apiResp.ErrorCode, apiResp.ErrorMessage),
    93  		Topics: make(map[string]ListPartitionReassignmentsResponseTopic),
    94  	}
    95  
    96  	for _, topicResult := range apiResp.Topics {
    97  		respTopic := ListPartitionReassignmentsResponseTopic{}
    98  		for _, partitionResult := range topicResult.Partitions {
    99  			respTopic.Partitions = append(
   100  				respTopic.Partitions,
   101  				ListPartitionReassignmentsResponsePartition{
   102  					PartitionIndex:   int(partitionResult.PartitionIndex),
   103  					Replicas:         int32ToIntArray(partitionResult.Replicas),
   104  					AddingReplicas:   int32ToIntArray(partitionResult.AddingReplicas),
   105  					RemovingReplicas: int32ToIntArray(partitionResult.RemovingReplicas),
   106  				},
   107  			)
   108  		}
   109  		resp.Topics[topicResult.Name] = respTopic
   110  	}
   111  
   112  	return resp, nil
   113  }
   114  
   115  func intToInt32Array(arr []int) []int32 {
   116  	if arr == nil {
   117  		return nil
   118  	}
   119  	res := make([]int32, len(arr))
   120  	for i := range arr {
   121  		res[i] = int32(arr[i])
   122  	}
   123  	return res
   124  }
   125  
   126  func int32ToIntArray(arr []int32) []int {
   127  	if arr == nil {
   128  		return nil
   129  	}
   130  	res := make([]int, len(arr))
   131  	for i := range arr {
   132  		res[i] = int(arr[i])
   133  	}
   134  	return res
   135  }