scriptoriumadmin/featurebase
1
0
Fork 0
mirror of https://github.com/featurebasedb/featurebase.git synced 2026-08-28 10:54:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

cli/kafka updates

Browse source
This commit is contained in:
jacob 2023-04-03 11:15:32 -05:00
parent c6dc1f0aed
commit aefbc6493b
27 changed files with 926 additions and 555 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
cli/batch/sql.go
View file

@ -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]

488
cli/cli_kafka_integration_test.go Normal file
View file

@ -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 <TABLE> 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 <TABLE> 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 <TABLE> 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 <TABLE> 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 <TABLE> 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, "<TABLE>", 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
}

343
cli/cli_kafka_runner_test.go
View file

@ -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
}

74
cli/internal/kafka_util.go
View file

@ -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
}

2
cli/kafka.go
View file

@ -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
}

84
cli/kafka/config.go
View file

@ -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")

26
cli/kafka/config_test.go
View file

@ -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)
}

24
cli/kafka/runner.go
View file

@ -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 {

0
cli/kafka/config_test_data/config00.toml → cli/kafka/testdata/config/config00.toml vendored
View file

0
cli/kafka/config_test_data/config01.toml → cli/kafka/testdata/config/config01.toml vendored
View file

27
cli/kafka/testdata/config/config02.toml vendored Normal file
View file

@ -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"

30
cli/kafka/testdata/config/config03.toml vendored Normal file
View file

@ -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"

3
cli/kafka/runner_test_data/config/json/config00.toml → cli/kafka/testdata/runner/config/config00.toml vendored
View file

@ -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"

3
cli/kafka/runner_test_data/config/json/config01.toml → cli/kafka/testdata/runner/config/config01.toml vendored
View file

@ -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"

3
cli/kafka/runner_test_data/config/json/config02.toml → cli/kafka/testdata/runner/config/config02.toml vendored
View file

@ -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"

3
cli/kafka/runner_test_data/config/json/config03.toml → cli/kafka/testdata/runner/config/config03.toml vendored
View file

@ -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"

15
cli/kafka/testdata/runner/config/config04.toml vendored Normal file
View file

@ -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

27
cli/kafka/testdata/runner/config/config05.toml vendored Normal file
View file

@ -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"

2
cli/kafka/runner_test_data/data/json/data00.json → cli/kafka/testdata/runner/data/data00.json vendored
View file

@ -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"}

10
cli/kafka/testdata/runner/data/data01.json vendored Normal file
View file

@ -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"}}

6
cli/kafka/testdata/runner/data/data02.json vendored Normal file
View file

@ -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"}

32
cli/kafka/testdata/runner/schema/schema01.json vendored Normal file
View file

@ -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"}
]
}

24
idk/ingest.go
View file

@ -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
}
}
}

7
idk/kafka/cmd.go
View file

@ -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

75
idk/kafka/source.go
View file

@ -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 {

123
idk/kafka_sasl/source.go
View file

@ -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")
}
}

23
wire_response.go
View file

@ -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