package kafkagen import ( "encoding/binary" "fmt" "log" "os" 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" liavro "github.com/linkedin/goavro/v2" "github.com/pkg/errors" ) type Main struct { idk.ConfluentCommand `flag:"!embed"` SchemaFile string Subject string Topic string DryRun bool configMap *confluent.ConfigMap } func NewMain() (*Main, error) { var err error m := &Main{ConfluentCommand: idk.ConfluentCommand{}} m.SchemaFile = "bigschema.json" m.Subject = "bigschema" m.SchemaRegistryURL = "localhost:8081" m.KafkaBootstrapServers = []string{"localhost:9092"} m.Topic = "defaulttopic" m.configMap, err = common.SetupConfluent(&m.ConfluentCommand) return m, err } func (m *Main) Run() (err error) { if m.DryRun { log.Printf("%+v\n", m) return nil } licodec, err := decodeSchema(m.SchemaFile) if err != nil { return errors.Wrap(err, "decoding schema") } schemaID, err := m.postSchema(m.SchemaFile, m.Subject) if err != nil { return errors.Wrap(err, "psting schema") } fields := []string{"abc", "db", "user_id", "all_users", "has_deleted_date", "central_group", "custom_audiences", "desktop_boolean", "desktop_frequency", "desktop_recency", "product_boolean_historical_forestry_cravings_or_bugles", "ddd_category_total_current_rhinocerous_checking", "ddd_category_total_current_rhinocerous_thedog_cheetah", "survey1234", "days_since_last_logon", "elephant_added_for_account"} // make a bunch of data and insert it records := [][]interface{}{ {"2", "1", 159, map[string]interface{}{"boolean": true}, map[string]interface{}{"boolean": false}, map[string]interface{}{"string": "cgr"}, map[string]interface{}{"array": []string{"a", "b"}}, nil, map[string]interface{}{"int": 7}, nil, nil, map[string]interface{}{"float": 5.4}, nil, map[string]interface{}{"org.test.survey1234": "yes"}, map[string]interface{}{"float": 8.0}, nil}, } // kakfa.SetupSasl(configMap *confluent.ConfigMap, SASLConfig SaslConfig) (err error) { p, err := confluent.NewProducer(m.configMap) if err != nil { fmt.Printf("Failed to create producer: %s\n", err) os.Exit(1) } defer p.Close() doneChan := make(chan bool) go func() { defer close(doneChan) cnt := 0 for e := range p.Events() { switch ev := e.(type) { case *confluent.Message: m := ev if m.TopicPartition.Error != nil { fmt.Printf("Delivery failed: %v\n", m.TopicPartition.Error) } cnt += 1 if cnt%10000 == 0 { fmt.Printf("Delivered %v messages to topic %v\n", cnt, m.TopicPartition.Topic) } if cnt == len(records) { return } default: fmt.Printf("Ignored event: %s\n", ev) } } }() for _, vals := range records { rec, err := makeRecord(fields, vals) if err != nil { return errors.Wrap(err, "making record") } err = m.putRecordKafka(p, schemaID, licodec, "akey", rec) if err != nil { return errors.Wrap(err, "putting record") } } <-doneChan // wait till all messages are acked p.Flush(10 * 1000) return nil } func (m *Main) putRecordKafka(p *confluent.Producer, schemaID int, schema *liavro.Codec, key string, record map[string]interface{}) 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 err } p.ProduceChannel() <- &confluent.Message{ TopicPartition: confluent.TopicPartition{Topic: &m.Topic, Partition: confluent.PartitionAny}, Key: []byte(key), Value: buf, } return nil } func readSchema(filename string) (string, error) { bytes, err := os.ReadFile(filename) if err != nil { return "", err } return string(bytes), nil } func decodeSchema(filename string) (*liavro.Codec, error) { s, err := readSchema(filename) if err != nil { return nil, errors.Wrap(err, "reading schema") } codec, err := liavro.NewCodec(s) if err != nil { return nil, err } return codec, nil } func (m *Main) postSchema(schemaFile, subj string) (schemaID int, err error) { var auth *csrc.BasicAuth if m.SchemaRegistryUsername != "" { auth = &csrc.BasicAuth{ KafkaSchemaApiKey: m.SchemaRegistryUsername, KafkaSchemaApiSecret: m.SchemaRegistryPassword, } } schemaClient := csrc.NewClient("http://"+m.SchemaRegistryURL, nil, auth) schemaStr, err := readSchema(schemaFile) if err != nil { return 0, errors.Wrap(err, "reading schema file") } resp, err := schemaClient.PostSubjects(subj, schemaStr) if err != nil { return 0, errors.Wrap(err, "posting schema") } return resp.ID, nil } func makeRecord(fields []string, vals []interface{}) (map[string]interface{}, error) { if len(fields) != len(vals) { return nil, errors.Errorf("have %d fields and %d vals", len(fields), len(vals)) } ret := make(map[string]interface{}) for i, field := range fields { ret[field] = vals[i] } return ret, nil }