kakfa从入门到放弃(二): golang编程操作kafka

一、kafka代理及创建topic

package amqp

import (
	"context"
	"errors"

	"project/pkg/logger"
	"github.com/Shopify/sarama"
)

const (
	prefix = "hello_" // topic前缀
)

var (
	TopicTTLPolicy    = "delete"
	TopicTTLRetention = "259200000" // 3d->259200000ms; 2d ->172800000ms;
)

// Option 对kafka连接参数的处理.
type Option func(*sarama.Config)

type KafkaClient struct {
   
	conf        *sarama.Config
	cli         sarama.ClusterAdmin
	lg          logger.Logger
	TopicPrefix string
}

// NewKafkaClient 实例化kafka连接实例
// retrun: kafkaClient, kafkaCloseFunc, error
func NewKafkaClient(addrs []string, logger logger.Logger) (*KafkaClient, func(), error) {
   
	conf := sarama.NewConfig()
	conf.Version = sarama.DefaultVersion

	cli, err := sarama.NewClusterAdmin(addrs, conf)
	if err != nil {
   
		return nil, func() {
   }, err
	}
	kfk := &KafkaClient{
   
		conf:        conf,
		cli:         cli,
		lg:          logger,
		TopicPrefix: prefix,
	}
	return kfk, kfk.close, nil
}

func (kc *KafkaClient) close() {
   
	if kc.cli != nil {
   
		_ = kc.cli.Close()
	}
}

// GetTopics 获取当前kafka中的topic
func (kc *KafkaClient) GetTopics(ctx context.Context) (topicList []string, err error) {
   
	curTopics, err := kc.cli.ListTopics()
	if err != nil || len(curTopics) == 0 {
   
		return
	}
	for name := range curTopics {
   
		topicList = append(topicList, name)
	}
	return
}

// GetBrokers 获取集群中的borker list
func (kc *KafkaClient) GetBrokers(ctx context.Context) (brokers []*sarama.Broker, err error) {
   
	brokers, _, err = kc.cli.DescribeCluster()
	return
}

// TopicConfig 定义创建topic时需要的参数
type TopicConfig struct {
   
	TTLRetention string
	TTLPolicy    string
	PartitionNum int
	ReplicNum    int
}

func NewTopicConfig(partitionNum, replicNum int) *TopicConfig {
   
	return &TopicConfig{
   
		PartitionNum: partitionNum,
		ReplicNum:    replicNum,
		TTLRetention: TopicTTLRetention,
		TTLPolicy:    TopicTTLPolicy,
	}
}

// CreateTopics 创建topics
func (kc *KafkaClient) CreateTopics(ctx context.Context, tarTopics []string, tpConf *TopicConfig) (errs map[string]error) {
   
	errs = map[string]error{
   }
	for _, name := range tarTopics {
   
		if err := kc.CreateTopic(ctx, name, tpConf); err != nil {
   
			errs[name] = err
		}
	}
	return
}

// CreateTopic 创建单个Topic
func (kc *KafkaClient) CreateTopic(ctx context.Context, topicName string, tpConf *TopicConfig) (err error) {
   
	topicName = kc.TopicPrefix + topicName
	topics, err := kc.cli.DescribeTopics([]string{
   topicName})
	if err != nil {
   
		return
	}
	topic := topics[0]
	if errors.Is(topic.Err, sarama.ErrNoError) {
   
		// 已存在
		kc.lg.Infof("topic[%s] existing", topicName)
		return
	}
	return kc.cli.CreateTopic(topicName, &sarama.TopicDetail{
   
		NumPartitions:     int32(tpConf.PartitionNum),
		ReplicationFactor: int16(tpConf.ReplicNum),
		ConfigEntries: map[string]*string{
   
			"cleanup.policy":      &tpConf.TTLPolicy,
			"delete.retention.ms": &tpConf.TTLRetention,
		},
	}, false)
}

单元测试:

package amqp

import (
	"context"
	"testing"

	"project/pkg/logger"
	"github.com/stretchr/testify/assert"
)

var (
	kfkAddrs = []string{
   "localhost:9093"}
)

func TestKafksClient(t *testing.T) {
   
	ctx := context.Background()
	logger, lgErr := logger.NewLogger(logger.SetDebug(true))
	assert.NoError(t, lgErr)
	kc, closeFunc, kcErr := NewKafkaClient(kfkAddrs, logger)
	assert.NoError(t, kcErr)
	defer closeFunc() // nolint

	t.Run("broker", func(t *testing.T) {
   
		brokers, bErr := kc.GetBrokers(ctx)
		assert.NoError(t, bErr)
		assert.EqualValues(t, len(kfkAddrs), len(brokers))
	})

	t.Run("CreateTopic", func(t *testing.T) {
   
		cases := []struct {
    // nolint
			name              string
			topics            []string
			partitionNum      int // 分区
			replicationFactor int // 副本
			wantErr           bool
			checkFunc         func(t *testing.T)
		}{
   
			{
   
				name:              "firstTopic",
				topics:            []string{
   "hello"},
				partitionNum:      3,
				replicationFactor: 1,
				wantErr:           false,
				checkFunc: func(t *testing.T) {
   
					topics, tpErr := kc.GetTopics(ctx)
					assert.NoError(t, tpErr)
					assert.Truef(t, NewStringSet(topics...).Has(kc.TopicPrefix+"hello"), "firstTopic")
				},
			},
			{
   
				name:              "bigReplication",
				topics:            []string{
   "hello1"},
				partitionNum:      3,
				replicationFactor: 3,
				wantErr:           true,
				checkFunc:         func(t *testing.T) {
   },
			},
			{
   
				name:              "someTopics",
				topics:            []string{
   "this", "is", "topic", "hello"},
				partitionNum:      3,
				replicationFactor: 1,
				wantErr:           false,
				checkFunc: func(t *testing.T) {
   
					topics, tpErr := kc.GetTopics(ctx)
					assert.NoError(t, tpErr)
					assert.Truef(t,
						NewStringSet