github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/sink/kafka/cluster_admin_client_mock_impl.go (about)

     1  // Copyright 2021 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package kafka
    15  
    16  import (
    17  	"context"
    18  	"fmt"
    19  	"strconv"
    20  
    21  	"github.com/IBM/sarama"
    22  	"github.com/pingcap/tiflow/pkg/errors"
    23  )
    24  
    25  const (
    26  	// DefaultMockTopicName specifies the default mock topic name.
    27  	DefaultMockTopicName = "mock_topic"
    28  	// DefaultMockPartitionNum is the default partition number of default mock topic.
    29  	DefaultMockPartitionNum = 3
    30  	// defaultMockControllerID specifies the default mock controller ID.
    31  	defaultMockControllerID = 1
    32  	// topic replication factor must be 3 for Confluent Cloud Kafka.
    33  	defaultReplicationFactor = 3
    34  )
    35  
    36  const (
    37  	// defaultMaxMessageBytes specifies the default max message bytes,
    38  	// default to 1048576, identical to kafka broker's `message.max.bytes` and topic's `max.message.bytes`
    39  	// see: https://kafka.apache.org/documentation/#brokerconfigs_message.max.bytes
    40  	// see: https://kafka.apache.org/documentation/#topicconfigs_max.message.bytes
    41  	defaultMaxMessageBytes = "1048588"
    42  
    43  	// defaultMinInsyncReplicas specifies the default `min.insync.replicas` for broker and topic.
    44  	defaultMinInsyncReplicas = "1"
    45  )
    46  
    47  var (
    48  	// BrokerMessageMaxBytes is the broker's `message.max.bytes`
    49  	BrokerMessageMaxBytes = defaultMaxMessageBytes
    50  	// TopicMaxMessageBytes is the topic's `max.message.bytes`
    51  	TopicMaxMessageBytes = defaultMaxMessageBytes
    52  	// MinInSyncReplicas is the `min.insync.replicas`
    53  	MinInSyncReplicas = defaultMinInsyncReplicas
    54  )
    55  
    56  type topicDetail struct {
    57  	TopicDetail
    58  	fetchesRemainingUntilVisible int
    59  }
    60  
    61  // ClusterAdminClientMockImpl mock implements the admin client interface.
    62  type ClusterAdminClientMockImpl struct {
    63  	topics map[string]*topicDetail
    64  	// Cluster controller ID.
    65  	controllerID  int
    66  	brokerConfigs map[string]string
    67  	topicConfigs  map[string]map[string]string
    68  }
    69  
    70  // NewClusterAdminClientMockImpl news a ClusterAdminClientMockImpl struct with default configurations.
    71  func NewClusterAdminClientMockImpl() *ClusterAdminClientMockImpl {
    72  	topics := make(map[string]*topicDetail)
    73  	topics[DefaultMockTopicName] = &topicDetail{
    74  		fetchesRemainingUntilVisible: 0,
    75  		TopicDetail: TopicDetail{
    76  			Name:          DefaultMockTopicName,
    77  			NumPartitions: 3,
    78  		},
    79  	}
    80  
    81  	brokerConfigs := make(map[string]string)
    82  	brokerConfigs[BrokerMessageMaxBytesConfigName] = BrokerMessageMaxBytes
    83  	brokerConfigs[MinInsyncReplicasConfigName] = MinInSyncReplicas
    84  
    85  	topicConfigs := make(map[string]map[string]string)
    86  	topicConfigs[DefaultMockTopicName] = make(map[string]string)
    87  	topicConfigs[DefaultMockTopicName][TopicMaxMessageBytesConfigName] = TopicMaxMessageBytes
    88  	topicConfigs[DefaultMockTopicName][MinInsyncReplicasConfigName] = MinInSyncReplicas
    89  
    90  	return &ClusterAdminClientMockImpl{
    91  		topics:        topics,
    92  		controllerID:  defaultMockControllerID,
    93  		brokerConfigs: brokerConfigs,
    94  		topicConfigs:  topicConfigs,
    95  	}
    96  }
    97  
    98  // GetAllBrokers implement the ClusterAdminClient interface
    99  func (c *ClusterAdminClientMockImpl) GetAllBrokers(context.Context) ([]Broker, error) {
   100  	return nil, nil
   101  }
   102  
   103  // GetBrokerConfig implement the ClusterAdminClient interface
   104  func (c *ClusterAdminClientMockImpl) GetBrokerConfig(
   105  	_ context.Context,
   106  	configName string,
   107  ) (string, error) {
   108  	value, ok := c.brokerConfigs[configName]
   109  	if !ok {
   110  		return "", errors.ErrKafkaConfigNotFound.GenWithStack(
   111  			"cannot find the `%s` from the broker's configuration", configName)
   112  	}
   113  	return value, nil
   114  }
   115  
   116  // GetTopicConfig implement the ClusterAdminClient interface
   117  func (c *ClusterAdminClientMockImpl) GetTopicConfig(ctx context.Context, topicName string, configName string) (string, error) {
   118  	if _, ok := c.topics[topicName]; !ok {
   119  		return "", errors.ErrKafkaConfigNotFound.GenWithStack("cannot find the `%s` from the topic's configuration", topicName)
   120  	}
   121  	value, ok := c.topicConfigs[topicName][configName]
   122  	if !ok {
   123  		return "", errors.ErrKafkaConfigNotFound.GenWithStack(
   124  			"cannot find the `%s` from the topic's configuration", configName)
   125  	}
   126  	return value, nil
   127  }
   128  
   129  // SetRemainingFetchesUntilTopicVisible is used to control the visibility of a specific topic.
   130  // It is used to mock the topic creation delay.
   131  func (c *ClusterAdminClientMockImpl) SetRemainingFetchesUntilTopicVisible(
   132  	topicName string,
   133  	fetchesRemainingUntilVisible int,
   134  ) error {
   135  	topic, ok := c.topics[topicName]
   136  	if !ok {
   137  		return fmt.Errorf("No such topic as %s", topicName)
   138  	}
   139  
   140  	topic.fetchesRemainingUntilVisible = fetchesRemainingUntilVisible
   141  	return nil
   142  }
   143  
   144  // GetTopicsMeta implement the ClusterAdminClient interface
   145  func (c *ClusterAdminClientMockImpl) GetTopicsMeta(
   146  	_ context.Context,
   147  	topics []string,
   148  	_ bool,
   149  ) (map[string]TopicDetail, error) {
   150  	result := make(map[string]TopicDetail, len(topics))
   151  	for _, topic := range topics {
   152  		details, ok := c.topics[topic]
   153  		if ok {
   154  			if details.fetchesRemainingUntilVisible > 0 {
   155  				details.fetchesRemainingUntilVisible--
   156  				continue
   157  			}
   158  			result[topic] = details.TopicDetail
   159  		}
   160  	}
   161  	return result, nil
   162  }
   163  
   164  // GetTopicsPartitionsNum implement the ClusterAdminClient interface
   165  func (c *ClusterAdminClientMockImpl) GetTopicsPartitionsNum(
   166  	_ context.Context, topics []string,
   167  ) (map[string]int32, error) {
   168  	result := make(map[string]int32, len(topics))
   169  	for _, topic := range topics {
   170  		result[topic] = c.topics[topic].NumPartitions
   171  	}
   172  	return result, nil
   173  }
   174  
   175  // CreateTopic adds topic into map.
   176  func (c *ClusterAdminClientMockImpl) CreateTopic(
   177  	_ context.Context,
   178  	detail *TopicDetail,
   179  	_ bool,
   180  ) error {
   181  	if detail.ReplicationFactor > defaultReplicationFactor {
   182  		return sarama.ErrInvalidReplicationFactor
   183  	}
   184  
   185  	_, minInsyncReplicaConfigFound := c.brokerConfigs[MinInsyncReplicasConfigName]
   186  	// For Confluent Cloud, min.insync.replica is invisible and replication factor must be 3.
   187  	// Otherwise, ErrPolicyViolation is expected to be returned.
   188  	if !minInsyncReplicaConfigFound &&
   189  		detail.ReplicationFactor != defaultReplicationFactor {
   190  		return sarama.ErrPolicyViolation
   191  	}
   192  
   193  	c.topics[detail.Name] = &topicDetail{
   194  		TopicDetail: *detail,
   195  	}
   196  	return nil
   197  }
   198  
   199  // DeleteTopic deletes a topic, only used for testing.
   200  func (c *ClusterAdminClientMockImpl) DeleteTopic(topicName string) {
   201  	delete(c.topics, topicName)
   202  }
   203  
   204  // Close do nothing.
   205  func (c *ClusterAdminClientMockImpl) Close() {}
   206  
   207  // SetMinInsyncReplicas sets the MinInsyncReplicas for broker and default topic.
   208  func (c *ClusterAdminClientMockImpl) SetMinInsyncReplicas(minInsyncReplicas string) {
   209  	c.topicConfigs[DefaultMockTopicName][MinInsyncReplicasConfigName] = minInsyncReplicas
   210  	c.brokerConfigs[MinInsyncReplicasConfigName] = minInsyncReplicas
   211  }
   212  
   213  // GetDefaultMockTopicName returns the default topic name
   214  func (c *ClusterAdminClientMockImpl) GetDefaultMockTopicName() string {
   215  	return DefaultMockTopicName
   216  }
   217  
   218  // GetBrokerMessageMaxBytes returns broker's `message.max.bytes`
   219  func (c *ClusterAdminClientMockImpl) GetBrokerMessageMaxBytes() int {
   220  	messageMaxBytes, _ := strconv.Atoi(BrokerMessageMaxBytes)
   221  	return messageMaxBytes
   222  }
   223  
   224  // GetTopicMaxMessageBytes returns topic's `max.message.bytes`
   225  func (c *ClusterAdminClientMockImpl) GetTopicMaxMessageBytes() int {
   226  	maxMessageBytes, _ := strconv.Atoi(TopicMaxMessageBytes)
   227  	return maxMessageBytes
   228  }
   229  
   230  // DropBrokerConfig remove all broker level configuration for test purpose.
   231  func (c *ClusterAdminClientMockImpl) DropBrokerConfig(configName string) {
   232  	delete(c.brokerConfigs, configName)
   233  }