cli kafka runner updates

cli/batch/sql.go

@@ -175,10 +175,26 @@ func buildBulkInsert(tbl *dax.Table, fields []*dax.Field, ids []interface{}, row
 	sb.WriteString(strings.Join(flds, ","))
 
 	// MAP
-	keyType := dax.BaseTypeID
-	if tbl.StringKeys() {
-		keyType = dax.BaseTypeString
+
+	// We need to set ID value below based on ids. Rather than changing what we
+	// get by changing IDK, decide how to format based on the type of ids. We
+	// get []uint64 for not keyed indexes or []byte for keyed indexes.
+	// Additionally, IDK doesn't give us a valid table for some reason (the keys
+	// options is always false). We will also use ids type to determine how to
+	// map values in the MAP clause
+	var fmtStr string
+	var keyType string
+	if len(ids) > 0 {
+		switch ids[0].(type) {
+		case uint64:
+			fmtStr = "%d"
+			keyType = dax.BaseTypeID
+		case []byte:
+			fmtStr = "%s"
+			keyType = dax.BaseTypeString
+		}
 	}
+
 	sb.WriteString(`) MAP ('$._id' `)
 	sb.WriteString(keyType)
 	sb.WriteString(`,`)
@@ -191,9 +207,10 @@ func buildBulkInsert(tbl *dax.Table, fields []*dax.Field, ids []interface{}, row
 	// bulk insert as one line in the NDJSON payload. We re-use the map for each
 	// row.
 	m := make(map[string]interface{})
+
 	for i := range rows {
 		// Write the ID value.
-		m[string(dax.PrimaryKeyFieldName)] = ids[i]
+		m[string(dax.PrimaryKeyFieldName)] = fmt.Sprintf(fmtStr, ids[i])
 		// Write the rest of the data values.
 		for col := range rows[i] {
 			m[fmt.Sprintf("col_%d", col)] = rows[i][col]

cli/cli_kafka_runner_test.go

@@ -0,0 +1,343 @@
+package cli
+
+import (
+	"bufio"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/featurebasedb/featurebase/v3/cli/internal"
+	"github.com/featurebasedb/featurebase/v3/cli/kafka"
+	"github.com/featurebasedb/featurebase/v3/errors"
+	"github.com/featurebasedb/featurebase/v3/logger"
+	"github.com/stretchr/testify/require"
+)
+
+// These variables are used throught tests to refer to resouces needed for the
+// tests (i.e. where to find the kafka service)
+var (
+	pilosaHost     string
+	pilosaTLSHost  string
+	pilosaGrpcHost string
+	kafkaHost      string
+	registryHost   string
+	certPath       string
+)
+
+// The variables above are set here. If local is set to true, the tests assume
+// all services are running locally. Otherwise, we use the name of the
+// containers that these tests will have access to in CI.
+func init() {
+	local := true
+	var ok bool
+	if pilosaHost, ok = os.LookupEnv("IDK_TEST_PILOSA_HOST"); !ok {
+		if local {
+			pilosaHost = "localhost:10101"
+		} else {
+			pilosaHost = "pilosa:10101"
+		}
+	}
+	if pilosaTLSHost, ok = os.LookupEnv("IDK_TEST_PILOSA_TLS_HOST"); !ok {
+		pilosaTLSHost = "https://pilosa-tls:10111"
+	}
+	if pilosaGrpcHost, ok = os.LookupEnv("IDK_TEST_PILOSA_GRPC_HOST"); !ok {
+		if local {
+			pilosaGrpcHost = "localhost:20101"
+		} else {
+			pilosaGrpcHost = "pilosa:20101"
+		}
+	}
+	if kafkaHost, ok = os.LookupEnv("IDK_TEST_KAFKA_HOST"); !ok {
+		if local {
+			kafkaHost = "localhost:9092"
+		} else {
+			kafkaHost = "kafka:9092"
+		}
+	}
+	if registryHost, ok = os.LookupEnv("IDK_TEST_REGISTRY_HOST"); !ok {
+		if local {
+			registryHost = "localhost:8081"
+		} else {
+			registryHost = "schema-registry:8081"
+		}
+	}
+	if certPath, ok = os.LookupEnv("IDK_TEST_CERT_PATH"); !ok {
+		certPath = "/certs"
+	}
+}
+
+// Struct used for TestRunner which contains the information needed to ingest
+// data to kafka, configure the runner, and test that the runner successfully
+// ran.
+type KafkaRunnerTest struct {
+	ConfigFile      string      // path to the configuration file used for the runner
+	DataFile        string      // path to the data file to populate kafka with
+	Encode          string      // how to encode the data from the data file
+	CreateTableStmt string      // statement used to create table prior to ingest
+	Tests           []TestQuery // list of test which are 2-tuples of query and expected results
+
+}
+
+// Struct which contains a query to run agaisnt featurebase and the expected
+// results from that query.
+type TestQuery struct {
+	Query        string
+	ExpectedResp string
+}
+
+// The paths (from the folder that contains this file) to folders with kafka
+// runner configurations files and data to be sent and consumed from kafka.
+var confPathPrefix string = "./kafka/runner_test_data/config/"
+var dataPathPrefix string = "./kafka/runner_test_data/data/"
+
+// A slice of KafkaRunnerTest structs that will be used in TestKafkaRunner test
+// function.
+var kafkaRunnerTests = []KafkaRunnerTest{
+	{ // id keys
+		ConfigFile: "config00.toml",
+		DataFile:   "data00.json",
+		Encode:     kafka.JSON,
+		Tests: []TestQuery{
+			{
+				Query:        "Extract(Sort(All(), field=name), Rows(name), Rows(age), Rows(hobbies))",
+				ExpectedResp: `{"results":[{"fields":[{"name":"name","type":"string"},{"name":"age","type":"int64"},{"name":"hobbies","type":"[]string"}],"columns":[{"column":1,"rows":["a",20,["hob2","hob1"]]},{"column":2,"rows":["b",21,["hob2","hob3"]]},{"column":3,"rows":["c",22,["hob3","hob4"]]},{"column":4,"rows":["d",23,["hob4","hob5"]]},{"column":5,"rows":["e",24,["hob5","hob6"]]},{"column":6,"rows":["f",26,["hob6","hob7"]]}]}]}`,
+			},
+		},
+		CreateTableStmt: "(_id ID, name String, age Int, hobbies StringSet)",
+	},
+	{ // string keys
+		ConfigFile: "config01.toml",
+		DataFile:   "data00.json",
+		Encode:     kafka.JSON,
+		Tests: []TestQuery{
+			{
+				Query:        "Extract(Sort(All(), field=name), Rows(name), Rows(age), Rows(hobbies))",
+				ExpectedResp: `{"results":[{"fields":[{"name":"name","type":"string"},{"name":"age","type":"int64"},{"name":"hobbies","type":"[]string"}],"columns":[{"column":"1","rows":["a",20,["hob1","hob2"]]},{"column":"2","rows":["b",21,["hob2","hob3"]]},{"column":"3","rows":["c",22,["hob3","hob4"]]},{"column":"4","rows":["d",23,["hob4","hob5"]]},{"column":"5","rows":["e",24,["hob5","hob6"]]},{"column":"6","rows":["f",26,["hob6","hob7"]]}]}]}`,
+			},
+		},
+		CreateTableStmt: "(_id String, name String, age Int, hobbies StringSet)",
+	},
+	{ // two string keys
+		ConfigFile: "config02.toml",
+		DataFile:   "data00.json",
+		Encode:     kafka.JSON,
+		Tests: []TestQuery{
+			{
+				Query:        "Extract(Sort(All(), field=name), Rows(id), Rows(name), Rows(age), Rows(hobbies))",
+				ExpectedResp: `{"results":[{"fields":[{"name":"id","type":"string"},{"name":"name","type":"string"},{"name":"age","type":"int64"},{"name":"hobbies","type":"[]string"}],"columns":[{"column":"1|a","rows":["1","a",20,["hob1","hob2"]]},{"column":"2|b","rows":["2","b",21,["hob2","hob3"]]},{"column":"3|c","rows":["3","c",22,["hob3","hob4"]]},{"column":"4|d","rows":["4","d",23,["hob4","hob5"]]},{"column":"5|e","rows":["5","e",24,["hob5","hob6"]]},{"column":"6|f","rows":["6","f",26,["hob6","hob7"]]}]}]}`,
+			},
+		},
+		CreateTableStmt: "(_id String, id String, name String, age Int, hobbies StringSet)",
+	},
+	{ // string, id, and int
+		ConfigFile: "config03.toml",
+		DataFile:   "data00.json",
+		Encode:     kafka.JSON,
+		Tests: []TestQuery{
+			{
+				Query:        "Extract(Sort(All(), field=name), Rows(id), Rows(name), Rows(age), Rows(hobbies))",
+				ExpectedResp: `{"results":[{"fields":[{"name":"id","type":"uint64"},{"name":"name","type":"string"},{"name":"age","type":"int64"},{"name":"hobbies","type":"[]string"}],"columns":[{"column":"1|a|20","rows":[1,"a",20,["hob2","hob1"]]},{"column":"2|b|21","rows":[2,"b",21,["hob2","hob3"]]},{"column":"3|c|22","rows":[3,"c",22,["hob3","hob4"]]},{"column":"4|d|23","rows":[4,"d",23,["hob4","hob5"]]},{"column":"5|e|24","rows":[5,"e",24,["hob5","hob6"]]},{"column":"6|f|26","rows":[6,"f",26,["hob6","hob7"]]}]}]}`,
+			},
+		},
+		CreateTableStmt: "(_id String, id id, name String, age Int, hobbies StringSet)",
+	},
+}
+
+// TestKafkaRunner takes as input a slice of RunnerTest structs
+// For each RunnerTest, TestRunner:
+//  1. Creates a kafkaRunner based on a config file
+//  2. Reads data from a data file
+//  3. Encodes that data
+//  4. Writes the data to kafka
+//  5. Runs the kafkaRunner
+//  6. Confirms that the data was written to FeatureBase as expected
+func TestKafkaRunner(t *testing.T) {
+	for _, test := range kafkaRunnerTests {
+
+		// define final path to config and data files
+		var subpath string
+		switch encode := test.Encode; encode {
+		case kafka.JSON:
+			subpath = kafka.JSON
+		case kafka.Avro:
+			subpath = kafka.Avro
+		default:
+			t.Fatalf("unsupported encoding type")
+		}
+		KafkaConfig := confPathPrefix + subpath + "/" + test.ConfigFile
+		DataFile := dataPathPrefix + subpath + "/" + test.DataFile
+
+		// create command and kafka runner
+		fbsql := NewCommand(logger.StderrLogger)
+		setupKafkaRunnerCommand(t, fbsql)
+
+		// read config file so we can build a table before inserting data
+		config, err := kafka.ConfigStructFromFile(KafkaConfig)
+		if err != nil {
+			t.Fatalf("building kafka.Config struct: %s", err)
+		}
+
+		// create table to write to (kafka runner does not currently create tables)
+		createTable := strings.NewReader(fmt.Sprintf("CREATE TABLE %s %s;", config.Table, test.CreateTableStmt))
+		_, err = fbsql.Queryer.Query(fbsql.organizationID, fbsql.databaseID, createTable)
+		if err != nil {
+			t.Fatalf("creating table: %s", err)
+		}
+
+		// delete all tables after test (note this is after all tests not after each test)
+		defer func() {
+			dropTable := strings.NewReader(fmt.Sprintf("DROP TABLE %s;", config.Table))
+			_, err = fbsql.Queryer.Query(fbsql.organizationID, fbsql.databaseID, dropTable)
+			if err != nil {
+				t.Errorf("dropping table: %s", err)
+			}
+		}()
+
+		// create runner and hijack hosts based on test env (i.e. not from configuration file or other)
+		kafkaRunner, err := fbsql.newKafkaRunner(KafkaConfig)
+		if err != nil {
+			t.Fatalf("creating runner: %s", err)
+		}
+		kafkaRunner.FeaturebaseHosts = []string{pilosaHost}
+		kafkaRunner.SchemaRegistryURL = registryHost
+		kafkaRunner.KafkaBootstrapServers = []string{kafkaHost}
+		kafkaRunner.FeaturebaseGRPCHosts = []string{pilosaGrpcHost}
+
+		// read data that will go to kafka
+		var records []map[string]interface{}
+		records, err = recordsFromFile(DataFile)
+		if err != nil {
+			t.Errorf("reading raw messages from data file: %s", err)
+		}
+
+		// write data to kafka
+		var messages [][]byte
+		switch encode := test.Encode; encode {
+		case kafka.JSON:
+			messages, err = encodeJSONMessages(records)
+		case kafka.Avro:
+			t.Errorf("unsupported encoding type")
+		default:
+			t.Errorf("unsupported encoding type")
+		}
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		err = internal.ProduceMessages(kafkaRunner.Topics[0], strings.Join(kafkaRunner.KafkaBootstrapServers, ","), messages)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer func() {
+			// and delete the topic
+			err = internal.DeleteTopic(kafkaRunner.Topics[0], strings.Join(kafkaRunner.KafkaBootstrapServers, ","))
+			if err != nil {
+				t.Fatal(err)
+			}
+		}()
+
+		// Run it!!
+		kafkaRunner.MaxMsgs = uint64(len(messages))
+		err = kafkaRunner.Run()
+		if err != nil {
+			t.Fatalf("running runner: %s", err)
+		}
+
+		// and check the results...
+		for _, q := range test.Tests {
+
+			client := &http.Client{}
+			var data = strings.NewReader(q.Query)
+			req, err := http.NewRequest("POST", fmt.Sprintf("http://localhost:10101/index/%s/query", kafkaRunner.Index), data)
+			if err != nil {
+				t.Fatal(err)
+			}
+			req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+			resp, err := client.Do(req)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer resp.Body.Close()
+			respBody, err := io.ReadAll(resp.Body)
+			if err != nil {
+				t.Fatal(err)
+			}
+			require.JSONEq(t, q.ExpectedResp, string(respBody))
+		}
+
+	}
+}
+
+// Reads lines from a file and builds an in memory structure This function
+// expects the file formated as new line delimited JSON
+func recordsFromFile(pathToRecords string) (records []map[string]interface{}, err error) {
+	var data map[string]interface{}
+
+	recordsFile, err := os.Open(pathToRecords)
+	if err != nil {
+		return nil, fmt.Errorf("opening records file: %s", err)
+	}
+	defer recordsFile.Close()
+
+	s := bufio.NewScanner(recordsFile)
+	for s.Scan() {
+		err := json.Unmarshal(s.Bytes(), &data)
+		if err != nil {
+			return nil, fmt.Errorf("unmarshal json: %s", err)
+		}
+		records = append(records, data)
+		data = make(map[string]interface{})
+	}
+	return records, nil
+}
+
+// Mainly emmulate what happens in cli/cli.go Command.run() Unfortunately we
+// cannot just use that function because we need to create our own kafkaRunner
+// and update hosts and ports based on the test environment
+func setupKafkaRunnerCommand(t *testing.T, fbsql *Command) {
+	fbsql.Config.Host = strings.Split(pilosaHost, ":")[0]
+	fbsql.Config.Port = strings.Split(pilosaHost, ":")[1]
+	if err := fbsql.setupConfig(); err != nil {
+		t.Fatalf("setting up config: %s", err)
+	}
+
+	// Check to see if Command needs to run in non-interactive mode.
+	if len(fbsql.Commands) > 0 ||
+		len(fbsql.Files) > 0 ||
+		fbsql.Config.KafkaConfig != "" ||
+		fbsql.Config.CSV {
+		fbsql.nonInteractiveMode = true
+	}
+
+	if err := fbsql.setupClient(); err != nil {
+		t.Fatalf("setting up client: %s", err)
+	}
+
+	// Print the connection info.
+	if !fbsql.nonInteractiveMode {
+		fbsql.printConnInfo()
+	}
+
+	if err := fbsql.connectToDatabase(fbsql.database); err != nil {
+		t.Fatalf("connecting to database: %s", err)
+		// We intentionally do not return err here.
+	}
+}
+
+// Take a slice of JSON objects and convert them to a slice of byte slices that
+// can be sent to and stored in kafka
+func encodeJSONMessages(records []map[string]interface{}) ([][]byte, error) {
+	var messages [][]byte
+	for _, record := range records {
+		encoded, err := json.Marshal(record)
+		if err != nil {
+			return nil, errors.Errorf("marshaling records to JSON: %s", err)
+		}
+		messages = append(messages, encoded)
+	}
+	return messages, nil
+}

cli/internal/kafka_util.go

@@ -0,0 +1,74 @@
+package internal
+
+import (
+	"context"
+
+	"github.com/confluentinc/confluent-kafka-go/kafka"
+	"github.com/featurebasedb/featurebase/v3/errors"
+)
+
+func ProduceMessages(topic string, bootstrapServers string, messages [][]byte) error {
+
+	// create producer, defer closing
+	p, err := kafka.NewProducer(&kafka.ConfigMap{"bootstrap.servers": bootstrapServers})
+	if err != nil {
+		return errors.Errorf("creating producer: %s", err)
+	}
+	defer p.Close()
+
+	// create admin client needed to create topic, defer closing
+	ac, err := kafka.NewAdminClient(&kafka.ConfigMap{"bootstrap.servers": bootstrapServers})
+	if err != nil {
+		return errors.Errorf("creating admin client: %s", err)
+	}
+	defer ac.Close()
+
+	// create a topic
+	var ts = []kafka.TopicSpecification{
+		{
+			Topic:             topic,
+			NumPartitions:     1,
+			ReplicationFactor: 1,
+		},
+	}
+	ac.CreateTopics(context.Background(), ts, nil)
+	// caller must delete if needed
+
+	for _, message := range messages {
+		err := p.Produce(&kafka.Message{
+			TopicPartition: kafka.TopicPartition{Topic: &topic, Partition: kafka.PartitionAny},
+			Value:          message,
+		}, nil)
+		if err != nil {
+			return errors.Errorf("producing messages: %s", err)
+		}
+		p.Flush(3000)
+	}
+
+	return nil
+
+}
+
+func DeleteTopic(topic string, bootstrapServers string) error {
+	// create admin client needed to create topic, defer closing
+	ac, err := kafka.NewAdminClient(&kafka.ConfigMap{"bootstrap.servers": bootstrapServers})
+	if err != nil {
+		return errors.Errorf("creating admin client: %s", err)
+	}
+	defer ac.Close()
+
+	// delete a topic, catch any errors
+	results, err := ac.DeleteTopics(context.Background(), []string{topic}, nil)
+	if err != nil {
+		return errors.Errorf("deleting topic: %s", err)
+	}
+
+	for _, result := range results {
+		if result.Error.Code() != kafka.ErrNoError {
+			return errors.Errorf("fatal error deleting topic: %s", result.String())
+		}
+	}
+
+	return nil
+
+}

cli/kafka.go

@@ -1,35 +1,22 @@
 package cli
 
 import (
-	"fmt"
-
 	"github.com/featurebasedb/featurebase/v3/cli/batch"
 	"github.com/featurebasedb/featurebase/v3/cli/kafka"
 	"github.com/featurebasedb/featurebase/v3/errors"
-	"github.com/spf13/viper"
 )
 
 func (cmd *Command) newKafkaRunner(cfgFile string) (*kafka.Runner, error) {
-	// Read the kafka config file.
-	v := viper.New()
-	v.SetConfigFile(cfgFile)
-	v.SetConfigType("toml")
-	err := v.ReadInConfig()
-	if err != nil {
-		return nil, fmt.Errorf("error reading configuration file '%s': %v", cfgFile, err)
-	}
 
-	cfg := kafka.Config{}
-	if err := v.Unmarshal(&cfg); err != nil {
-		return nil, errors.Wrap(err, "unmarshalling config")
+	cfg, err := kafka.ConfigStructFromFile(cfgFile)
+	if err != nil {
+		return nil, err
 	}
 
 	if err := kafka.ValidateConfig(cfg); err != nil {
 		return nil, errors.Wrap(err, "validating config")
 	}
 
-	// Create a new config with defaults.
-
 	// Look up fields based on table provided in the config.
 	wqr, err := cmd.executeQuery(newRawQuery("SHOW COLUMNS FROM " + cfg.Table))
 	if err != nil {
@@ -57,7 +44,7 @@ func (cmd *Command) newKafkaRunner(cfgFile string) (*kafka.Runner, error) {
 		return nil, errors.Wrap(err, "cleaning config")
 	}
 
-	flds, err := kafka.ConfigToFields(cfg)
+	flds, err := kafka.ConfigToFields(cfg, idkCfg.PrimaryKeys)
 	if err != nil {
 		return nil, errors.Wrap(err, "getting fields from config")
 	}

cli/kafka/config.go

@@ -8,6 +8,14 @@ import (
 	"github.com/featurebasedb/featurebase/v3/dax"
 	"github.com/featurebasedb/featurebase/v3/idk"
 	"github.com/pkg/errors"
+	"github.com/spf13/viper"
+)
+
+// Kafka message encoding types supported - changes here should be
+// propogated to the Config struct and ValidateConfig error message
+const (
+	JSON string = "json"
+	Avro string = "avro"
 )
 
 // Config is the user-facing configuration for kafka support in the CLI. This is
@@ -23,6 +31,8 @@ type Config struct {
 
 	Table  string  `mapstructure:"table" help:"Destination table name."`
 	Fields []Field `mapstructure:"fields"`
+
+	Encode string `mapstructure:"encode" help:"Encoding format (currently supported formats: avro, json)"`
 }
 
 // Field is a user-facing configuration field.
@@ -44,18 +54,73 @@ type ConfigForIDK struct {
 	BatchMaxStaleness time.Duration
 	Timeout           time.Duration
 
-	Table   string
-	IDField string
-	Fields  []idk.RawField
+	Table       string
+	IDField     string
+	PrimaryKeys []string
+	Fields      []idk.RawField
+
+	Encode string
+}
+
+// Generate a Config struct based on a configuration file
+// Default values for Config struct are defined here
+func ConfigStructFromFile(cfgFile string) (cfg Config, err error) {
+	// configure viper
+	v := viper.New()
+	v.SetConfigFile(cfgFile)
+	v.SetConfigType("toml")
+
+	// set defaults
+	v.SetDefault("hosts", []string{"localhost:9092"})
+	v.SetDefault("group", "default-featurebase-group")
+	v.SetDefault("batch-size", 1)
+	v.SetDefault("batch-max-staleness", 5*time.Second)
+	v.SetDefault("timeout", 5*time.Second)
+	v.SetDefault("encode", JSON)
+
+	// Read the kafka config file.
+	err = v.ReadInConfig()
+	if err != nil {
+		return cfg, fmt.Errorf("error reading configuration file '%s': %v", cfgFile, err)
+	}
+
+	if err := v.Unmarshal(&cfg); err != nil {
+		return cfg, errors.Wrap(err, "unmarshalling config")
+	}
+
+	return
+
 }
 
 // ValidateConfig validates the config is usable.
+// Note that different encoding methods require
+// different configurations
 func ValidateConfig(c Config) error {
+
+	// validate common
 	if c.Table == "" {
 		return errors.Errorf("table is required")
-	} else if len(c.Topics) == 0 {
+	}
+
+	if len(c.Topics) == 0 {
 		return errors.Errorf("at least one topic is required")
-	} else if len(c.Fields) > 0 {
+	}
+
+	// validate on a by encoding basis
+	switch c.Encode {
+	case JSON:
+		return validateConfigJSON(c)
+	case Avro:
+		return validateConfigAvro(c)
+	}
+
+	return nil
+
+}
+
+func validateConfigJSON(c Config) error {
+
+	if len(c.Fields) > 0 {
 		// We only need to do these checks if any fields are specified at all.
 		// If no fields are specified, that's ok because then we default to
 		// using fields based off the existing table.
@@ -70,46 +135,64 @@ func ValidateConfig(c Config) error {
 				if c.Fields[i].Name == "" {
 					return errors.Errorf("a name attribute (which isn't equal to \"\") should exist for all fields")
 				}
+				if c.Fields[i].SourceType == "" {
+					return errors.Errorf("a source-type attribute (which isn't equal to \"\") should exist for all fields")
+				}
 			}
-			if found != 1 {
-				return errors.Errorf("exactly one primary key field is required")
+			if found < 1 {
+				return errors.Errorf("at least one primary key field is required")
 			}
 		}
 	}
+
+	return nil
+}
+
+// Only primary key fields required
+func validateConfigAvro(c Config) error {
 	return nil
 }
 
 // ConvertConfig converts a Config to one that suitable for IDK.
 func ConvertConfig(c Config) (ConfigForIDK, error) {
-	// Set a default kafka host in case one isn't provided.
-	hosts := []string{"localhost:9092"}
-	if len(c.Hosts) > 0 {
-		hosts = c.Hosts
-	}
 
 	// Copy all the shared members from Config to ConfigForIDK.
 	out := ConfigForIDK{
-		Hosts:             hosts,
+		Hosts:             c.Hosts,
 		Group:             c.Group,
 		Topics:            c.Topics,
 		BatchSize:         c.BatchSize,
 		BatchMaxStaleness: c.BatchMaxStaleness,
 		Timeout:           c.Timeout,
 		Table:             c.Table,
+		Encode:            c.Encode,
 	}
 
 	if len(c.Fields) == 0 {
 		return out, errors.New("fields cannot be empty")
 	}
 
-	// rawFields wil be the same as c.Fields, but possibly enhanced.
+	// rawFields will be the same as c.Fields, but possibly enhanced.
 	rawFields := make([]idk.RawField, 0, len(c.Fields))
 
-	var foundPK bool
+	var stringKeys bool
+	primaryKeys := []string{}
 	for _, fld := range c.Fields {
+
+		// handle primary key
 		if fld.PrimaryKey {
-			out.IDField = fld.Name
-			foundPK = true
+			switch keyType := fld.SourceType; keyType {
+			case dax.BaseTypeID:
+				// no-op
+			case dax.BaseTypeIDSet, dax.BaseTypeStringSet, dax.BaseTypeIDSetQ, dax.BaseTypeStringSetQ:
+				// key should be field that cannot contain multiple values
+				return out, errors.Errorf("Invalid")
+			default:
+				// unless
+				stringKeys = true
+
+			}
+			primaryKeys = append(primaryKeys, fld.Name)
 		}
 
 		typ, quals, err := dax.SplitFieldType(fld.SourceType)
@@ -149,8 +232,16 @@ func ConvertConfig(c Config) (ConfigForIDK, error) {
 
 		rawFields = append(rawFields, rawFld)
 	}
-	if !foundPK {
+
+	// Should have at least one primary key.  If there is more than one primary key OR
+	// using a string field as they key then use string keys (i.e. table will be keyed)
+	// Else, use ids (i.e. table will not be keyed)
+	if len(primaryKeys) < 1 {
 		return out, errors.New("primary-key not found in fields")
+	} else if stringKeys || len(primaryKeys) > 1 {
+		out.PrimaryKeys = primaryKeys
+	} else {
+		out.IDField = primaryKeys[0]
 	}
 
 	out.Fields = rawFields
@@ -160,13 +251,17 @@ func ConvertConfig(c Config) (ConfigForIDK, error) {
 
 // ConfigToFields returns a list of *dax.Field based on the IDField and Fields
 // in the Config.
-func ConfigToFields(c Config) ([]*dax.Field, error) {
+func ConfigToFields(c Config, primaryKeys []string) ([]*dax.Field, error) {
 	// We don't know if a primary key will be found, so we can't set the
 	// capacity to `len(c.Fields)-1`.
 	out := make([]*dax.Field, 0, len(c.Fields))
 
 	for _, fld := range c.Fields {
-		if fld.PrimaryKey {
+		// When we have a single primary key, don't also store that value as a
+		// field in FeatureBase. However, when we have more than one primary
+		// key, store all the values used for the compound key as fields in
+		// FeatureBase.
+		if fld.PrimaryKey && len(primaryKeys) < 2 {
 			continue
 		}
 		typ, quals, err := dax.SplitFieldType(fld.SourceType)

cli/kafka/config_test.go

@@ -0,0 +1,39 @@
+package kafka
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type configStructFromFileTest struct {
+	configFilePath string
+	expectedStruct string
+}
+
+func TestConfigStructFromFile(t *testing.T) {
+	tests := []configStructFromFileTest{
+		{ // confirm struct fields are being set in general
+			configFilePath: "./config_test_data/config01.toml",
+			expectedStruct: `{"Hosts":["kafka:9090"],"Group":"testGroup","Topics":["testTopic"],"BatchSize":35,"BatchMaxStaleness":25000000000,"Timeout":16000000000,"Table":"testTable","Fields":[{"Name":"id","SourceType":"id","SourcePath":null,"PrimaryKey":true},{"Name":"name","SourceType":"string","SourcePath":["test","path"],"PrimaryKey":false},{"Name":"age","SourceType":"int","SourcePath":null,"PrimaryKey":false},{"Name":"hobbies","SourceType":"stringset","SourcePath":null,"PrimaryKey":false}],"Encode":"json"}`,
+		},
+		{ // confirm defaults are being set
+			configFilePath: "./config_test_data/config00.toml",
+			expectedStruct: `{"Hosts":["localhost:9092"],"Group":"default-featurebase-group","Topics":["topic00"],"BatchSize":1,"BatchMaxStaleness":5000000000,"Timeout":5000000000,"Table":"table00","Fields":[{"Name":"id","SourceType":"string","SourcePath":null,"PrimaryKey":true},{"Name":"name","SourceType":"string","SourcePath":null,"PrimaryKey":false},{"Name":"age","SourceType":"int","SourcePath":null,"PrimaryKey":false},{"Name":"hobbies","SourceType":"stringset","SourcePath":null,"PrimaryKey":false}],"Encode":"json"}`,
+		},
+	}
+
+	for _, test := range tests {
+		cfg, err := ConfigStructFromFile(test.configFilePath)
+		if err != nil {
+			t.Fatalf("getting config struct from file: %s", err)
+		}
+		json, err := json.Marshal(cfg)
+		if err != nil {
+			t.Fatalf("marshaling config struct: %s", err)
+		}
+		require.JSONEq(t, test.expectedStruct, string(json))
+	}
+
+}

cli/kafka/config_test_data/config00.toml

@@ -0,0 +1,20 @@
+## confirm defaults are being set correctly
+topics = "topic00"
+table = "table00"
+
+[[fields]]
+name		= "id"
+source-type	= "string"
+primary-key	= true
+
+[[fields]]
+name		= "name"
+source-type	= "string"
+
+[[fields]]
+name		= "age"
+source-type	= "int"
+
+[[fields]]
+name		= "hobbies"
+source-type	= "stringset"

cli/kafka/config_test_data/config01.toml

@@ -0,0 +1,26 @@
+hosts = ["kafka:9090"]
+group = "testGroup"
+topics = "testTopic"
+table = "testTable"
+batch-size = 35
+batch-max-staleness = "25s"
+timeout = "16s"
+encode = "json"
+
+[[fields]]
+name		= "id"
+source-type	= "id"
+primary-key	= true
+
+[[fields]]
+name		= "name"
+source-type	= "string"
+source-path = ["test", "path"]
+
+[[fields]]
+name		= "age"
+source-type	= "int"
+
+[[fields]]
+name		= "hobbies"
+source-type	= "stringset"

cli/kafka/runner.go

@@ -7,7 +7,9 @@ import (
 	fbbatch "github.com/featurebasedb/featurebase/v3/batch"
 	"github.com/featurebasedb/featurebase/v3/errors"
 	"github.com/featurebasedb/featurebase/v3/idk"
-	"github.com/featurebasedb/featurebase/v3/idk/kafka_static"
+	"github.com/featurebasedb/featurebase/v3/idk/common"
+	"github.com/featurebasedb/featurebase/v3/idk/kafka"
+	"github.com/featurebasedb/featurebase/v3/idk/kafka_sasl"
 	"github.com/featurebasedb/featurebase/v3/logger"
 )
 
@@ -15,53 +17,90 @@ import (
 // idk.kafka_static.Main in that it embeds idk.Main and contains additional
 // functionality specific to its use case.
 type Runner struct {
-	idk.Main   `flag:"!embed"`
-	KafkaHosts []string       `help:"Comma separated list of host:port pairs for Kafka."`
-	Group      string         `help:"Kafka group."`
-	Topics     []string       `help:"Kafka topics to read from."`
-	Timeout    time.Duration  `help:"Time to wait for more records from Kafka before flushing a batch. 0 to disable."`
-	Header     []idk.RawField `help:"Header configuration."`
+	idk.Main             `flag:"!embed"`
+	idk.ConfluentCommand `flag:"!embed"`
+	Hosts                []string       `help:"Comma separated list of host:port pairs for Kafka."`
+	Group                string         `help:"Kafka group."`
+	Topics               []string       `help:"Kafka topics to read from."`
+	Timeout              time.Duration  `help:"Time to wait for more records from Kafka before flushing a batch. 0 to disable."`
+	Header               []idk.RawField `help:"Header configuration."`
 }
 
+// Configure and return *runner with common elements of all kafka runners
 func NewRunner(cfg ConfigForIDK, batcher fbbatch.Batcher, logWriter io.Writer) *Runner {
 	idkMain := idk.NewMain()
 	idkMain.IDField = cfg.IDField
+	idkMain.PrimaryKeyFields = cfg.PrimaryKeys
 	idkMain.Index = cfg.Table
 	idkMain.Batcher = batcher
 	idkMain.BatchSize = cfg.BatchSize
 	idkMain.BatchMaxStaleness = cfg.BatchMaxStaleness
 	idkMain.SetBasic()
 	idkMain.SetLog(logger.NewStandardLogger(logWriter))
+	idkMain.OffsetMode = true
+	idkMain.Namespace = "sql_kafka_runner"
+	idkMain.Pprof = "" // don't initialize pprof until we actually use it in tests
 
 	kr := &Runner{
-		Main:       *idkMain,
-		KafkaHosts: cfg.Hosts,
-		Group:      cfg.Group,
-		Topics:     cfg.Topics,
-		Header:     cfg.Fields,
-		Timeout:    cfg.Timeout,
+		Main:    *idkMain,
+		Hosts:   cfg.Hosts,
+		Group:   cfg.Group,
+		Topics:  cfg.Topics,
+		Header:  cfg.Fields,
+		Timeout: cfg.Timeout,
 	}
-	kr.OffsetMode = true
-	kr.Main.Namespace = "cli_kafka_runner"
-	kr.Main.Pprof = "" // don't initialize pprof until we actually use it in tests
-	kr.NewSource = func() (idk.Source, error) {
-		source := kafka_static.NewSource()
-		source.Hosts = kr.KafkaHosts
-		source.Group = kr.Group
-		source.Topics = kr.Topics
-		source.Log = kr.Main.Log()
-		// source.TLS = m.KafkaTLS
-		source.Timeout = kr.Timeout
-		// source.SkipOld = m.SkipOld
-		source.HeaderFields = kr.Header
-		// source.S3Region = m.S3Region
-		// source.AllowMissingFields = m.AllowMissingFields
 
-		err := source.Open()
+	// NewSource should be set based on the encoding of the source (e.g. JSON, Avro)
+	if cfg.Encode == Avro {
+		kr.GetAvroNewSource()
+	} else if cfg.Encode == JSON {
+		kr.GetJSONNewSource()
+	}
+
+	return kr
+}
+
+func (r *Runner) GetJSONNewSource() {
+	r.NewSource = func() (idk.Source, error) {
+		source := kafka_sasl.NewSource()
+		source.KafkaBootstrapServers = r.Hosts
+		source.Group = r.Group
+		source.Topics = r.Topics
+		source.Log = r.Main.Log()
+		source.Timeout = r.Timeout
+		source.HeaderFields = r.Header
+		cfg, err := common.SetupConfluent(&r.ConfluentCommand)
+		if err != nil {
+			return nil, err
+		}
+		source.ConfigMap = cfg
+
+		err = source.Open()
+		if err != nil {
+			return nil, errors.Wrap(err, "opening source")
+		}
+		return source, nil
+	}
+}
+
+func (r *Runner) GetAvroNewSource() {
+	r.NewSource = func() (idk.Source, error) {
+		source := kafka.NewSource()
+		source.KafkaBootstrapServers = r.Hosts
+		source.Group = r.Group
+		source.Topics = r.Topics
+		source.Log = r.Main.Log()
+		source.Timeout = r.Timeout
+		cfg, err := common.SetupConfluent(&r.ConfluentCommand)
+		if err != nil {
+			return nil, err
+		}
+		source.ConfigMap = cfg
+
+		err = source.Open()
 		if err != nil {
 			return nil, errors.Wrap(err, "opening source")
 		}
 		return source, nil
 	}
-	return kr
 }

cli/kafka/runner_test_data/config/json/config00.toml

@@ -0,0 +1,25 @@
+hosts = ["localhost:9092"]
+group = "grp"
+topics = "topic00"
+table = "table00"
+batch-size = 1
+batch-max-staleness = "5s"
+timeout = "5s"
+encode = "json"
+
+[[fields]]
+name		= "id"
+source-type	= "id"
+primary-key	= true
+
+[[fields]]
+name		= "name"
+source-type	= "string"
+
+[[fields]]
+name		= "age"
+source-type	= "int"
+
+[[fields]]
+name		= "hobbies"
+source-type	= "stringset"

cli/kafka/runner_test_data/config/json/config01.toml

@@ -0,0 +1,25 @@
+hosts = ["localhost:9092"]
+group = "grp"
+topics = "topic01"
+table = "table01"
+batch-size = 1
+batch-max-staleness = "5s"
+timeout = "5s"
+encode = "json"
+
+[[fields]]
+name		= "id"
+source-type	= "string"
+primary-key	= true
+
+[[fields]]
+name		= "name"
+source-type	= "string"
+
+[[fields]]
+name		= "age"
+source-type	= "int"
+
+[[fields]]
+name		= "hobbies"
+source-type	= "stringset"

cli/kafka/runner_test_data/config/json/config02.toml

@@ -0,0 +1,26 @@
+hosts = ["localhost:9092"]
+group = "grp"
+topics = "topic02"
+table = "table02"
+batch-size = 1
+batch-max-staleness = "5s"
+timeout = "5s"
+encode = "json"
+
+[[fields]]
+name		= "id"
+source-type	= "string"
+primary-key	= true
+
+[[fields]]
+name		= "name"
+source-type	= "string"
+primary-key	= true
+
+[[fields]]
+name		= "age"
+source-type	= "int"
+
+[[fields]]
+name		= "hobbies"
+source-type	= "stringset"

cli/kafka/runner_test_data/config/json/config03.toml

@@ -0,0 +1,27 @@
+hosts = ["localhost:9092"]
+group = "grp"
+topics = "topic03"
+table = "table03"
+batch-size = 1
+batch-max-staleness = "5s"
+timeout = "5s"
+encode = "json"
+
+[[fields]]
+name		= "id"
+source-type	= "id"
+primary-key	= true
+
+[[fields]]
+name		= "name"
+source-type	= "string"
+primary-key	= true
+
+[[fields]]
+name		= "age"
+source-type	= "int"
+primary-key	= true
+
+[[fields]]
+name		= "hobbies"
+source-type	= "stringset"

cli/kafka/runner_test_data/data/json/data00.json

@@ -0,0 +1,6 @@
+{"id": 1, "name": "a", "age": 20, "hobbies": "hob1,hob2"}
+{"id": 2, "name": "b", "age": 21, "hobbies": "hob3,hob2"}
+{"id": 3, "name": "c", "age": 22, "hobbies": "hob3,hob4"}
+{"id": 4, "name": "d", "age": 23, "hobbies": "hob4,hob5"}
+{"id": 5, "name": "e", "age": 24, "hobbies": "hob5,hob6"}
+{"id": 6, "name": "f", "age": 26, "hobbies": "hob6,hob7"}

idk/kafka_sasl/source.go

@@ -26,9 +26,16 @@ type Source struct {
 	Log                logger.Logger
 	Timeout            time.Duration
 	SkipOld            bool
-	Header             string
 	AllowMissingFields bool
 
+	// Header is a file referencing a file containing JSON header configuration.
+	Header string
+
+	// HeaderFields can be provided instead of Header. It is a slice of
+	// RawFields which will be marshalled and parsed the same way a JSON object
+	// in Header would be. It is used only if a Header is not provided.
+	HeaderFields []idk.RawField
+
 	schema []idk.Field
 	paths  idk.PathTable
 
@@ -193,14 +200,24 @@ func (r *Record) Data() []interface{} {
 
 // Open initializes the kafka source.
 func (s *Source) Open() error {
-	if len(s.Header) == 0 {
-		return errors.New("needs header specification file")
+	if len(s.Header) == 0 && len(s.HeaderFields) == 0 {
+		return errors.New("needs header specification file (file or fields)")
 	}
 
-	headerData, err := os.ReadFile(s.Header)
-	if err != nil {
-		return errors.Wrap(err, "reading header file")
+	var headerData []byte
+	var err error
+	if s.Header != "" {
+		headerData, err = os.ReadFile(s.Header)
+		if err != nil {
+			return errors.Wrap(err, "reading header file")
+		}
+	} else {
+		headerData, err = json.Marshal(s.HeaderFields)
+		if err != nil {
+			return errors.Wrap(err, "marshalling header fields")
+		}
 	}
+
 	schema, paths, err := idk.ParseHeader(headerData)
 	if err != nil {
 		return errors.Wrap(err, "processing header")