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