diff --git a/cli/batch/sql.go b/cli/batch/sql.go index 1f84420be..aadcb199e 100644 --- a/cli/batch/sql.go +++ b/cli/batch/sql.go @@ -175,10 +175,26 @@ func buildBulkInsert(tbl *dax.Table, fields []*dax.Field, ids []interface{}, row sb.WriteString(strings.Join(flds, ",")) // MAP - keyType := dax.BaseTypeID - if tbl.StringKeys() { - keyType = dax.BaseTypeString + + // We need to set ID value below based on ids. Rather than changing what we + // get by changing IDK, decide how to format based on the type of ids. We + // get []uint64 for not keyed indexes or []byte for keyed indexes. + // Additionally, IDK doesn't give us a valid table for some reason (the keys + // options is always false). We will also use ids type to determine how to + // map values in the MAP clause + var fmtStr string + var keyType string + if len(ids) > 0 { + switch ids[0].(type) { + case uint64: + fmtStr = "%d" + keyType = dax.BaseTypeID + case []byte: + fmtStr = "%s" + keyType = dax.BaseTypeString + } } + sb.WriteString(`) MAP ('$._id' `) sb.WriteString(keyType) sb.WriteString(`,`) @@ -191,9 +207,10 @@ func buildBulkInsert(tbl *dax.Table, fields []*dax.Field, ids []interface{}, row // bulk insert as one line in the NDJSON payload. We re-use the map for each // row. m := make(map[string]interface{}) + for i := range rows { // Write the ID value. - m[string(dax.PrimaryKeyFieldName)] = ids[i] + m[string(dax.PrimaryKeyFieldName)] = fmt.Sprintf(fmtStr, ids[i]) // Write the rest of the data values. for col := range rows[i] { m[fmt.Sprintf("col_%d", col)] = rows[i][col] diff --git a/cli/cli_kafka_runner_test.go b/cli/cli_kafka_runner_test.go new file mode 100644 index 000000000..3bfd66bd9 --- /dev/null +++ b/cli/cli_kafka_runner_test.go @@ -0,0 +1,343 @@ +package cli + +import ( + "bufio" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "strings" + "testing" + + "github.com/featurebasedb/featurebase/v3/cli/internal" + "github.com/featurebasedb/featurebase/v3/cli/kafka" + "github.com/featurebasedb/featurebase/v3/errors" + "github.com/featurebasedb/featurebase/v3/logger" + "github.com/stretchr/testify/require" +) + +// These variables are used throught tests to refer to resouces needed for the +// tests (i.e. where to find the kafka service) +var ( + pilosaHost string + pilosaTLSHost string + pilosaGrpcHost string + kafkaHost string + registryHost string + certPath string +) + +// The variables above are set here. If local is set to true, the tests assume +// all services are running locally. Otherwise, we use the name of the +// containers that these tests will have access to in CI. +func init() { + local := true + var ok bool + if pilosaHost, ok = os.LookupEnv("IDK_TEST_PILOSA_HOST"); !ok { + if local { + pilosaHost = "localhost:10101" + } else { + pilosaHost = "pilosa:10101" + } + } + if pilosaTLSHost, ok = os.LookupEnv("IDK_TEST_PILOSA_TLS_HOST"); !ok { + pilosaTLSHost = "https://pilosa-tls:10111" + } + if pilosaGrpcHost, ok = os.LookupEnv("IDK_TEST_PILOSA_GRPC_HOST"); !ok { + if local { + pilosaGrpcHost = "localhost:20101" + } else { + pilosaGrpcHost = "pilosa:20101" + } + } + if kafkaHost, ok = os.LookupEnv("IDK_TEST_KAFKA_HOST"); !ok { + if local { + kafkaHost = "localhost:9092" + } else { + kafkaHost = "kafka:9092" + } + } + if registryHost, ok = os.LookupEnv("IDK_TEST_REGISTRY_HOST"); !ok { + if local { + registryHost = "localhost:8081" + } else { + registryHost = "schema-registry:8081" + } + } + if certPath, ok = os.LookupEnv("IDK_TEST_CERT_PATH"); !ok { + certPath = "/certs" + } +} + +// Struct used for TestRunner which contains the information needed to ingest +// data to kafka, configure the runner, and test that the runner successfully +// ran. +type KafkaRunnerTest struct { + ConfigFile string // path to the configuration file used for the runner + DataFile string // path to the data file to populate kafka with + Encode string // how to encode the data from the data file + CreateTableStmt string // statement used to create table prior to ingest + Tests []TestQuery // list of test which are 2-tuples of query and expected results + +} + +// Struct which contains a query to run agaisnt featurebase and the expected +// results from that query. +type TestQuery struct { + Query string + ExpectedResp string +} + +// The paths (from the folder that contains this file) to folders with kafka +// runner configurations files and data to be sent and consumed from kafka. +var confPathPrefix string = "./kafka/runner_test_data/config/" +var dataPathPrefix string = "./kafka/runner_test_data/data/" + +// A slice of KafkaRunnerTest structs that will be used in TestKafkaRunner test +// function. +var kafkaRunnerTests = []KafkaRunnerTest{ + { // id keys + ConfigFile: "config00.toml", + DataFile: "data00.json", + Encode: kafka.JSON, + Tests: []TestQuery{ + { + Query: "Extract(Sort(All(), field=name), Rows(name), Rows(age), Rows(hobbies))", + ExpectedResp: `{"results":[{"fields":[{"name":"name","type":"string"},{"name":"age","type":"int64"},{"name":"hobbies","type":"[]string"}],"columns":[{"column":1,"rows":["a",20,["hob2","hob1"]]},{"column":2,"rows":["b",21,["hob2","hob3"]]},{"column":3,"rows":["c",22,["hob3","hob4"]]},{"column":4,"rows":["d",23,["hob4","hob5"]]},{"column":5,"rows":["e",24,["hob5","hob6"]]},{"column":6,"rows":["f",26,["hob6","hob7"]]}]}]}`, + }, + }, + CreateTableStmt: "(_id ID, name String, age Int, hobbies StringSet)", + }, + { // string keys + ConfigFile: "config01.toml", + DataFile: "data00.json", + Encode: kafka.JSON, + Tests: []TestQuery{ + { + Query: "Extract(Sort(All(), field=name), Rows(name), Rows(age), Rows(hobbies))", + ExpectedResp: `{"results":[{"fields":[{"name":"name","type":"string"},{"name":"age","type":"int64"},{"name":"hobbies","type":"[]string"}],"columns":[{"column":"1","rows":["a",20,["hob1","hob2"]]},{"column":"2","rows":["b",21,["hob2","hob3"]]},{"column":"3","rows":["c",22,["hob3","hob4"]]},{"column":"4","rows":["d",23,["hob4","hob5"]]},{"column":"5","rows":["e",24,["hob5","hob6"]]},{"column":"6","rows":["f",26,["hob6","hob7"]]}]}]}`, + }, + }, + CreateTableStmt: "(_id String, name String, age Int, hobbies StringSet)", + }, + { // two string keys + ConfigFile: "config02.toml", + DataFile: "data00.json", + Encode: kafka.JSON, + Tests: []TestQuery{ + { + Query: "Extract(Sort(All(), field=name), Rows(id), Rows(name), Rows(age), Rows(hobbies))", + ExpectedResp: `{"results":[{"fields":[{"name":"id","type":"string"},{"name":"name","type":"string"},{"name":"age","type":"int64"},{"name":"hobbies","type":"[]string"}],"columns":[{"column":"1|a","rows":["1","a",20,["hob1","hob2"]]},{"column":"2|b","rows":["2","b",21,["hob2","hob3"]]},{"column":"3|c","rows":["3","c",22,["hob3","hob4"]]},{"column":"4|d","rows":["4","d",23,["hob4","hob5"]]},{"column":"5|e","rows":["5","e",24,["hob5","hob6"]]},{"column":"6|f","rows":["6","f",26,["hob6","hob7"]]}]}]}`, + }, + }, + CreateTableStmt: "(_id String, id String, name String, age Int, hobbies StringSet)", + }, + { // string, id, and int + ConfigFile: "config03.toml", + DataFile: "data00.json", + Encode: kafka.JSON, + Tests: []TestQuery{ + { + Query: "Extract(Sort(All(), field=name), Rows(id), Rows(name), Rows(age), Rows(hobbies))", + ExpectedResp: `{"results":[{"fields":[{"name":"id","type":"uint64"},{"name":"name","type":"string"},{"name":"age","type":"int64"},{"name":"hobbies","type":"[]string"}],"columns":[{"column":"1|a|20","rows":[1,"a",20,["hob2","hob1"]]},{"column":"2|b|21","rows":[2,"b",21,["hob2","hob3"]]},{"column":"3|c|22","rows":[3,"c",22,["hob3","hob4"]]},{"column":"4|d|23","rows":[4,"d",23,["hob4","hob5"]]},{"column":"5|e|24","rows":[5,"e",24,["hob5","hob6"]]},{"column":"6|f|26","rows":[6,"f",26,["hob6","hob7"]]}]}]}`, + }, + }, + CreateTableStmt: "(_id String, id id, name String, age Int, hobbies StringSet)", + }, +} + +// TestKafkaRunner takes as input a slice of RunnerTest structs +// For each RunnerTest, TestRunner: +// 1. Creates a kafkaRunner based on a config file +// 2. Reads data from a data file +// 3. Encodes that data +// 4. Writes the data to kafka +// 5. Runs the kafkaRunner +// 6. Confirms that the data was written to FeatureBase as expected +func TestKafkaRunner(t *testing.T) { + for _, test := range kafkaRunnerTests { + + // define final path to config and data files + var subpath string + switch encode := test.Encode; encode { + case kafka.JSON: + subpath = kafka.JSON + case kafka.Avro: + subpath = kafka.Avro + default: + t.Fatalf("unsupported encoding type") + } + KafkaConfig := confPathPrefix + subpath + "/" + test.ConfigFile + DataFile := dataPathPrefix + subpath + "/" + test.DataFile + + // create command and kafka runner + fbsql := NewCommand(logger.StderrLogger) + setupKafkaRunnerCommand(t, fbsql) + + // read config file so we can build a table before inserting data + config, err := kafka.ConfigStructFromFile(KafkaConfig) + if err != nil { + t.Fatalf("building kafka.Config struct: %s", err) + } + + // create table to write to (kafka runner does not currently create tables) + createTable := strings.NewReader(fmt.Sprintf("CREATE TABLE %s %s;", config.Table, test.CreateTableStmt)) + _, err = fbsql.Queryer.Query(fbsql.organizationID, fbsql.databaseID, createTable) + if err != nil { + t.Fatalf("creating table: %s", err) + } + + // delete all tables after test (note this is after all tests not after each test) + defer func() { + dropTable := strings.NewReader(fmt.Sprintf("DROP TABLE %s;", config.Table)) + _, err = fbsql.Queryer.Query(fbsql.organizationID, fbsql.databaseID, dropTable) + if err != nil { + t.Errorf("dropping table: %s", err) + } + }() + + // create runner and hijack hosts based on test env (i.e. not from configuration file or other) + kafkaRunner, err := fbsql.newKafkaRunner(KafkaConfig) + if err != nil { + t.Fatalf("creating runner: %s", err) + } + kafkaRunner.FeaturebaseHosts = []string{pilosaHost} + kafkaRunner.SchemaRegistryURL = registryHost + kafkaRunner.KafkaBootstrapServers = []string{kafkaHost} + kafkaRunner.FeaturebaseGRPCHosts = []string{pilosaGrpcHost} + + // read data that will go to kafka + var records []map[string]interface{} + records, err = recordsFromFile(DataFile) + if err != nil { + t.Errorf("reading raw messages from data file: %s", err) + } + + // write data to kafka + var messages [][]byte + switch encode := test.Encode; encode { + case kafka.JSON: + messages, err = encodeJSONMessages(records) + case kafka.Avro: + t.Errorf("unsupported encoding type") + default: + t.Errorf("unsupported encoding type") + } + if err != nil { + t.Fatal(err) + } + + err = internal.ProduceMessages(kafkaRunner.Topics[0], strings.Join(kafkaRunner.KafkaBootstrapServers, ","), messages) + if err != nil { + t.Fatal(err) + } + defer func() { + // and delete the topic + err = internal.DeleteTopic(kafkaRunner.Topics[0], strings.Join(kafkaRunner.KafkaBootstrapServers, ",")) + if err != nil { + t.Fatal(err) + } + }() + + // Run it!! + kafkaRunner.MaxMsgs = uint64(len(messages)) + err = kafkaRunner.Run() + if err != nil { + t.Fatalf("running runner: %s", err) + } + + // and check the results... + for _, q := range test.Tests { + + client := &http.Client{} + var data = strings.NewReader(q.Query) + req, err := http.NewRequest("POST", fmt.Sprintf("http://localhost:10101/index/%s/query", kafkaRunner.Index), data) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + respBody, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatal(err) + } + require.JSONEq(t, q.ExpectedResp, string(respBody)) + } + + } +} + +// Reads lines from a file and builds an in memory structure This function +// expects the file formated as new line delimited JSON +func recordsFromFile(pathToRecords string) (records []map[string]interface{}, err error) { + var data map[string]interface{} + + recordsFile, err := os.Open(pathToRecords) + if err != nil { + return nil, fmt.Errorf("opening records file: %s", err) + } + defer recordsFile.Close() + + s := bufio.NewScanner(recordsFile) + for s.Scan() { + err := json.Unmarshal(s.Bytes(), &data) + if err != nil { + return nil, fmt.Errorf("unmarshal json: %s", err) + } + records = append(records, data) + data = make(map[string]interface{}) + } + return records, nil +} + +// Mainly emmulate what happens in cli/cli.go Command.run() Unfortunately we +// cannot just use that function because we need to create our own kafkaRunner +// and update hosts and ports based on the test environment +func setupKafkaRunnerCommand(t *testing.T, fbsql *Command) { + fbsql.Config.Host = strings.Split(pilosaHost, ":")[0] + fbsql.Config.Port = strings.Split(pilosaHost, ":")[1] + if err := fbsql.setupConfig(); err != nil { + t.Fatalf("setting up config: %s", err) + } + + // Check to see if Command needs to run in non-interactive mode. + if len(fbsql.Commands) > 0 || + len(fbsql.Files) > 0 || + fbsql.Config.KafkaConfig != "" || + fbsql.Config.CSV { + fbsql.nonInteractiveMode = true + } + + if err := fbsql.setupClient(); err != nil { + t.Fatalf("setting up client: %s", err) + } + + // Print the connection info. + if !fbsql.nonInteractiveMode { + fbsql.printConnInfo() + } + + if err := fbsql.connectToDatabase(fbsql.database); err != nil { + t.Fatalf("connecting to database: %s", err) + // We intentionally do not return err here. + } +} + +// Take a slice of JSON objects and convert them to a slice of byte slices that +// can be sent to and stored in kafka +func encodeJSONMessages(records []map[string]interface{}) ([][]byte, error) { + var messages [][]byte + for _, record := range records { + encoded, err := json.Marshal(record) + if err != nil { + return nil, errors.Errorf("marshaling records to JSON: %s", err) + } + messages = append(messages, encoded) + } + return messages, nil +} diff --git a/cli/internal/kafka_util.go b/cli/internal/kafka_util.go new file mode 100644 index 000000000..e07f8a5f2 --- /dev/null +++ b/cli/internal/kafka_util.go @@ -0,0 +1,74 @@ +package internal + +import ( + "context" + + "github.com/confluentinc/confluent-kafka-go/kafka" + "github.com/featurebasedb/featurebase/v3/errors" +) + +func ProduceMessages(topic string, bootstrapServers string, messages [][]byte) error { + + // create producer, defer closing + p, err := kafka.NewProducer(&kafka.ConfigMap{"bootstrap.servers": bootstrapServers}) + if err != nil { + return errors.Errorf("creating producer: %s", err) + } + defer p.Close() + + // create admin client needed to create topic, defer closing + ac, err := kafka.NewAdminClient(&kafka.ConfigMap{"bootstrap.servers": bootstrapServers}) + if err != nil { + return errors.Errorf("creating admin client: %s", err) + } + defer ac.Close() + + // create a topic + var ts = []kafka.TopicSpecification{ + { + Topic: topic, + NumPartitions: 1, + ReplicationFactor: 1, + }, + } + ac.CreateTopics(context.Background(), ts, nil) + // caller must delete if needed + + for _, message := range messages { + err := p.Produce(&kafka.Message{ + TopicPartition: kafka.TopicPartition{Topic: &topic, Partition: kafka.PartitionAny}, + Value: message, + }, nil) + if err != nil { + return errors.Errorf("producing messages: %s", err) + } + p.Flush(3000) + } + + return nil + +} + +func DeleteTopic(topic string, bootstrapServers string) error { + // create admin client needed to create topic, defer closing + ac, err := kafka.NewAdminClient(&kafka.ConfigMap{"bootstrap.servers": bootstrapServers}) + if err != nil { + return errors.Errorf("creating admin client: %s", err) + } + defer ac.Close() + + // delete a topic, catch any errors + results, err := ac.DeleteTopics(context.Background(), []string{topic}, nil) + if err != nil { + return errors.Errorf("deleting topic: %s", err) + } + + for _, result := range results { + if result.Error.Code() != kafka.ErrNoError { + return errors.Errorf("fatal error deleting topic: %s", result.String()) + } + } + + return nil + +} diff --git a/cli/kafka.go b/cli/kafka.go index 9908dd56d..2cfd0b7aa 100644 --- a/cli/kafka.go +++ b/cli/kafka.go @@ -1,35 +1,22 @@ package cli import ( - "fmt" - "github.com/featurebasedb/featurebase/v3/cli/batch" "github.com/featurebasedb/featurebase/v3/cli/kafka" "github.com/featurebasedb/featurebase/v3/errors" - "github.com/spf13/viper" ) func (cmd *Command) newKafkaRunner(cfgFile string) (*kafka.Runner, error) { - // Read the kafka config file. - v := viper.New() - v.SetConfigFile(cfgFile) - v.SetConfigType("toml") - err := v.ReadInConfig() - if err != nil { - return nil, fmt.Errorf("error reading configuration file '%s': %v", cfgFile, err) - } - cfg := kafka.Config{} - if err := v.Unmarshal(&cfg); err != nil { - return nil, errors.Wrap(err, "unmarshalling config") + cfg, err := kafka.ConfigStructFromFile(cfgFile) + if err != nil { + return nil, err } if err := kafka.ValidateConfig(cfg); err != nil { return nil, errors.Wrap(err, "validating config") } - // Create a new config with defaults. - // Look up fields based on table provided in the config. wqr, err := cmd.executeQuery(newRawQuery("SHOW COLUMNS FROM " + cfg.Table)) if err != nil { @@ -57,7 +44,7 @@ func (cmd *Command) newKafkaRunner(cfgFile string) (*kafka.Runner, error) { return nil, errors.Wrap(err, "cleaning config") } - flds, err := kafka.ConfigToFields(cfg) + flds, err := kafka.ConfigToFields(cfg, idkCfg.PrimaryKeys) if err != nil { return nil, errors.Wrap(err, "getting fields from config") } diff --git a/cli/kafka/config.go b/cli/kafka/config.go index 444d0642f..25390e37f 100644 --- a/cli/kafka/config.go +++ b/cli/kafka/config.go @@ -8,6 +8,14 @@ import ( "github.com/featurebasedb/featurebase/v3/dax" "github.com/featurebasedb/featurebase/v3/idk" "github.com/pkg/errors" + "github.com/spf13/viper" +) + +// Kafka message encoding types supported - changes here should be +// propogated to the Config struct and ValidateConfig error message +const ( + JSON string = "json" + Avro string = "avro" ) // Config is the user-facing configuration for kafka support in the CLI. This is @@ -23,6 +31,8 @@ type Config struct { Table string `mapstructure:"table" help:"Destination table name."` Fields []Field `mapstructure:"fields"` + + Encode string `mapstructure:"encode" help:"Encoding format (currently supported formats: avro, json)"` } // Field is a user-facing configuration field. @@ -44,18 +54,73 @@ type ConfigForIDK struct { BatchMaxStaleness time.Duration Timeout time.Duration - Table string - IDField string - Fields []idk.RawField + Table string + IDField string + PrimaryKeys []string + Fields []idk.RawField + + Encode string +} + +// Generate a Config struct based on a configuration file +// Default values for Config struct are defined here +func ConfigStructFromFile(cfgFile string) (cfg Config, err error) { + // configure viper + v := viper.New() + v.SetConfigFile(cfgFile) + v.SetConfigType("toml") + + // set defaults + v.SetDefault("hosts", []string{"localhost:9092"}) + v.SetDefault("group", "default-featurebase-group") + v.SetDefault("batch-size", 1) + v.SetDefault("batch-max-staleness", 5*time.Second) + v.SetDefault("timeout", 5*time.Second) + v.SetDefault("encode", JSON) + + // Read the kafka config file. + err = v.ReadInConfig() + if err != nil { + return cfg, fmt.Errorf("error reading configuration file '%s': %v", cfgFile, err) + } + + if err := v.Unmarshal(&cfg); err != nil { + return cfg, errors.Wrap(err, "unmarshalling config") + } + + return + } // ValidateConfig validates the config is usable. +// Note that different encoding methods require +// different configurations func ValidateConfig(c Config) error { + + // validate common if c.Table == "" { return errors.Errorf("table is required") - } else if len(c.Topics) == 0 { + } + + if len(c.Topics) == 0 { return errors.Errorf("at least one topic is required") - } else if len(c.Fields) > 0 { + } + + // validate on a by encoding basis + switch c.Encode { + case JSON: + return validateConfigJSON(c) + case Avro: + return validateConfigAvro(c) + } + + return nil + +} + +func validateConfigJSON(c Config) error { + + if len(c.Fields) > 0 { // We only need to do these checks if any fields are specified at all. // If no fields are specified, that's ok because then we default to // using fields based off the existing table. @@ -70,46 +135,64 @@ func ValidateConfig(c Config) error { if c.Fields[i].Name == "" { return errors.Errorf("a name attribute (which isn't equal to \"\") should exist for all fields") } + if c.Fields[i].SourceType == "" { + return errors.Errorf("a source-type attribute (which isn't equal to \"\") should exist for all fields") + } } - if found != 1 { - return errors.Errorf("exactly one primary key field is required") + if found < 1 { + return errors.Errorf("at least one primary key field is required") } } } + + return nil +} + +// Only primary key fields required +func validateConfigAvro(c Config) error { return nil } // ConvertConfig converts a Config to one that suitable for IDK. func ConvertConfig(c Config) (ConfigForIDK, error) { - // Set a default kafka host in case one isn't provided. - hosts := []string{"localhost:9092"} - if len(c.Hosts) > 0 { - hosts = c.Hosts - } // Copy all the shared members from Config to ConfigForIDK. out := ConfigForIDK{ - Hosts: hosts, + Hosts: c.Hosts, Group: c.Group, Topics: c.Topics, BatchSize: c.BatchSize, BatchMaxStaleness: c.BatchMaxStaleness, Timeout: c.Timeout, Table: c.Table, + Encode: c.Encode, } if len(c.Fields) == 0 { return out, errors.New("fields cannot be empty") } - // rawFields wil be the same as c.Fields, but possibly enhanced. + // rawFields will be the same as c.Fields, but possibly enhanced. rawFields := make([]idk.RawField, 0, len(c.Fields)) - var foundPK bool + var stringKeys bool + primaryKeys := []string{} for _, fld := range c.Fields { + + // handle primary key if fld.PrimaryKey { - out.IDField = fld.Name - foundPK = true + switch keyType := fld.SourceType; keyType { + case dax.BaseTypeID: + // no-op + case dax.BaseTypeIDSet, dax.BaseTypeStringSet, dax.BaseTypeIDSetQ, dax.BaseTypeStringSetQ: + // key should be field that cannot contain multiple values + return out, errors.Errorf("Invalid") + default: + // unless + stringKeys = true + + } + primaryKeys = append(primaryKeys, fld.Name) } typ, quals, err := dax.SplitFieldType(fld.SourceType) @@ -149,8 +232,16 @@ func ConvertConfig(c Config) (ConfigForIDK, error) { rawFields = append(rawFields, rawFld) } - if !foundPK { + + // Should have at least one primary key. If there is more than one primary key OR + // using a string field as they key then use string keys (i.e. table will be keyed) + // Else, use ids (i.e. table will not be keyed) + if len(primaryKeys) < 1 { return out, errors.New("primary-key not found in fields") + } else if stringKeys || len(primaryKeys) > 1 { + out.PrimaryKeys = primaryKeys + } else { + out.IDField = primaryKeys[0] } out.Fields = rawFields @@ -160,13 +251,17 @@ func ConvertConfig(c Config) (ConfigForIDK, error) { // ConfigToFields returns a list of *dax.Field based on the IDField and Fields // in the Config. -func ConfigToFields(c Config) ([]*dax.Field, error) { +func ConfigToFields(c Config, primaryKeys []string) ([]*dax.Field, error) { // We don't know if a primary key will be found, so we can't set the // capacity to `len(c.Fields)-1`. out := make([]*dax.Field, 0, len(c.Fields)) for _, fld := range c.Fields { - if fld.PrimaryKey { + // When we have a single primary key, don't also store that value as a + // field in FeatureBase. However, when we have more than one primary + // key, store all the values used for the compound key as fields in + // FeatureBase. + if fld.PrimaryKey && len(primaryKeys) < 2 { continue } typ, quals, err := dax.SplitFieldType(fld.SourceType) diff --git a/cli/kafka/config_test.go b/cli/kafka/config_test.go new file mode 100644 index 000000000..2bea16853 --- /dev/null +++ b/cli/kafka/config_test.go @@ -0,0 +1,39 @@ +package kafka + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +type configStructFromFileTest struct { + configFilePath string + expectedStruct string +} + +func TestConfigStructFromFile(t *testing.T) { + tests := []configStructFromFileTest{ + { // confirm struct fields are being set in general + configFilePath: "./config_test_data/config01.toml", + expectedStruct: `{"Hosts":["kafka:9090"],"Group":"testGroup","Topics":["testTopic"],"BatchSize":35,"BatchMaxStaleness":25000000000,"Timeout":16000000000,"Table":"testTable","Fields":[{"Name":"id","SourceType":"id","SourcePath":null,"PrimaryKey":true},{"Name":"name","SourceType":"string","SourcePath":["test","path"],"PrimaryKey":false},{"Name":"age","SourceType":"int","SourcePath":null,"PrimaryKey":false},{"Name":"hobbies","SourceType":"stringset","SourcePath":null,"PrimaryKey":false}],"Encode":"json"}`, + }, + { // confirm defaults are being set + configFilePath: "./config_test_data/config00.toml", + expectedStruct: `{"Hosts":["localhost:9092"],"Group":"default-featurebase-group","Topics":["topic00"],"BatchSize":1,"BatchMaxStaleness":5000000000,"Timeout":5000000000,"Table":"table00","Fields":[{"Name":"id","SourceType":"string","SourcePath":null,"PrimaryKey":true},{"Name":"name","SourceType":"string","SourcePath":null,"PrimaryKey":false},{"Name":"age","SourceType":"int","SourcePath":null,"PrimaryKey":false},{"Name":"hobbies","SourceType":"stringset","SourcePath":null,"PrimaryKey":false}],"Encode":"json"}`, + }, + } + + for _, test := range tests { + cfg, err := ConfigStructFromFile(test.configFilePath) + if err != nil { + t.Fatalf("getting config struct from file: %s", err) + } + json, err := json.Marshal(cfg) + if err != nil { + t.Fatalf("marshaling config struct: %s", err) + } + require.JSONEq(t, test.expectedStruct, string(json)) + } + +} diff --git a/cli/kafka/config_test_data/config00.toml b/cli/kafka/config_test_data/config00.toml new file mode 100644 index 000000000..ccdb572af --- /dev/null +++ b/cli/kafka/config_test_data/config00.toml @@ -0,0 +1,20 @@ +## confirm defaults are being set correctly +topics = "topic00" +table = "table00" + +[[fields]] +name = "id" +source-type = "string" +primary-key = true + +[[fields]] +name = "name" +source-type = "string" + +[[fields]] +name = "age" +source-type = "int" + +[[fields]] +name = "hobbies" +source-type = "stringset" diff --git a/cli/kafka/config_test_data/config01.toml b/cli/kafka/config_test_data/config01.toml new file mode 100644 index 000000000..827177b03 --- /dev/null +++ b/cli/kafka/config_test_data/config01.toml @@ -0,0 +1,26 @@ +hosts = ["kafka:9090"] +group = "testGroup" +topics = "testTopic" +table = "testTable" +batch-size = 35 +batch-max-staleness = "25s" +timeout = "16s" +encode = "json" + +[[fields]] +name = "id" +source-type = "id" +primary-key = true + +[[fields]] +name = "name" +source-type = "string" +source-path = ["test", "path"] + +[[fields]] +name = "age" +source-type = "int" + +[[fields]] +name = "hobbies" +source-type = "stringset" diff --git a/cli/kafka/runner.go b/cli/kafka/runner.go index cacde76bb..7b4e2e83f 100644 --- a/cli/kafka/runner.go +++ b/cli/kafka/runner.go @@ -7,7 +7,9 @@ import ( fbbatch "github.com/featurebasedb/featurebase/v3/batch" "github.com/featurebasedb/featurebase/v3/errors" "github.com/featurebasedb/featurebase/v3/idk" - "github.com/featurebasedb/featurebase/v3/idk/kafka_static" + "github.com/featurebasedb/featurebase/v3/idk/common" + "github.com/featurebasedb/featurebase/v3/idk/kafka" + "github.com/featurebasedb/featurebase/v3/idk/kafka_sasl" "github.com/featurebasedb/featurebase/v3/logger" ) @@ -15,53 +17,90 @@ import ( // idk.kafka_static.Main in that it embeds idk.Main and contains additional // functionality specific to its use case. type Runner struct { - idk.Main `flag:"!embed"` - KafkaHosts []string `help:"Comma separated list of host:port pairs for Kafka."` - Group string `help:"Kafka group."` - Topics []string `help:"Kafka topics to read from."` - Timeout time.Duration `help:"Time to wait for more records from Kafka before flushing a batch. 0 to disable."` - Header []idk.RawField `help:"Header configuration."` + idk.Main `flag:"!embed"` + idk.ConfluentCommand `flag:"!embed"` + Hosts []string `help:"Comma separated list of host:port pairs for Kafka."` + Group string `help:"Kafka group."` + Topics []string `help:"Kafka topics to read from."` + Timeout time.Duration `help:"Time to wait for more records from Kafka before flushing a batch. 0 to disable."` + Header []idk.RawField `help:"Header configuration."` } +// Configure and return *runner with common elements of all kafka runners func NewRunner(cfg ConfigForIDK, batcher fbbatch.Batcher, logWriter io.Writer) *Runner { idkMain := idk.NewMain() idkMain.IDField = cfg.IDField + idkMain.PrimaryKeyFields = cfg.PrimaryKeys idkMain.Index = cfg.Table idkMain.Batcher = batcher idkMain.BatchSize = cfg.BatchSize idkMain.BatchMaxStaleness = cfg.BatchMaxStaleness idkMain.SetBasic() idkMain.SetLog(logger.NewStandardLogger(logWriter)) + idkMain.OffsetMode = true + idkMain.Namespace = "sql_kafka_runner" + idkMain.Pprof = "" // don't initialize pprof until we actually use it in tests kr := &Runner{ - Main: *idkMain, - KafkaHosts: cfg.Hosts, - Group: cfg.Group, - Topics: cfg.Topics, - Header: cfg.Fields, - Timeout: cfg.Timeout, + Main: *idkMain, + Hosts: cfg.Hosts, + Group: cfg.Group, + Topics: cfg.Topics, + Header: cfg.Fields, + Timeout: cfg.Timeout, } - kr.OffsetMode = true - kr.Main.Namespace = "cli_kafka_runner" - kr.Main.Pprof = "" // don't initialize pprof until we actually use it in tests - kr.NewSource = func() (idk.Source, error) { - source := kafka_static.NewSource() - source.Hosts = kr.KafkaHosts - source.Group = kr.Group - source.Topics = kr.Topics - source.Log = kr.Main.Log() - // source.TLS = m.KafkaTLS - source.Timeout = kr.Timeout - // source.SkipOld = m.SkipOld - source.HeaderFields = kr.Header - // source.S3Region = m.S3Region - // source.AllowMissingFields = m.AllowMissingFields - err := source.Open() + // NewSource should be set based on the encoding of the source (e.g. JSON, Avro) + if cfg.Encode == Avro { + kr.GetAvroNewSource() + } else if cfg.Encode == JSON { + kr.GetJSONNewSource() + } + + return kr +} + +func (r *Runner) GetJSONNewSource() { + r.NewSource = func() (idk.Source, error) { + source := kafka_sasl.NewSource() + source.KafkaBootstrapServers = r.Hosts + source.Group = r.Group + source.Topics = r.Topics + source.Log = r.Main.Log() + source.Timeout = r.Timeout + source.HeaderFields = r.Header + cfg, err := common.SetupConfluent(&r.ConfluentCommand) + if err != nil { + return nil, err + } + source.ConfigMap = cfg + + err = source.Open() + if err != nil { + return nil, errors.Wrap(err, "opening source") + } + return source, nil + } +} + +func (r *Runner) GetAvroNewSource() { + r.NewSource = func() (idk.Source, error) { + source := kafka.NewSource() + source.KafkaBootstrapServers = r.Hosts + source.Group = r.Group + source.Topics = r.Topics + source.Log = r.Main.Log() + source.Timeout = r.Timeout + cfg, err := common.SetupConfluent(&r.ConfluentCommand) + if err != nil { + return nil, err + } + source.ConfigMap = cfg + + err = source.Open() if err != nil { return nil, errors.Wrap(err, "opening source") } return source, nil } - return kr } diff --git a/cli/kafka/runner_test_data/config/json/config00.toml b/cli/kafka/runner_test_data/config/json/config00.toml new file mode 100644 index 000000000..472987fe2 --- /dev/null +++ b/cli/kafka/runner_test_data/config/json/config00.toml @@ -0,0 +1,25 @@ +hosts = ["localhost:9092"] +group = "grp" +topics = "topic00" +table = "table00" +batch-size = 1 +batch-max-staleness = "5s" +timeout = "5s" +encode = "json" + +[[fields]] +name = "id" +source-type = "id" +primary-key = true + +[[fields]] +name = "name" +source-type = "string" + +[[fields]] +name = "age" +source-type = "int" + +[[fields]] +name = "hobbies" +source-type = "stringset" diff --git a/cli/kafka/runner_test_data/config/json/config01.toml b/cli/kafka/runner_test_data/config/json/config01.toml new file mode 100644 index 000000000..737012cb3 --- /dev/null +++ b/cli/kafka/runner_test_data/config/json/config01.toml @@ -0,0 +1,25 @@ +hosts = ["localhost:9092"] +group = "grp" +topics = "topic01" +table = "table01" +batch-size = 1 +batch-max-staleness = "5s" +timeout = "5s" +encode = "json" + +[[fields]] +name = "id" +source-type = "string" +primary-key = true + +[[fields]] +name = "name" +source-type = "string" + +[[fields]] +name = "age" +source-type = "int" + +[[fields]] +name = "hobbies" +source-type = "stringset" diff --git a/cli/kafka/runner_test_data/config/json/config02.toml b/cli/kafka/runner_test_data/config/json/config02.toml new file mode 100644 index 000000000..5f0b687f1 --- /dev/null +++ b/cli/kafka/runner_test_data/config/json/config02.toml @@ -0,0 +1,26 @@ +hosts = ["localhost:9092"] +group = "grp" +topics = "topic02" +table = "table02" +batch-size = 1 +batch-max-staleness = "5s" +timeout = "5s" +encode = "json" + +[[fields]] +name = "id" +source-type = "string" +primary-key = true + +[[fields]] +name = "name" +source-type = "string" +primary-key = true + +[[fields]] +name = "age" +source-type = "int" + +[[fields]] +name = "hobbies" +source-type = "stringset" \ No newline at end of file diff --git a/cli/kafka/runner_test_data/config/json/config03.toml b/cli/kafka/runner_test_data/config/json/config03.toml new file mode 100644 index 000000000..1fd9e08b7 --- /dev/null +++ b/cli/kafka/runner_test_data/config/json/config03.toml @@ -0,0 +1,27 @@ +hosts = ["localhost:9092"] +group = "grp" +topics = "topic03" +table = "table03" +batch-size = 1 +batch-max-staleness = "5s" +timeout = "5s" +encode = "json" + +[[fields]] +name = "id" +source-type = "id" +primary-key = true + +[[fields]] +name = "name" +source-type = "string" +primary-key = true + +[[fields]] +name = "age" +source-type = "int" +primary-key = true + +[[fields]] +name = "hobbies" +source-type = "stringset" \ No newline at end of file diff --git a/cli/kafka/runner_test_data/data/json/data00.json b/cli/kafka/runner_test_data/data/json/data00.json new file mode 100644 index 000000000..c6e720ddc --- /dev/null +++ b/cli/kafka/runner_test_data/data/json/data00.json @@ -0,0 +1,6 @@ +{"id": 1, "name": "a", "age": 20, "hobbies": "hob1,hob2"} +{"id": 2, "name": "b", "age": 21, "hobbies": "hob3,hob2"} +{"id": 3, "name": "c", "age": 22, "hobbies": "hob3,hob4"} +{"id": 4, "name": "d", "age": 23, "hobbies": "hob4,hob5"} +{"id": 5, "name": "e", "age": 24, "hobbies": "hob5,hob6"} +{"id": 6, "name": "f", "age": 26, "hobbies": "hob6,hob7"} \ No newline at end of file diff --git a/idk/kafka_sasl/source.go b/idk/kafka_sasl/source.go index 8a27c7305..de9b0967f 100644 --- a/idk/kafka_sasl/source.go +++ b/idk/kafka_sasl/source.go @@ -26,9 +26,16 @@ type Source struct { Log logger.Logger Timeout time.Duration SkipOld bool - Header string AllowMissingFields bool + // Header is a file referencing a file containing JSON header configuration. + Header string + + // HeaderFields can be provided instead of Header. It is a slice of + // RawFields which will be marshalled and parsed the same way a JSON object + // in Header would be. It is used only if a Header is not provided. + HeaderFields []idk.RawField + schema []idk.Field paths idk.PathTable @@ -193,14 +200,24 @@ func (r *Record) Data() []interface{} { // Open initializes the kafka source. func (s *Source) Open() error { - if len(s.Header) == 0 { - return errors.New("needs header specification file") + if len(s.Header) == 0 && len(s.HeaderFields) == 0 { + return errors.New("needs header specification file (file or fields)") } - headerData, err := os.ReadFile(s.Header) - if err != nil { - return errors.Wrap(err, "reading header file") + var headerData []byte + var err error + if s.Header != "" { + headerData, err = os.ReadFile(s.Header) + if err != nil { + return errors.Wrap(err, "reading header file") + } + } else { + headerData, err = json.Marshal(s.HeaderFields) + if err != nil { + return errors.Wrap(err, "marshalling header fields") + } } + schema, paths, err := idk.ParseHeader(headerData) if err != nil { return errors.Wrap(err, "processing header")