From aefbc6493bf223986c3caa99238a5d8f453e0f40 Mon Sep 17 00:00:00 2001 From: jacob Date: Mon, 3 Apr 2023 11:15:32 -0500 Subject: [PATCH] cli/kafka updates --- cli/batch/sql.go | 27 +- cli/cli_kafka_integration_test.go | 488 ++++++++++++++++++ cli/cli_kafka_runner_test.go | 343 ------------ cli/internal/kafka_util.go | 74 --- cli/kafka.go | 2 +- cli/kafka/config.go | 84 +-- cli/kafka/config_test.go | 26 +- cli/kafka/runner.go | 24 +- .../config}/config00.toml | 0 .../config}/config01.toml | 0 cli/kafka/testdata/config/config02.toml | 27 + cli/kafka/testdata/config/config03.toml | 30 ++ .../runner/config}/config00.toml | 3 +- .../runner/config}/config01.toml | 3 +- .../runner/config}/config02.toml | 3 +- .../runner/config}/config03.toml | 3 +- .../testdata/runner/config/config04.toml | 15 + .../testdata/runner/config/config05.toml | 27 + .../json => testdata/runner/data}/data00.json | 2 +- cli/kafka/testdata/runner/data/data01.json | 10 + cli/kafka/testdata/runner/data/data02.json | 6 + .../testdata/runner/schema/schema01.json | 32 ++ idk/ingest.go | 24 +- idk/kafka/cmd.go | 7 +- idk/kafka/source.go | 75 +-- idk/kafka_sasl/source.go | 123 ++++- wire_response.go | 23 + 27 files changed, 926 insertions(+), 555 deletions(-) create mode 100644 cli/cli_kafka_integration_test.go delete mode 100644 cli/cli_kafka_runner_test.go delete mode 100644 cli/internal/kafka_util.go rename cli/kafka/{config_test_data => testdata/config}/config00.toml (100%) rename cli/kafka/{config_test_data => testdata/config}/config01.toml (100%) create mode 100644 cli/kafka/testdata/config/config02.toml create mode 100644 cli/kafka/testdata/config/config03.toml rename cli/kafka/{runner_test_data/config/json => testdata/runner/config}/config00.toml (86%) rename cli/kafka/{runner_test_data/config/json => testdata/runner/config}/config01.toml (86%) rename cli/kafka/{runner_test_data/config/json => testdata/runner/config}/config02.toml (86%) rename cli/kafka/{runner_test_data/config/json => testdata/runner/config}/config03.toml (87%) create mode 100644 cli/kafka/testdata/runner/config/config04.toml create mode 100644 cli/kafka/testdata/runner/config/config05.toml rename cli/kafka/{runner_test_data/data/json => testdata/runner/data}/data00.json (83%) create mode 100644 cli/kafka/testdata/runner/data/data01.json create mode 100644 cli/kafka/testdata/runner/data/data02.json create mode 100644 cli/kafka/testdata/runner/schema/schema01.json diff --git a/cli/batch/sql.go b/cli/batch/sql.go index aadcb199e..5e4e870f6 100644 --- a/cli/batch/sql.go +++ b/cli/batch/sql.go @@ -176,23 +176,10 @@ func buildBulkInsert(tbl *dax.Table, fields []*dax.Field, ids []interface{}, row // MAP - // 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 - } + keyType := dax.BaseTypeID + if tbl.StringKeys() { + keyType = dax.BaseTypeString } sb.WriteString(`) MAP ('$._id' `) @@ -210,7 +197,13 @@ func buildBulkInsert(tbl *dax.Table, fields []*dax.Field, ids []interface{}, row for i := range rows { // Write the ID value. - m[string(dax.PrimaryKeyFieldName)] = fmt.Sprintf(fmtStr, ids[i]) + if keyType == dax.BaseTypeID { + m[string(dax.PrimaryKeyFieldName)] = ids[i] + } else { + // ids for key index can be string or []byte + m[string(dax.PrimaryKeyFieldName)] = fmt.Sprintf("%s", 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_integration_test.go b/cli/cli_kafka_integration_test.go new file mode 100644 index 000000000..a05314811 --- /dev/null +++ b/cli/cli_kafka_integration_test.go @@ -0,0 +1,488 @@ +package cli_test + +import ( + "bufio" + "context" + "encoding/binary" + "encoding/json" + "fmt" + "io/ioutil" + "os" + "strconv" + "strings" + "testing" + + featurebase "github.com/featurebasedb/featurebase/v3" + "github.com/featurebasedb/featurebase/v3/cli" + "github.com/featurebasedb/featurebase/v3/cli/kafka" + "github.com/featurebasedb/featurebase/v3/errors" + "github.com/featurebasedb/featurebase/v3/idk/kafka/csrc" + "github.com/featurebasedb/featurebase/v3/logger" + "github.com/stretchr/testify/require" + + confluent "github.com/confluentinc/confluent-kafka-go/kafka" + + avro "github.com/linkedin/goavro/v2" +) + +// Struct that contains the services required to run the kafka runner tests +type KafkaRunnerTestServices struct { + featurebaseHost string + featurebaseGRPCHost string + kafkaHost string + registryHost string +} + +func envOr(envName, defaultVal string) string { + if val, ok := os.LookupEnv(envName); !ok { + return defaultVal + } else { + return val + } +} + +func getKafkaRunnerTestServices() *KafkaRunnerTestServices { + return &KafkaRunnerTestServices{ + featurebaseHost: envOr("KAFKA_RUNNER_TEST_FEATUREBASE_HOST", "localhost:10101"), + featurebaseGRPCHost: envOr("KAFKA_RUNNER_TEST_FEATUREBASEGRPC_HOST", "localhost:20101"), + kafkaHost: envOr("KAFKA_RUNNER_TEST_KAFKA_HOST", "localhost:9092"), + registryHost: envOr("KAFKA_RUNNER_TEST_REGISTRY_HOST", "localhost:8081"), + } +} + +// 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 + CreateTableStmt string // statement used to create table prior to ingest + SchemaFile string + 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 +} + +// A slice of KafkaRunnerTest structs that will be used in TestKafkaRunner test +// function. +var kafkaRunnerTests = []kafkaRunnerTest{ + { // id keys + ConfigFile: "config00.toml", + DataFile: "data00.json", + Tests: []testQuery{ + { + Query: "select * from order by name", + ExpectedResp: `[[1,"a",20,["hob1","hob2"]],[2,"b",21,["hob2","hob3"]],[3,"c",22,["hob3","hob4"]],[4,"d",23,["hob4","hob5"]],[5,"e",24,["hob5","hob6"]],[6,"f",26,["hob6","hob7"]]]`, + }, + }, + CreateTableStmt: "(_id ID, name String, age Int, hobbies StringSet)", + }, + { // string keys + ConfigFile: "config01.toml", + DataFile: "data00.json", + Tests: []testQuery{ + { + Query: "select * from
order by name", + ExpectedResp: `[["1","a",20,["hob1","hob2"]],["2","b",21,["hob2","hob3"]],["3","c",22,["hob3","hob4"]],["4","d",23,["hob4","hob5"]],["5","e",24,["hob5","hob6"]],["6","f",26,["hob6","hob7"]]]`, + }, + }, + CreateTableStmt: "(_id String, name String, age Int, hobbies StringSet)", + }, + { // two string keys + ConfigFile: "config02.toml", + DataFile: "data00.json", + Tests: []testQuery{ + { + Query: "select * from
order by name", + ExpectedResp: `[["1|a","1","a",20,["hob1","hob2"]],["2|b","2","b",21,["hob2","hob3"]],["3|c","3","c",22,["hob3","hob4"]],["4|d","4","d",23,["hob4","hob5"]],["5|e","5","e",24,["hob5","hob6"]],["6|f","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", + Tests: []testQuery{ + { + Query: "select * from
order by name", + ExpectedResp: `[["1|a|20",1,"a",20,["hob1","hob2"]],["2|b|21",2,"b",21,["hob2","hob3"]],["3|c|22",3,"c",22,["hob3","hob4"]],["4|d|23",4,"d",23,["hob4","hob5"]],["5|e|24",5,"e",24,["hob5","hob6"]],["6|f|26",6,"f",26,["hob6","hob7"]]]`, + }, + }, + CreateTableStmt: "(_id String, id id, name String, age Int, hobbies StringSet)", + }, + { // missing values + ConfigFile: "config05.toml", + DataFile: "data02.json", + Tests: []testQuery{ + { + Query: "select * from
order by name", + ExpectedResp: `[["3",null,22,["hob3","hob4"]],["1","a",20,["hob1","hob2"]],["2","b",21,["hob2","hob3"]],["4","d",null,["hob4","hob5"]],["5","e",24,null],["6","f",26,["hob6","hob7"]]]`, + }, + }, + CreateTableStmt: "(_id String, name String, age Int, hobbies StringSet)", + }, + /*{ // id keys + ConfigFile: "config04.toml", + DataFile: "data01.json", + Encode: "avro", + SchemaFile: "schema01.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 string, string_string string, string_bytes string, pk2 stringset, stringset_bytes stringset, idset_long idset, decimal_double decimal(2), timestamp_bytes_ts timestamp, idset_longarray idset, dateint_bytes_ts int, bools stringset, stringset_string stringset, stringset_stringarray stringset, idset_int idset, timestamp_bytes_int timestamp, int_long int, id_long id, id_int id, idset_intarray idset, decimal_float decimal(2), bools-exists stringset, stringset_bytesarray stringset, int_int int, decimal_bytes decimal(2), pk1 stringset)", + },*/ +} + +// TestKafkaRunner takes as input a slice of kafkaRunnerTest structs. +// For each kafkaRunnerTest, TestKafkaRunner: +// 1. Creates and configues a new cli.Command +// 2. Reads kafka messages from a data file +// 3. Encodes that data +// 4. Writes the data to kafka +// 5. Runs the cli.Command +// 6. Confirms that the data was written to FeatureBase as expected +func TestKafkaRunner(t *testing.T) { + if testing.Short() { + t.Skip("skipping integration test") + } + + // Get host and port pair for services required for test (e.g. kafka and + // featurebase) + services := getKafkaRunnerTestServices() + + for _, test := range kafkaRunnerTests { + + // define path to config and data files + kafkaConfig := "./kafka/testdata/runner/config/" + test.ConfigFile + dataFile := "./kafka/testdata/runner/data/" + test.DataFile + schemaFile := "./kafka/testdata/runner/schema/" + test.SchemaFile + + // 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) + } + + // copy config file and replace values as defined in findAndReplace + var findAndReplace = map[string]string{ + "KAFKA_SERVICE": services.kafkaHost, + "SCHEMA_REGISTRY_SERVICE": services.registryHost, + "MAX_MESSAGES": strconv.Itoa(len(records)), + } + if err := createTempFindAndReplace(kafkaConfig, findAndReplace); err != nil { + t.Fatalf("creating temp config file: %s", err) + } + defer os.Remove(kafkaConfig + ".tmp") + + // create new command + fbsql := cli.NewCommand(logger.StderrLogger) + + fbsql.Config.Host = strings.Split(services.featurebaseHost, ":")[0] + fbsql.Config.Port = strings.Split(services.featurebaseHost, ":")[1] + fbsql.Run(context.Background()) // creates fbsql's Queryer which we can then use below + fbsql.Config.KafkaConfig = kafkaConfig + ".tmp" // when fbsql comes back, add kafka config + + config, err := kafka.ConfigFromFile(kafkaConfig + ".tmp") + if err != nil { + t.Fatal(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)) + wqr, err := fbsql.Queryer.Query("", "", createTable) + if err != nil { + t.Fatalf("creating table: %s", err) + } + + if wqr.Error != "" { + t.Error(wqr.Error) + } + + // 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("", "", dropTable) + if err != nil { + t.Errorf("dropping table: %s", err) + } + }() + + // write data to kafka + var messages [][]byte + switch encode := config.Encode; encode { + case "json": + messages, err = encodeJSONMessages(records) + if err != nil { + t.Fatal(err) + } + case "avro": + // get avro schema as string + schemaBytes, err := os.ReadFile(schemaFile) + if err != nil { + t.Fatal(err) + } + schema := string(schemaBytes) + + // post the schema to schema registry + schemaID, err := postSchema(schema, "kafka-runner-subject", services.registryHost) + if err != nil { + t.Fatal(err) + } + + // encode avro messages + messages, err = encodeAvroMessages(records, schema, schemaID) + if err != nil { + t.Fatal(err) + } + default: + t.Errorf("unsupported encoding type") + } + + if err = createTopic(config.Topics[0], services.kafkaHost); err != nil { + t.Fatal(err) + } + + if err = produceMessages(config.Topics[0], services.kafkaHost, messages); err != nil { + t.Fatal(err) + } + defer func() { + if err = deleteTopic(config.Topics[0], services.kafkaHost); err != nil { + t.Fatal(err) + } + }() + + if err = fbsql.Run(context.Background()); err != nil { + t.Fatal(err) + } + + // and check the results... + for _, q := range test.Tests { + query := strings.Replace(q.Query, "
", config.Table, 1) + if resp, err := fbsql.Queryer.Query("", "", strings.NewReader(query)); err != nil { + t.Fatal(err) + } else { + verifyQueryReponse(t, resp, q.ExpectedResp) + } + } + } +} + +// createTempFindAndReplace reads a file (source), replaces substrings specified +// in mapping, and writes the output to the same path as souce, appending +// ".tmp". +func createTempFindAndReplace(source string, mapping map[string]string) error { + //Read all the contents of the original file + bytesRead, err := ioutil.ReadFile(source) + if err != nil { + return errors.Errorf("%s", err) + } + + stringRead := string(bytesRead) + + for find, replace := range mapping { + stringRead = strings.Replace(stringRead, find, replace, 1) + } + + //Copy all the contents to the desitination file + if err = ioutil.WriteFile(source+".tmp", []byte(stringRead), 0755); err != nil { + return errors.Errorf("%s", err) + } + + return nil + +} + +// verifyQueryResponse compares the JSON binary representation of the Data field +// in an WireQueryResponse against some expected string. All lists in the Data +// field of the WireQueryResponse are sorted and then serialized. The test that +// calls this function fails if the query responses are different. In this case, +// the diff is displayed. +func verifyQueryReponse(t *testing.T, wqr *featurebase.WireQueryResponse, expectedQuery string) { + // we need to sort slices so json compare is accurate + var data [][]interface{} + for _, line := range wqr.Data { + var newline []interface{} + for _, element := range line { + switch newElement := element.(type) { + case featurebase.StringSet: + newline = append(newline, newElement.SortedStringSlice()) + case featurebase.IDSet: + newline = append(newline, newElement.SortedInt64Slice()) + default: + newline = append(newline, element) + } + } + data = append(data, newline) + } + js, err := json.Marshal(data) + if err != nil { + t.Fatal(err) + } + //t.Fatal(string(js)) + require.JSONEq(t, expectedQuery, string(js)) +} + +// recordsFromFile reads lines from a file and builds an in memory structure. +// The file pointed to by pathToRecords must be 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 +} + +// encodeJSONMessages takes a slice of JSON objects and convert them to a slice +// of byte slices which are the binary representation of JSON. +func encodeJSONMessages(records []map[string]interface{}) ([][]byte, error) { + messages := make([][]byte, len(records)) + for i, record := range records { + encoded, err := json.Marshal(record) + if err != nil { + return nil, errors.Errorf("marshaling records to JSON: %s", err) + } + messages[i] = encoded + } + return messages, nil +} + +// encodeAvroMessages takes a slice of JSON objects and convert them to a slice +// of byte slices which are the binary avro encoding based on schema with a +// specific schemaID. +func encodeAvroMessages(records []map[string]interface{}, schema string, schemaID int) ([][]byte, error) { + + // get a thing which can encode a byte slice based on a schema + avroEncoder, err := avro.NewCodec(schema) + if err != nil { + return nil, errors.Errorf("getting avro encoder: %s", err) + } + + messages := make([][]byte, len(records)) + for i, record := range records { + buf := make([]byte, 5, 1000) + buf[0] = 0 + binary.BigEndian.PutUint32(buf[1:], uint32(schemaID)) + buf, err = avroEncoder.BinaryFromNative(buf, record) + if err != nil { + return nil, errors.Errorf("avro encoding record: %s", err) + } + messages[i] = buf + } + return messages, nil +} + +// postSchema takes an avro schema and subject and posts it to schema registry +// at a specific url. +func postSchema(schema, subj, schemaRegistryURL string) (schemaID int, err error) { + schemaClient := csrc.NewClient("http://"+schemaRegistryURL, nil, nil) + resp, err := schemaClient.PostSubjects(subj, schema) + if err != nil { + return 0, errors.Wrap(err, "posting schema") + } + return resp.ID, nil +} + +// createTopic creates a kafka topic on the kafka server pointed to by +// bootstrapServers +func createTopic(topic string, bootstrapServers string) error { + + // create admin client needed to create topic, defer closing + ac, err := confluent.NewAdminClient(&confluent.ConfigMap{"bootstrap.servers": bootstrapServers}) + if err != nil { + return errors.Errorf("creating admin client: %s", err) + } + defer ac.Close() + + // create a topic + var ts = []confluent.TopicSpecification{ + { + Topic: topic, + NumPartitions: 1, + ReplicationFactor: 1, + }, + } + _, err = ac.CreateTopics(context.Background(), ts, nil) + if err != nil { + return errors.Errorf("creating topics: %s", err) + } + // caller must delete if needed + + return nil +} + +// produceMessages produces messages to a kafka topic on the kafka server +// pointed to by bootstrapServer. Messages should be byte slices. +func produceMessages(topic string, bootstrapServers string, messages [][]byte) error { + + // create producer, defer closing + p, err := confluent.NewProducer(&confluent.ConfigMap{"bootstrap.servers": bootstrapServers}) + if err != nil { + return errors.Errorf("creating producer: %s", err) + } + defer p.Close() + + for _, message := range messages { + err := p.Produce(&confluent.Message{ + TopicPartition: confluent.TopicPartition{Topic: &topic, Partition: confluent.PartitionAny}, + Value: message, + }, nil) + if err != nil { + return errors.Errorf("producing messages: %s", err) + } + p.Flush(3000) + } + + return nil + +} + +// deleteTopic deletes a kafka topic on the kafka server pointed to by +// bootstrapServers +func deleteTopic(topic string, bootstrapServers string) error { + + // create admin client needed to create topic, defer closing + ac, err := confluent.NewAdminClient(&confluent.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() != confluent.ErrNoError { + return errors.Errorf("fatal error deleting topic: %s", result.String()) + } + } + + return nil + +} diff --git a/cli/cli_kafka_runner_test.go b/cli/cli_kafka_runner_test.go deleted file mode 100644 index 3bfd66bd9..000000000 --- a/cli/cli_kafka_runner_test.go +++ /dev/null @@ -1,343 +0,0 @@ -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 deleted file mode 100644 index e07f8a5f2..000000000 --- a/cli/internal/kafka_util.go +++ /dev/null @@ -1,74 +0,0 @@ -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 2cfd0b7aa..8834eb7d3 100644 --- a/cli/kafka.go +++ b/cli/kafka.go @@ -8,7 +8,7 @@ import ( func (cmd *Command) newKafkaRunner(cfgFile string) (*kafka.Runner, error) { - cfg, err := kafka.ConfigStructFromFile(cfgFile) + cfg, err := kafka.ConfigFromFile(cfgFile) if err != nil { return nil, err } diff --git a/cli/kafka/config.go b/cli/kafka/config.go index 25390e37f..4e09b9dc2 100644 --- a/cli/kafka/config.go +++ b/cli/kafka/config.go @@ -14,8 +14,8 @@ import ( // 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" + encodingTypeJSON = "json" + encodingTypeAvro = "avro" ) // Config is the user-facing configuration for kafka support in the CLI. This is @@ -32,7 +32,10 @@ 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)"` + Encode string `mapstructure:"encode" help:"Encoding format (currently supported formats: avro, json)"` + AllowMissingFields bool `mapstructure:"allow-missing-fields" help:"allow missing fields in messages from kafka"` + MaxMessages int `mapstructure:"max-messages" help:"max messages read from kakfka"` + ConfluentConfig string `mapstructure:"confluent-config" help:"max messages read from kakfka"` } // Field is a user-facing configuration field. @@ -59,12 +62,15 @@ type ConfigForIDK struct { PrimaryKeys []string Fields []idk.RawField - Encode string + Encode string + AllowMissingFields bool + MaxMessages int + ConfluentConfig 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) { +// ConfigFromFile returns a Config struct based on a configuration file Default +// values for Config struct are defined here +func ConfigFromFile(cfgFile string) (cfg Config, err error) { // configure viper v := viper.New() v.SetConfigFile(cfgFile) @@ -76,7 +82,7 @@ func ConfigStructFromFile(cfgFile string) (cfg Config, err error) { v.SetDefault("batch-size", 1) v.SetDefault("batch-max-staleness", 5*time.Second) v.SetDefault("timeout", 5*time.Second) - v.SetDefault("encode", JSON) + v.SetDefault("encode", encodingTypeJSON) // Read the kafka config file. err = v.ReadInConfig() @@ -92,9 +98,8 @@ func ConfigStructFromFile(cfgFile string) (cfg Config, err error) { } -// ValidateConfig validates the config is usable. -// Note that different encoding methods require -// different configurations +// ValidateConfig validates the config is usable. Note that different encoding +// methods require different configurations func ValidateConfig(c Config) error { // validate common @@ -108,9 +113,9 @@ func ValidateConfig(c Config) error { // validate on a by encoding basis switch c.Encode { - case JSON: + case encodingTypeJSON: return validateConfigJSON(c) - case Avro: + case encodingTypeAvro: return validateConfigAvro(c) } @@ -150,6 +155,24 @@ func validateConfigJSON(c Config) error { // Only primary key fields required func validateConfigAvro(c Config) error { + + // for avro encoded messages, we just need to know what avro fields are + // going to be used for the primary key. We'll check that there is at least + // one field. For every field, we'll check that it has a name attribute and + // has primary-key set. + if len(c.Fields) < 1 { + return errors.New("at least one field is required for avro encoded messages") + } + for i := range c.Fields { + if !c.Fields[i].PrimaryKey { + return errors.New("each field must be a primary key for avro encoded messages") + } + if c.Fields[i].Name == "" { + return errors.New("a name attribute (which isn't equal to \"\") should exist for all fields") + } + + } + return nil } @@ -158,14 +181,17 @@ func ConvertConfig(c Config) (ConfigForIDK, error) { // Copy all the shared members from Config to ConfigForIDK. out := ConfigForIDK{ - Hosts: c.Hosts, - Group: c.Group, - Topics: c.Topics, - BatchSize: c.BatchSize, - BatchMaxStaleness: c.BatchMaxStaleness, - Timeout: c.Timeout, - Table: c.Table, - Encode: c.Encode, + Hosts: c.Hosts, + Group: c.Group, + Topics: c.Topics, + BatchSize: c.BatchSize, + BatchMaxStaleness: c.BatchMaxStaleness, + Timeout: c.Timeout, + Table: c.Table, + Encode: c.Encode, + AllowMissingFields: c.AllowMissingFields, + MaxMessages: c.MaxMessages, + ConfluentConfig: c.ConfluentConfig, } if len(c.Fields) == 0 { @@ -183,14 +209,14 @@ func ConvertConfig(c Config) (ConfigForIDK, error) { if fld.PrimaryKey { 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 + case dax.BaseTypeString, dax.BaseTypeInt: stringKeys = true - + default: + // IDK can handle other field types as primary keys but limiting + // here to the ones above for now. ID and string are the ones + // that make sense and existing users also use int fields so I'm + // including that as well. + return out, errors.Errorf("primary-key fields must be \"id\", \"string\", or \"int\": got field %s which is type %s", fld.Name, keyType) } primaryKeys = append(primaryKeys, fld.Name) } @@ -234,7 +260,7 @@ func ConvertConfig(c Config) (ConfigForIDK, error) { } // 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) + // using a string field as the 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") diff --git a/cli/kafka/config_test.go b/cli/kafka/config_test.go index 2bea16853..5088950ed 100644 --- a/cli/kafka/config_test.go +++ b/cli/kafka/config_test.go @@ -7,25 +7,33 @@ import ( "github.com/stretchr/testify/require" ) -type configStructFromFileTest struct { +type configFromFileTest 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"}`, +func TestConfigFromFile(t *testing.T) { + tests := []configFromFileTest{ + { // confirm json config struct fields are being set + configFilePath: "./testdata/config/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","AllowMissingFields":false,"MaxMessages":0,"ConfluentConfig":""}`, }, { // 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"}`, + configFilePath: "./testdata/config/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","AllowMissingFields":false,"MaxMessages":0,"ConfluentConfig":""}`, + }, + { // confirm avro config struct field are being set + configFilePath: "./testdata/config/config02.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":"avro","AllowMissingFields":false,"MaxMessages":0,"ConfluentConfig":""}`, + }, + { // confirm avro config struct field are being set + configFilePath: "./testdata/config/config03.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":"avro","AllowMissingFields":true,"MaxMessages":100,"ConfluentConfig":"./test/confluent/config.json"}`, }, } for _, test := range tests { - cfg, err := ConfigStructFromFile(test.configFilePath) + cfg, err := ConfigFromFile(test.configFilePath) if err != nil { t.Fatalf("getting config struct from file: %s", err) } diff --git a/cli/kafka/runner.go b/cli/kafka/runner.go index 7b4e2e83f..4c59da79d 100644 --- a/cli/kafka/runner.go +++ b/cli/kafka/runner.go @@ -35,6 +35,7 @@ func NewRunner(cfg ConfigForIDK, batcher fbbatch.Batcher, logWriter io.Writer) * idkMain.Batcher = batcher idkMain.BatchSize = cfg.BatchSize idkMain.BatchMaxStaleness = cfg.BatchMaxStaleness + idkMain.MaxMsgs = uint64(cfg.MaxMessages) idkMain.SetBasic() idkMain.SetLog(logger.NewStandardLogger(logWriter)) idkMain.OffsetMode = true @@ -51,16 +52,16 @@ func NewRunner(cfg ConfigForIDK, batcher fbbatch.Batcher, logWriter io.Writer) * } // 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() + if cfg.Encode == encodingTypeAvro { + kr.GetAvroNewSource(cfg) + } else if cfg.Encode == encodingTypeJSON { + kr.GetJSONNewSource(cfg) } return kr } -func (r *Runner) GetJSONNewSource() { +func (r *Runner) GetJSONNewSource(cfg ConfigForIDK) { r.NewSource = func() (idk.Source, error) { source := kafka_sasl.NewSource() source.KafkaBootstrapServers = r.Hosts @@ -69,11 +70,13 @@ func (r *Runner) GetJSONNewSource() { source.Log = r.Main.Log() source.Timeout = r.Timeout source.HeaderFields = r.Header - cfg, err := common.SetupConfluent(&r.ConfluentCommand) + source.AllowMissingFields = cfg.AllowMissingFields + source.KafkaConfiguration = cfg.ConfluentConfig + confluentCfg, err := common.SetupConfluent(&r.ConfluentCommand) if err != nil { return nil, err } - source.ConfigMap = cfg + source.ConfigMap = confluentCfg err = source.Open() if err != nil { @@ -83,7 +86,7 @@ func (r *Runner) GetJSONNewSource() { } } -func (r *Runner) GetAvroNewSource() { +func (r *Runner) GetAvroNewSource(cfg ConfigForIDK) { r.NewSource = func() (idk.Source, error) { source := kafka.NewSource() source.KafkaBootstrapServers = r.Hosts @@ -91,11 +94,12 @@ func (r *Runner) GetAvroNewSource() { source.Topics = r.Topics source.Log = r.Main.Log() source.Timeout = r.Timeout - cfg, err := common.SetupConfluent(&r.ConfluentCommand) + source.KafkaConfiguration = cfg.ConfluentConfig + confluentcfg, err := common.SetupConfluent(&r.ConfluentCommand) if err != nil { return nil, err } - source.ConfigMap = cfg + source.ConfigMap = confluentcfg err = source.Open() if err != nil { diff --git a/cli/kafka/config_test_data/config00.toml b/cli/kafka/testdata/config/config00.toml similarity index 100% rename from cli/kafka/config_test_data/config00.toml rename to cli/kafka/testdata/config/config00.toml diff --git a/cli/kafka/config_test_data/config01.toml b/cli/kafka/testdata/config/config01.toml similarity index 100% rename from cli/kafka/config_test_data/config01.toml rename to cli/kafka/testdata/config/config01.toml diff --git a/cli/kafka/testdata/config/config02.toml b/cli/kafka/testdata/config/config02.toml new file mode 100644 index 000000000..0c090726d --- /dev/null +++ b/cli/kafka/testdata/config/config02.toml @@ -0,0 +1,27 @@ +hosts = ["kafka:9090"] +group = "testGroup" +topics = "testTopic" +table = "testTable" +batch-size = 35 +batch-max-staleness = "25s" +timeout = "16s" +encode = "avro" +schemaRegistryHost = "localhost:8081" + +[[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/testdata/config/config03.toml b/cli/kafka/testdata/config/config03.toml new file mode 100644 index 000000000..323bfadaf --- /dev/null +++ b/cli/kafka/testdata/config/config03.toml @@ -0,0 +1,30 @@ +hosts = ["kafka:9090"] +group = "testGroup" +topics = "testTopic" +table = "testTable" +batch-size = 35 +batch-max-staleness = "25s" +timeout = "16s" +encode = "avro" +schemaRegistryHost = "localhost:8081" +max-messages = 100 +allow-missing-fields = true +confluent-config = "./test/confluent/config.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_test_data/config/json/config00.toml b/cli/kafka/testdata/runner/config/config00.toml similarity index 86% rename from cli/kafka/runner_test_data/config/json/config00.toml rename to cli/kafka/testdata/runner/config/config00.toml index 472987fe2..d1eb17cde 100644 --- a/cli/kafka/runner_test_data/config/json/config00.toml +++ b/cli/kafka/testdata/runner/config/config00.toml @@ -1,4 +1,4 @@ -hosts = ["localhost:9092"] +hosts = ["KAFKA_SERVICE"] group = "grp" topics = "topic00" table = "table00" @@ -6,6 +6,7 @@ batch-size = 1 batch-max-staleness = "5s" timeout = "5s" encode = "json" +max-messages = MAX_MESSAGES [[fields]] name = "id" diff --git a/cli/kafka/runner_test_data/config/json/config01.toml b/cli/kafka/testdata/runner/config/config01.toml similarity index 86% rename from cli/kafka/runner_test_data/config/json/config01.toml rename to cli/kafka/testdata/runner/config/config01.toml index 737012cb3..bc627e9a2 100644 --- a/cli/kafka/runner_test_data/config/json/config01.toml +++ b/cli/kafka/testdata/runner/config/config01.toml @@ -1,4 +1,4 @@ -hosts = ["localhost:9092"] +hosts = ["KAFKA_SERVICE"] group = "grp" topics = "topic01" table = "table01" @@ -6,6 +6,7 @@ batch-size = 1 batch-max-staleness = "5s" timeout = "5s" encode = "json" +max-messages = MAX_MESSAGES [[fields]] name = "id" diff --git a/cli/kafka/runner_test_data/config/json/config02.toml b/cli/kafka/testdata/runner/config/config02.toml similarity index 86% rename from cli/kafka/runner_test_data/config/json/config02.toml rename to cli/kafka/testdata/runner/config/config02.toml index 5f0b687f1..aea4c7456 100644 --- a/cli/kafka/runner_test_data/config/json/config02.toml +++ b/cli/kafka/testdata/runner/config/config02.toml @@ -1,4 +1,4 @@ -hosts = ["localhost:9092"] +hosts = ["KAFKA_SERVICE"] group = "grp" topics = "topic02" table = "table02" @@ -6,6 +6,7 @@ batch-size = 1 batch-max-staleness = "5s" timeout = "5s" encode = "json" +max-messages = MAX_MESSAGES [[fields]] name = "id" diff --git a/cli/kafka/runner_test_data/config/json/config03.toml b/cli/kafka/testdata/runner/config/config03.toml similarity index 87% rename from cli/kafka/runner_test_data/config/json/config03.toml rename to cli/kafka/testdata/runner/config/config03.toml index 1fd9e08b7..d3f1ff457 100644 --- a/cli/kafka/runner_test_data/config/json/config03.toml +++ b/cli/kafka/testdata/runner/config/config03.toml @@ -1,4 +1,4 @@ -hosts = ["localhost:9092"] +hosts = ["KAFKA_SERVICE"] group = "grp" topics = "topic03" table = "table03" @@ -6,6 +6,7 @@ batch-size = 1 batch-max-staleness = "5s" timeout = "5s" encode = "json" +max-messages = MAX_MESSAGES [[fields]] name = "id" diff --git a/cli/kafka/testdata/runner/config/config04.toml b/cli/kafka/testdata/runner/config/config04.toml new file mode 100644 index 000000000..6df6d5ee8 --- /dev/null +++ b/cli/kafka/testdata/runner/config/config04.toml @@ -0,0 +1,15 @@ +hosts = ["KAFKA_SERVICE"] +group = "grp" +topics = "topic04" +table = "table04" +batch-size = 1 +batch-max-staleness = "5s" +timeout = "5s" +encode = "avro" +schemaRegistryHost = "localhost:8081" +max-messages = MAX_MESSAGES + +[[fields]] +name = "pk0" +source-type = "string" +primary-key = true \ No newline at end of file diff --git a/cli/kafka/testdata/runner/config/config05.toml b/cli/kafka/testdata/runner/config/config05.toml new file mode 100644 index 000000000..8e9beef61 --- /dev/null +++ b/cli/kafka/testdata/runner/config/config05.toml @@ -0,0 +1,27 @@ +hosts = ["KAFKA_SERVICE"] +group = "grp" +topics = "topic05" +table = "table05" +batch-size = 1 +batch-max-staleness = "5s" +timeout = "5s" +encode = "json" +allow-missing-fields = true +max-messages = MAX_MESSAGES + +[[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/data/json/data00.json b/cli/kafka/testdata/runner/data/data00.json similarity index 83% rename from cli/kafka/runner_test_data/data/json/data00.json rename to cli/kafka/testdata/runner/data/data00.json index c6e720ddc..2bb22c5dd 100644 --- a/cli/kafka/runner_test_data/data/json/data00.json +++ b/cli/kafka/testdata/runner/data/data00.json @@ -1,5 +1,5 @@ {"id": 1, "name": "a", "age": 20, "hobbies": "hob1,hob2"} -{"id": 2, "name": "b", "age": 21, "hobbies": "hob3,hob2"} +{"id": 2, "age": 21, "hobbies": "hob3,hob2", "name": "b"} {"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"} diff --git a/cli/kafka/testdata/runner/data/data01.json b/cli/kafka/testdata/runner/data/data01.json new file mode 100644 index 000000000..ab65999eb --- /dev/null +++ b/cli/kafka/testdata/runner/data/data01.json @@ -0,0 +1,10 @@ +{"pk0": "9z4aw", "pk1": "5ptDx", "pk2": "CKs1F", "stringset_string": {"string": "7EYSp"}, "string_string": {"string": "uirDR"}, "stringset_bytes": {"bytes": "gL2Hg"}, "string_bytes": {"bytes": "BmvHF"}, "stringset_stringarray": {"array": ["vbbuf", "VQs7y", "9z4aw", "h1iqc", "aQQxr"]}, "stringset_bytesarray": {"array": ["u2Yr4", "tvNOB", "iYeOV", "ZgkOB", "RPGAm"]}, "idset_long": {"long": 647}, "id_long": {"long": 792}, "idset_int": {"int": 898}, "id_int": {"int": 63}, "idset_longarray": {"array": [442, 167, 230, 344, 733]}, "idset_intarray": {"array": [442, 614, 394, 284, 344]}, "int_long": {"long": 584}, "int_int": {"int": 344}, "decimal_bytes": {"bytes": "\u0001\u00cb"}, "decimal_float": {"float": 3.23}, "decimal_double": {"double": 0.95}, "dateint_bytes_ts": {"bytes": "2023-02-16 07:53:59"}, "bool_bool": {"boolean": false}, "timestamp_bytes_ts": {"bytes": "2023-02-16 07:53:59"}, "timestamp_bytes_int": {"bytes": "1676555639"}} +{"pk0": "ASSAw", "pk1": "kauLy", "pk2": "oxjI0", "stringset_string": {"string": "iYeOV"}, "string_string": {"string": "LBTEU"}, "stringset_bytes": {"bytes": "5ptDx"}, "string_bytes": {"bytes": "EyQoi"}, "stringset_stringarray": {"array": ["iYeOV", "XzEHj", "rrkYB", "v31XN", "uirDR"]}, "stringset_bytesarray": {"array": ["X9jWC", "x5z8P", "PYE8V", "PYE8V", "vTwn4"]}, "idset_long": {"long": 484}, "id_long": {"long": 23}, "idset_int": {"int": 322}, "id_int": {"int": 320}, "idset_longarray": {"array": [792, 809, 168, 399, 639]}, "idset_intarray": {"array": [606, 293, 23, 358, 821]}, "int_long": {"long": 533}, "int_int": {"int": 884}, "decimal_bytes": {"bytes": "\u0001\u0029"}, "decimal_float": {"float": 4.32}, "decimal_double": {"double": 4.97}, "dateint_bytes_ts": {"bytes": "2023-02-22 14:32:23"}, "bool_bool": {"boolean": true}, "timestamp_bytes_ts": {"bytes": "2023-02-22 14:32:23"}, "timestamp_bytes_int": {"bytes": "1677097943"}} +{"pk0": "BmvHF", "pk1": "798ka", "pk2": "6TKzc", "stringset_string": {"string": "X9jWC"}, "string_string": {"string": "I1gXJ"}, "stringset_bytes": {"bytes": "gjWEI"}, "string_bytes": {"bytes": "thuky"}, "stringset_stringarray": {"array": ["tyP3m", "5ptDx", "TLaUE", "EyQoi", "Chgzr"]}, "stringset_bytesarray": {"array": ["VQs7y", "h1iqc", "F0uC4", "d0U7s", "byHh9"]}, "idset_long": {"long": 166}, "id_long": {"long": 320}, "idset_int": {"int": 232}, "id_int": {"int": 286}, "idset_longarray": {"array": [890, 975, 284, 289, 388]}, "idset_intarray": {"array": [865, 931, 614, 884, 322]}, "int_long": {"long": 857}, "int_int": {"int": 879}, "decimal_bytes": {"bytes": "\u0000\u00e6"}, "decimal_float": {"float": 2.84}, "decimal_double": {"double": 2.91}, "dateint_bytes_ts": {"bytes": "2023-01-31 11:11:30"}, "bool_bool": {"boolean": true }, "timestamp_bytes_ts": {"bytes": "2023-01-31 11:11:30"}, "timestamp_bytes_int": {"bytes": "1675185090"}} +{"pk0": "tElMR", "pk1": "ARlcJ", "pk2": "n9HUP", "stringset_string": {"string": "58KIR"}, "string_string": {"string": "FW39I"}, "stringset_bytes": {"bytes": "PNB4s"}, "string_bytes": {"bytes": "FW39I"}, "stringset_stringarray": {"array": ["X9jWC", "58KIR", "X9jWC", "6TKzc", "8MGwy"]}, "stringset_bytesarray": {"array": ["vhisL", "BmvHF", "eofzb", "TLaUE", "PNB4s"]}, "idset_long": {"long": 289}, "id_long": {"long": 695}, "idset_int": {"int": 791}, "id_int": {"int": 821}, "idset_longarray": {"array": [2, 680, 958, 289, 389]}, "idset_intarray": {"array": [606, 890, 387, 102, 220]}, "int_long": {"long": 289}, "int_int": {"int": 2}, "decimal_bytes": {"bytes": "\u005f"}, "decimal_float": {"float": 2.65}, "decimal_double": {"double": 2.19}, "dateint_bytes_ts": {"bytes": "2023-02-20 17:04:21"}, "bool_bool": {"boolean": false}, "timestamp_bytes_ts": {"bytes": "2023-02-20 17:04:21"}, "timestamp_bytes_int": {"bytes": "1676934261"}} +{"pk0": "RKE3c", "pk1": "6TKzc", "pk2": "RKE3c", "stringset_string": {"string": "dF6kx"}, "string_string": {"string": "TLaUE"}, "stringset_bytes": {"bytes": "dxKKn"}, "string_bytes": {"bytes": "YdwQY"}, "stringset_stringarray": {"array": ["5HIn2", "wNZ7o", "KdTtE", "x5z8P", "nVQrd"]}, "stringset_bytesarray": {"array": ["5HIn2", "tyP3m", "I6NST", "gjWEI", "Qylqq"]}, "idset_long": {"long": 39}, "id_long": {"long": 665}, "idset_int": {"int": 113}, "id_int": {"int": 681}, "idset_longarray": {"array": [857, 63, 172, 220, 358]}, "idset_intarray": {"array": [220, 731, 647, 778, 665]}, "int_long": {"long": 582}, "int_int": {"int": 690}, "decimal_bytes": {"bytes": "\u0001\u00d1"}, "decimal_float": {"float": 1.36}, "decimal_double": {"double": 2.72}, "dateint_bytes_ts": {"bytes": "2023-02-16 20:09:13"}, "bool_bool": {"boolean": true }, "timestamp_bytes_ts": {"bytes": "2023-02-16 20:09:13"}, "timestamp_bytes_int": {"bytes": "1676599753"}} +{"pk0": "yg8hY", "pk1": "tvNOB", "pk2": "byHh9", "stringset_string": {"string": "911oj"}, "string_string": {"string": "5HIn2"}, "stringset_bytes": {"bytes": "u2Yr4"}, "string_bytes": {"bytes": "qK5TE"}, "stringset_stringarray": {"array": ["nVQrd", "fjQK2", "m5d59", "dxKKn", "d0U7s"]}, "stringset_bytesarray": {"array": ["wNZ7o", "OKNV2", "F0uC4", "VBcyJ", "KMZnH"]}, "idset_long": {"long": 839}, "id_long": {"long": 809}, "idset_int": {"int": 533}, "id_int": {"int": 168}, "idset_longarray": {"array": [582, 629, 680, 63, 690]}, "idset_intarray": {"array": [969, 175, 172, 257, 115]}, "int_long": {"long": 433}, "int_int": {"int": 680}, "decimal_bytes": {"bytes": "\u0001\u00d6"}, "decimal_float": {"float": 1.27}, "decimal_double": {"double": 3.23}, "dateint_bytes_ts": {"bytes": "2023-01-30 06:56:05"}, "bool_bool": {"boolean": false}, "timestamp_bytes_ts": {"bytes": "2023-01-30 06:56:05"}, "timestamp_bytes_int": {"bytes": "1675083365"}} +{"pk0": "6TKzc", "pk1": "YKLk9", "pk2": "h1iqc", "stringset_string": {"string": "eofzb"}, "string_string": {"string": "n9HUP"}, "stringset_bytes": {"bytes": "t5f7R"}, "string_bytes": {"bytes": "5HIn2"}, "stringset_stringarray": {"array": ["yg8hY", "xE5jX", "C6xxn", "BmvHF", "PYE8V"]}, "stringset_bytesarray": {"array": ["6TKzc", "vK0WD", "xE5jX", "jVVfZ", "pjxqm"]}, "idset_long": {"long": 72}, "id_long": {"long": 387}, "idset_int": {"int": 676}, "id_int": {"int": 797}, "idset_longarray": {"array": [399, 322, 975, 730, 969]}, "idset_intarray": {"array": [958, 242, 778, 289, 797]}, "int_long": {"long": 430}, "int_int": {"int": 23}, "decimal_bytes": {"bytes": "\u0001\u00d6"}, "decimal_float": {"float": 3.35}, "decimal_double": {"double": 0.78}, "dateint_bytes_ts": {"bytes": "2023-02-23 05:04:34"}, "bool_bool": {"boolean": true }, "timestamp_bytes_ts": {"bytes": "2023-02-23 05:04:34"}, "timestamp_bytes_int": {"bytes": "1677150274"}} +{"pk0": "h1iqc", "pk1": "5ptDx", "pk2": "iYeOV", "stringset_string": {"string": "ASSAw"}, "string_string": {"string": "58KIR"}, "stringset_bytes": {"bytes": "eNKWF"}, "string_bytes": {"bytes": "x5z8P"}, "stringset_stringarray": {"array": ["pjxqm", "6TKzc", "ZgkOB", "eofzb", "RKE3c"]}, "stringset_bytesarray": {"array": ["BmvHF", "Qylqq", "5HIn2", "7EYSp", "yTeUQ"]}, "idset_long": {"long": 255}, "id_long": {"long": 647}, "idset_int": {"int": 821}, "id_int": {"int": 389}, "idset_longarray": {"array": [606, 110, 320, 63, 344]}, "idset_intarray": {"array": [29, 289, 388, 257, 606]}, "int_long": {"long": 110}, "int_int": {"int": 148}, "decimal_bytes": {"bytes": "\u0001\u001c"}, "decimal_float": {"float": 0.4}, "decimal_double": {"double": 4.41}, "dateint_bytes_ts": {"bytes": "2023-02-19 08:52:56"}, "bool_bool": {"boolean": false}, "timestamp_bytes_ts": {"bytes": "2023-02-19 08:52:56"}, "timestamp_bytes_int": {"bytes": "1676818376"}} +{"pk0": "DY2Ui", "pk1": "kUbdU", "pk2": "pjxqm", "stringset_string": {"string": "tyP3m"}, "string_string": {"string": "8MGwy"}, "stringset_bytes": {"bytes": "DDLN5"}, "string_bytes": {"bytes": "vTwn4"}, "stringset_stringarray": {"array": ["XzEHj", "8MGwy", "gjWEI", "xE5jX", "v31XN"]}, "stringset_bytesarray": {"array": ["d0U7s", "u2Yr4", "d0U7s", "sDdtS", "y2Y7b"]}, "idset_long": {"long": 984}, "id_long": {"long": 430}, "idset_int": {"int": 931}, "id_int": {"int": 297}, "idset_longarray": {"array": [975, 733, 113, 751, 772]}, "idset_intarray": {"array": [297, 72, 694, 898, 384]}, "int_long": {"long": 63}, "int_int": {"int": 388}, "decimal_bytes": {"bytes": "\u007e"}, "decimal_float": {"float": 0.83}, "decimal_double": {"double": 4.23}, "dateint_bytes_ts": {"bytes": "2023-02-12 18:37:16"}, "bool_bool": {"boolean": false}, "timestamp_bytes_ts": {"bytes": "2023-02-12 18:37:16"}, "timestamp_bytes_int": {"bytes": "1676248636"}} +{"pk0": "u2Yr4", "pk1": "sHaUv", "pk2": "x5z8P", "stringset_string": {"string": "7EYSp"}, "string_string": {"string": "ZgkOB"}, "stringset_bytes": {"bytes": "qK5TE"}, "string_bytes": {"bytes": "6iGIm"}, "stringset_stringarray": {"array": ["u2Yr4", "PYE8V", "VBcyJ", "Chgzr", "DY2Ui"]}, "stringset_bytesarray": {"array": ["YdwQY", "kUbdU", "aQQxr", "KdTtE", "MVNow"]}, "idset_long": {"long": 148}, "id_long": {"long": 115}, "idset_int": {"int": 890}, "id_int": {"int": 39}, "idset_longarray": {"array": [839, 63, 148, 984, 958]}, "idset_intarray": {"array": [731, 13, 167, 772, 629]}, "int_long": {"long": 13}, "int_int": {"int": 969}, "decimal_bytes": {"bytes": "\u0000\u009a"}, "decimal_float": {"float": 2.93}, "decimal_double": {"double": 2.29}, "dateint_bytes_ts": {"bytes": "2023-02-03 16:19:37"}, "bool_bool": {"boolean": false}, "timestamp_bytes_ts": {"bytes": "2023-02-03 16:19:37"}, "timestamp_bytes_int": {"bytes": "1675462777"}} \ No newline at end of file diff --git a/cli/kafka/testdata/runner/data/data02.json b/cli/kafka/testdata/runner/data/data02.json new file mode 100644 index 000000000..49b056000 --- /dev/null +++ b/cli/kafka/testdata/runner/data/data02.json @@ -0,0 +1,6 @@ +{"id": 1, "name": "a", "age": 20, "hobbies": "hob1,hob2"} +{"id": 2, "age": 21, "hobbies": "hob3,hob2", "name": "b"} +{"id": 3, "age": 22, "hobbies": "hob3,hob4"} +{"id": 4, "name": "d", "hobbies": "hob4,hob5"} +{"id": 5, "name": "e", "age": 24} +{"id": 6, "name": "f", "age": 26, "hobbies": "hob6,hob7"} \ No newline at end of file diff --git a/cli/kafka/testdata/runner/schema/schema01.json b/cli/kafka/testdata/runner/schema/schema01.json new file mode 100644 index 000000000..572f93bd4 --- /dev/null +++ b/cli/kafka/testdata/runner/schema/schema01.json @@ -0,0 +1,32 @@ +{ + "namespace": "org.test", + "type": "record", + "name": "all_type_schema", + "doc": "All supported avro types and property variations", + "fields": [ + {"name": "pk0", "type": "string"}, + {"name": "pk1", "type": "string"}, + {"name": "pk2", "type": "string"}, + {"name": "stringset_string", "type": ["string", "null"], "mutex": false }, + {"name": "string_string", "type": ["string", "null"], "mutex": true }, + {"name": "stringset_bytes", "type": ["bytes", "null"], "mutex": false}, + {"name": "string_bytes", "type": ["bytes", "null"] , "mutex": true }, + {"name": "stringset_stringarray", "type": [{"type": "array", "items": "string"}, "null"]}, + {"name": "stringset_bytesarray", "type": [{"type": "array", "items": "string"}, "null"]}, + {"name": "idset_long", "type": ["long", "null"], "mutex": false, "fieldType": "id"}, + {"name": "id_long", "type": ["long", "null"], "mutex": true, "fieldType": "id"}, + {"name": "idset_int", "type": ["int", "null"], "mutex": false, "fieldType": "id"}, + {"name": "id_int", "type": ["int", "null"], "mutex": true, "fieldType": "id"}, + {"name": "idset_longarray", "type": [{"type": "array", "items": "long"}, "null"], "fieldType": "id"}, + {"name": "idset_intarray", "type": [{"type": "array", "items": "int"}, "null"]}, + {"name": "int_long", "type": ["long", "null"], "fieldType": "int"}, + {"name": "int_int", "type": ["int", "null"], "fieldType": "int"}, + {"name": "decimal_bytes", "type": ["bytes", "null"], "fieldType": "decimal", "scale": 2}, + {"name": "decimal_float", "type": ["float", "null"], "fieldType": "decimal", "scale": 2}, + {"name": "decimal_double", "type": ["double", "null"], "fieldType": "decimal", "scale": 2}, + {"name": "dateint_bytes_ts", "type": ["bytes", "null"], "fieldType": "dateInt", "layout": "2006-01-02 15:04:05", "unit": "s", "epoch": "1970-01-01 00:00:00"}, + {"name": "bool_bool", "type": ["boolean", "null"]}, + {"name": "timestamp_bytes_ts", "type": ["bytes", "null"], "fieldType": "timestamp", "layout": "2006-01-02 15:04:05", "epoch": "1970-01-01 00:00:00"}, + {"name": "timestamp_bytes_int", "type": ["bytes", "null"], "fieldType": "timestamp", "unit": "s", "layout": "2006-01-02 15:04:05", "epoch": "1970-01-01 00:00:00"} + ] +} diff --git a/idk/ingest.go b/idk/ingest.go index a3f0f7144..3958556e5 100644 --- a/idk/ingest.go +++ b/idk/ingest.go @@ -301,13 +301,34 @@ func (m *Main) run() error { func (m *Main) clone() (*Main, error) { var index *pilosaclient.Index + noOpSchemaManager := false + + // If you have a schema manager, it does it's thing. Otherwise, it's a no + // opt manager. Then you get a default schema. If you get a default schema, + // you get a default index. This seems fine for IDK. However, for the CLI / + // SQL kafka runner, we use the m.index to create the dax.Table. If m.index + // is set to default values, then keys is false even when we don't want it + // to be. This was causing issue in buildBulkInsert in the batch package of + // the CLI. + switch m.SchemaManager.(type) { + case *nopSchemaManager: + noOpSchemaManager = true + } schema, err := m.SchemaManager.Schema() if err != nil { return nil, err } - index = schema.Index(m.Index) + if noOpSchemaManager && len(m.PrimaryKeyFields) > 0 { + // if IDK has PrimaryKeyFields, it expects keyed index + keys := pilosaclient.OptIndexKeys(true) + // most queries don't work with this set to false so set to true + exists := pilosaclient.OptIndexTrackExistence(true) + index = schema.Index(m.Index, keys, exists) + } else { + index = schema.Index(m.Index) + } // use a copy (schema race condition issues) mClone := *m @@ -659,7 +680,6 @@ initialFetch: for n := range lookupRow { lookupRow[n] = nil } - } } diff --git a/idk/kafka/cmd.go b/idk/kafka/cmd.go index 5e20ad76e..4117fe657 100644 --- a/idk/kafka/cmd.go +++ b/idk/kafka/cmd.go @@ -23,10 +23,9 @@ func NewMain() (*Main, error) { ConfluentCommand: idk.ConfluentCommand{ KafkaBootstrapServers: []string{"localhost:9092"}, }, - Group: "defaultgroup", - Topics: []string{"defaulttopic"}, - Timeout: time.Second, - ConsumerCloseTimeout: 30, + Group: "defaultgroup", + Topics: []string{"defaulttopic"}, + Timeout: time.Second, } m.SchemaRegistryURL = "http://" + defaultRegistryHost diff --git a/idk/kafka/source.go b/idk/kafka/source.go index 3ead55b75..00adf0bff 100644 --- a/idk/kafka/source.go +++ b/idk/kafka/source.go @@ -48,8 +48,7 @@ type Source struct { highmarks []confluent.TopicPartition client *confluent.Consumer recordChannel chan recordWithError - - ConfigMap *confluent.ConfigMap + ConfigMap *confluent.ConfigMap // lastSchemaID and lastSchema keep track of the most recent // schema in use. We expect this not to change often, but when it @@ -78,12 +77,13 @@ func NewSource() *Source { Group: "group0", Log: logger.NopLogger, - lastSchemaID: -1, - cache: make(map[int32]avro.Schema), - recordChannel: make(chan recordWithError), - quit: make(chan struct{}), - ConfigMap: &confluent.ConfigMap{}, - highmarks: make([]confluent.TopicPartition, 0), + lastSchemaID: -1, + cache: make(map[int32]avro.Schema), + recordChannel: make(chan recordWithError), + quit: make(chan struct{}), + ConfigMap: &confluent.ConfigMap{}, + highmarks: make([]confluent.TopicPartition, 0), + consumerCloseTimeout: 30, } src.SchemaRegistryURL = "http://" + defaultRegistryHost @@ -189,6 +189,10 @@ func (s *Source) toPDKRecord(vals map[string]interface{}) []interface{} { return data } +func (s *Source) CommitMessages(recs []confluent.TopicPartition) ([]confluent.TopicPartition, error) { + return s.client.CommitOffsets(recs) +} + type Record struct { src *Source topic string @@ -205,6 +209,10 @@ func (r *Record) StreamOffset() (string, uint64) { var _ idk.OffsetStreamRecord = &Record{} +func (r *Record) Schema() interface{} { + return r.avroSchema +} + func (r *Record) Commit(ctx context.Context) error { r.src.mu.Lock() defer r.src.mu.Unlock() @@ -226,7 +234,6 @@ func (r *Record) Commit(ctx context.Context) error { r.src.highmarks = r.src.highmarks[:0] // sort by increasing partition, decreasing offset - for _, x := range section { if s != *x.Topic || p != x.Partition { r.src.highmarks = append(r.src.highmarks, x) @@ -235,12 +242,10 @@ func (r *Record) Commit(ctx context.Context) error { s = *x.Topic } - committedOffsets, err := r.src.CommitMessages(r.src.highmarks) if err != nil { return errors.Wrap(err, "failed to commit messages") } - if r.src.Verbose { for _, o := range committedOffsets { r.src.Log.Debugf("t: %v p: %v o: %v", *o.Topic, o.Partition, o.Offset) @@ -249,6 +254,7 @@ func (r *Record) Commit(ctx context.Context) error { r.src.spool = remaining r.src.spoolBase = idx + return nil } @@ -256,14 +262,6 @@ func (r *Record) Data() []interface{} { return r.data } -func (r *Record) Schema() interface{} { - return r.avroSchema -} - -func (s *Source) CommitMessages(recs []confluent.TopicPartition) ([]confluent.TopicPartition, error) { - return s.client.CommitOffsets(recs) -} - // Open initializes the kafka source. (i.e. creating and configuring a consumer) // The configuration options for the confluentinc/confluent-kafka-go/kafka // libarary are: https://github.com/confluentinc/librdkafka/blob/master/CONFIGURATION.md @@ -286,6 +284,8 @@ func (s *Source) Open() error { } s.httpClient = getHTTPClient(tlsConfig) } + + // group if s.Group != "" { err = s.ConfigMap.SetKey("group.id", s.Group) if err != nil { @@ -321,6 +321,8 @@ func (s *Source) Open() error { return errors.Wrap(err, "new consumer") } + // by default, Kafka will use the stored offset (the latest committed message) and continue on from there. + // to skip old msgs, use rebalanceCbSkipOld to manually set offset to the end err = cl.SubscribeTopics(s.Topics, nil) if err != nil { return errors.Wrap(err, "subscribe topics") @@ -333,6 +335,7 @@ func (s *Source) Open() error { s.client = cl s.opened = true s.wg.Add(1) + go func() { s.generator() }() @@ -340,23 +343,6 @@ func (s *Source) Open() error { return nil } -func (s *Source) cleanRegistryURL() error { - // We can't immediately url.Parse the RegistryURL because parsing - // a host without a scheme is invalid. First we'll check for a - // scheme and add the default http:// if needed. - if !strings.Contains(s.SchemaRegistryURL, "://") { - s.SchemaRegistryURL = "http://" + s.SchemaRegistryURL - } - - SchemaRegistryURL, err := url.Parse(s.SchemaRegistryURL) - if err != nil { - return errors.Wrap(err, "parse registry URL") - } - - s.SchemaRegistryURL = SchemaRegistryURL.String() - return nil -} - func (c *Source) generator() { defer func() { close(c.recordChannel) @@ -474,6 +460,23 @@ func (s *Source) Close() error { return nil } +func (s *Source) cleanRegistryURL() error { + // We can't immediately url.Parse the RegistryURL because parsing + // a host without a scheme is invalid. First we'll check for a + // scheme and add the default http:// if needed. + if !strings.Contains(s.SchemaRegistryURL, "://") { + s.SchemaRegistryURL = "http://" + s.SchemaRegistryURL + } + + SchemaRegistryURL, err := url.Parse(s.SchemaRegistryURL) + if err != nil { + return errors.Wrap(err, "parse registry URL") + } + + s.SchemaRegistryURL = SchemaRegistryURL.String() + return nil +} + // TODO change name func (s *Source) decodeAvroValueWithSchemaRegistry(val []byte) (interface{}, avro.Schema, error) { if len(val) < 6 || val[0] != 0 { diff --git a/idk/kafka_sasl/source.go b/idk/kafka_sasl/source.go index de9b0967f..feb063c4e 100644 --- a/idk/kafka_sasl/source.go +++ b/idk/kafka_sasl/source.go @@ -1,8 +1,10 @@ package kafka_sasl import ( + "bytes" "context" "encoding/json" + "fmt" "io" "os" "sort" @@ -12,6 +14,7 @@ import ( 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/logger" "github.com/pkg/errors" ) @@ -21,12 +24,14 @@ import ( // achieve concurrency, create multiple Sources. type Source struct { idk.ConfluentCommand - Topics []string - Group string - Log logger.Logger - Timeout time.Duration - SkipOld bool - AllowMissingFields bool + Topics []string + Group string + Log logger.Logger + Timeout time.Duration + SkipOld bool + Verbose bool + AllowMissingFields bool + consumerCloseTimeout time.Duration // Header is a file referencing a file containing JSON header configuration. Header string @@ -58,11 +63,15 @@ type Source struct { // NewSource gets a new Source func NewSource() *Source { src := &Source{ - Topics: []string{"test"}, - Group: "group0", - Log: logger.NopLogger, - recordChannel: make(chan recordWithError), - quit: make(chan struct{}), + ConfluentCommand: idk.ConfluentCommand{}, + Topics: []string{"test"}, + Group: "group0", + Log: logger.NopLogger, + recordChannel: make(chan recordWithError), + quit: make(chan struct{}), + ConfigMap: &confluent.ConfigMap{}, + highmarks: make([]confluent.TopicPartition, 0), + consumerCloseTimeout: 30, } src.KafkaBootstrapServers = []string{"localhost:9092"} @@ -81,8 +90,9 @@ func (s *Source) Record() (idk.Record, error) { case context.DeadlineExceeded: return nil, idk.ErrFlush default: - return nil, errors.Wrap(rec.Err, "failed to fetch record from Kafka Confluent") + return nil, errors.Wrap(rec.Err, "failed to fetch record from Kafka") } + if rec.Record == nil { return nil, idk.ErrFlush } @@ -161,32 +171,41 @@ func (r *Record) Schema() interface{} { return nil } func (r *Record) Commit(ctx context.Context) error { r.src.mu.Lock() defer r.src.mu.Unlock() - idx, base := r.idx, r.src.spoolBase if idx < base { return errors.New("cannot commit a record that has already been committed") } section, remaining := r.src.spool[:idx-base], r.src.spool[idx-base:] - // sort by increasing partition, decreasing offset sort.Slice(section, func(i, j int) bool { - if section[i].Partition != section[j].Partition { + if *section[i].Topic != *section[j].Topic { + return *section[i].Topic < *section[j].Topic + } else if section[i].Partition != section[j].Partition { return section[i].Partition < section[j].Partition } return section[i].Offset > section[j].Offset }) - // calculate the high marks p := int32(-1) + s := "" r.src.highmarks = r.src.highmarks[:0] + + // sort by increasing partition, decreasing offset for _, x := range section { - if p != x.Partition { + if s != *x.Topic || p != x.Partition { r.src.highmarks = append(r.src.highmarks, x) } p = x.Partition + s = *x.Topic + } - _, err := r.src.CommitMessages(r.src.highmarks) + committedOffsets, err := r.src.CommitMessages(r.src.highmarks) if err != nil { return errors.Wrap(err, "failed to commit messages") } + if r.src.Verbose { + for _, o := range committedOffsets { + r.src.Log.Debugf("t: %v p: %v o: %v", *o.Topic, o.Partition, o.Offset) + } + } r.src.spool = remaining r.src.spoolBase = idx @@ -200,12 +219,17 @@ func (r *Record) Data() []interface{} { // Open initializes the kafka source. func (s *Source) Open() error { + cfg, err := common.SetupConfluent(&s.ConfluentCommand) + if err != nil { + return err + } + s.ConfigMap = cfg + if len(s.Header) == 0 && len(s.HeaderFields) == 0 { return errors.New("needs header specification file (file or fields)") } var headerData []byte - var err error if s.Header != "" { headerData, err = os.ReadFile(s.Header) if err != nil { @@ -232,21 +256,30 @@ func (s *Source) Open() error { return err } } + // when there is no initial offset in Kafka or if the current offset does not exist any more, // use this as starting offset: // "earliest": automatically reset the offset to the earliest offset // "latest": automatically reset the offset to the latest offset - err = s.ConfigMap.SetKey("auto.offset.reset", "earliest") + offset := "earliest" + if s.SkipOld { + offset = "latest" + } + err = s.ConfigMap.SetKey("auto.offset.reset", offset) if err != nil { return err } - if s.SkipOld { - err = s.ConfigMap.SetKey("auto.offset.reset", "latest") + if s.Verbose { + buf := bytes.NewBufferString("Confluent Config Map:") + encoder := json.NewEncoder(buf) + err = encoder.Encode(s.ConfigMap) if err != nil { return err } + s.Log.Debugf(buf.String()) + stv, iv := confluent.LibraryVersion() + s.Log.Debugf("version:(%v) %v", iv, stv) } - cl, err := confluent.NewConsumer(s.ConfigMap) if err != nil { return errors.Wrap(err, "new consumer") @@ -259,6 +292,10 @@ func (s *Source) Open() error { return errors.Wrap(err, "subscribe topics") } + if s.Verbose { + s.Log.Debugf("subscribed to %v", s.Topics) + } + s.client = cl s.opened = true s.wg.Add(1) @@ -280,6 +317,9 @@ func (c *Source) generator() { select { case <-c.quit: + if c.Verbose { + c.Log.Debugf("source quit") + } return default: ev := c.client.Poll(100) @@ -293,6 +333,9 @@ func (c *Source) generator() { case confluent.AssignedPartitions: err := c.client.Assign(e.Partitions) if err != nil { + if c.Verbose { + c.Log.Debugf("quit AssignedParitions") + } return } // If we received an `RevokedPartitions` event, we need to revoke this @@ -301,6 +344,9 @@ func (c *Source) generator() { case confluent.RevokedPartitions: err := c.client.Unassign() if err != nil { + if c.Verbose { + c.Log.Debugf("quit RevokeParkitions") + } return } @@ -312,6 +358,9 @@ func (c *Source) generator() { select { case c.recordChannel <- msg: case <-c.quit: + if c.Verbose { + c.Log.Debugf("source quit Error") + } return } @@ -328,9 +377,17 @@ func (c *Source) generator() { select { case c.recordChannel <- msg: case <-c.quit: + if c.Verbose { + c.Log.Debugf("source quit Message") + } return } + case confluent.OffsetsCommitted: + c.Log.Debugf("commited %s", e) default: + if c.Verbose { + c.Log.Debugf("ignored %#v", ev) + } continue // consumer doesn't care about all event types (e.g. OffsetsCommitted) } } @@ -341,10 +398,26 @@ func (c *Source) generator() { func (s *Source) Close() error { if s.client != nil { if s.opened { // only close opened sources + var err error + closedReturned := make(chan error, 1) + // send quit message to polling routine & wait for it to exit s.quit <- struct{}{} s.wg.Wait() - err := s.client.Close() - s.opened = false + s.Log.Debugf("Trying to close consumer %s...", s.client.String()) + go func() { + closedReturned <- s.client.Close() + }() + start := time.Now() + select { + case err = <-closedReturned: + if err == nil { + s.Log.Debugf("Successfully closed consumer %s!", s.client.String()) + s.opened = false + } + case <-time.After(time.Duration(s.consumerCloseTimeout * 1000 * 1000 * 1000)): + err = fmt.Errorf("unable to properly close consumer %s after %f seconds", s.client.String(), time.Since(start).Seconds()) + } + return errors.Wrap(err, "closing kafka consumer") } } diff --git a/wire_response.go b/wire_response.go index dfb08bc00..cf02a33db 100644 --- a/wire_response.go +++ b/wire_response.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "log" + "sort" "strings" "time" @@ -210,6 +211,17 @@ func (ii IDSet) String() string { return sb.String() } +// SortedInt64Slice returns the values in a IDSet field in a int64 slice that is +// sorted +func (ii IDSet) SortedInt64Slice() []int64 { + var idSetSlice = make([]int64, len(ii)) + for i, str := range ii { + idSetSlice[i] = str + } + sort.Slice(idSetSlice, func(i, j int) bool { return idSetSlice[i] < idSetSlice[j] }) + return idSetSlice +} + // StringSet is a return type specific to SQLResponse types. type StringSet []string @@ -227,6 +239,17 @@ func (ss StringSet) String() string { return sb.String() } +// SortedStringSlice returns the values in a StringSet field in a string slice +// that is sorted +func (ss StringSet) SortedStringSlice() []string { + var stringSetSlice = make([]string, len(ss)) + for i, str := range ss { + stringSetSlice[i] = str + } + sort.Strings(stringSetSlice) + return stringSetSlice +} + // ShowColumnsResponse returns a structure which is specific to a `SHOW COLUMNS` // statement, derived from the results in the WireQueryResponse. This is kind of // a crude way to unmarshal a WireQueryResponse into a type which is specific to