package kafka import ( "context" "encoding/binary" "os" "time" confluent "github.com/confluentinc/confluent-kafka-go/kafka" "github.com/featurebasedb/featurebase/v3/idk" "github.com/featurebasedb/featurebase/v3/idk/common" "github.com/featurebasedb/featurebase/v3/idk/kafka/csrc" "github.com/featurebasedb/featurebase/v3/logger" liavro "github.com/linkedin/goavro/v2" "github.com/pkg/errors" ) type PutCmd struct { idk.ConfluentCommand `flag:"!embed"` Topic string `help:"Kafka topic to post to."` Schema string SchemaFile string Subject string Data string DryRun bool `help:"Dry run - just flag parsing."` configMap *confluent.ConfigMap `flag:"-"` log logger.Logger NumPartitions int ReplicationFactor int } func (p *PutCmd) Log() logger.Logger { return p.log } func NewPutCmd() (*PutCmd, error) { var err error p := PutCmd{} p.ConfluentCommand = idk.ConfluentCommand{} p.KafkaBootstrapServers = []string{"localhost:9092"} p.SchemaRegistryURL = "localhost:8081" p.Topic = "defaulttopic" p.Schema = `{"type": "record","namespace": "c.e","name": "F","fields": [{"name":"id","type":"long"},{"name":"a","type": "boolean"},{"name": "b", "type": "float", "scale": 2}]}` p.SchemaFile = "" p.Subject = "test" p.Data = `{"id": 1, "a": true, "b": 1.43}` p.log = logger.NewVerboseLogger(os.Stderr) return &p, err } func (p *PutCmd) Run() (err error) { p.configMap, err = common.SetupConfluent(&p.ConfluentCommand) if err != nil { return errors.Wrap(err, "setting up confluent") } var auth *csrc.BasicAuth if p.SchemaRegistryUsername != "" { auth = &csrc.BasicAuth{ KafkaSchemaApiKey: p.SchemaRegistryUsername, KafkaSchemaApiSecret: p.SchemaRegistryPassword, } } client := csrc.NewClient(p.SchemaRegistryURL, nil, auth) schemaStr, err := p.getSchema() if err != nil { return errors.Wrap(err, "getting schema") } resp, err := client.PostSubjects(p.Subject, schemaStr) if err != nil { return errors.Wrap(err, "posting schema") } p.log.Printf("Posted schema ID: %d\n", resp.ID) licodec, err := liavro.NewCodec(schemaStr) if err != nil { return errors.Wrap(err, "li decoding schema") } rec, _, err := licodec.NativeFromTextual([]byte(p.Data)) if err != nil { return errors.Wrap(err, "decoding data") } p.log.Debugf("Decoded: %v\n", rec) data, err := endcodeAvro(int(resp.ID), licodec, rec.(map[string]interface{})) if err != nil { return errors.Wrap(err, "encoding data") } // Create Producer instance pr, err := confluent.NewProducer(p.configMap) if err != nil { return errors.Wrap(err, "failed to create producer") } defer pr.Close() err = CreateKafkaTopic(context.Background(), p.Topic, pr, p.NumPartitions, p.ReplicationFactor) if err != nil { return errors.Wrap(err, "creating topic") } delivery_chan := make(chan confluent.Event, 10) err = pr.Produce(&confluent.Message{ TopicPartition: confluent.TopicPartition{Topic: &p.Topic, Partition: confluent.PartitionAny}, Value: data, }, delivery_chan) e := <-delivery_chan m := e.(*confluent.Message) if m.TopicPartition.Error != nil { return m.TopicPartition.Error } close(delivery_chan) if err != nil { return errors.Wrap(err, "failed to write message to kafka") } return nil } func CreateKafkaTopic(ctx context.Context, topic string, p *confluent.Producer, numPartitions int, replicationFactor int) error { a, err := confluent.NewAdminClientFromProducer(p) if err != nil { return errors.Wrap(err, "failed to create new admin client from producer") } defer a.Close() // Contexts are used to abort or limit the amount of time // the Admin call blocks waiting for a result. ctx, cancel := context.WithCancel(ctx) defer cancel() // Create topics on cluster. // Set Admin options to wait up to 60s for the operation to finish on the remote cluster maxDur, err := time.ParseDuration("60s") if err != nil { return errors.Wrap(err, "ParseDuration(60s)") } // TODO: maybe add delete and retention period to config results, err := a.CreateTopics( ctx, // Multiple topics can be created simultaneously // by providing more TopicSpecification structs here. []confluent.TopicSpecification{{ Topic: topic, NumPartitions: numPartitions, ReplicationFactor: replicationFactor, }}, // Admin options confluent.SetAdminOperationTimeout(maxDur)) if err != nil { return errors.Wrap(err, "admin client request error") } for _, result := range results { if result.Error.Code() != confluent.ErrNoError && result.Error.Code() != confluent.ErrTopicAlreadyExists { return errors.Wrap(errors.New(result.Error.Error()), "failed to create topic") } } return nil } func (p *PutCmd) getSchema() (string, error) { if p.Schema == "" && p.SchemaFile == "" { return "", errors.New("need a string schema or schema file") } if p.Schema != "" { return p.Schema, nil } bytes, err := os.ReadFile(p.SchemaFile) if err != nil { return "", errors.Wrap(err, "reading schema file") } return string(bytes), nil } func endcodeAvro(schemaID int, schema *liavro.Codec, record map[string]interface{}) ([]byte, error) { buf := make([]byte, 5, 1000) buf[0] = 0 binary.BigEndian.PutUint32(buf[1:], uint32(schemaID)) buf, err := schema.BinaryFromNative(buf, record) if err != nil { return nil, errors.Errorf("encoding %v: %v", record, err) } return buf, nil }