Switch fbsql from using kafka_static to kafka_sasl

Since we likely want to use the ConfluentCommand support in the future,
this uses kafka_sasl package instead of kafka_static. That allows us to
continue using non-avro headers (which the kafka package uses), but still
use ConfluentCommand.

This also adds the Timeout logic to the kafka_sasl package.

.gitlab/.gitlab-ci.yml

@@ -193,7 +193,7 @@ build fbsql amd64:
     - export SOURCE_DATE_EPOCH=$(git log -1 --pretty=%ct)
     - date
     - GOOS="linux" GOARCH="amd64" make docker-build-fbsql BUILD_CGO=1
-    - GOOS="darwin" GOARCH="amd64" make docker-build-fbsql
+    - GOOS="darwin" GOARCH="amd64" make docker-build-fbsql BUILD_CGO=1
   artifacts:
     paths:
       - ./build/fbsql_*
@@ -210,7 +210,7 @@ build fbsql arm64:
     - export SOURCE_DATE_EPOCH=$(git log -1 --pretty=%ct)
     - date
     - GOOS="linux" GOARCH="arm64" make docker-build-fbsql BUILD_CGO=1
-    - GOOS="darwin" GOARCH="arm64" make docker-build-fbsql
+    - GOOS="darwin" GOARCH="arm64" make docker-build-fbsql BUILD_CGO=1
   artifacts:
     paths:
       - ./build/fbsql_*

Dockerfile-fbsql

@@ -28,7 +28,7 @@ ARG SOURCE_DATE_EPOCH
 WORKDIR /featurebase/
 
 ENV SOURCE_DATE_EPOCH=${SOURCE_DATE_EPOCH}
-RUN make build-fbsql GO_BUILD_FLAGS="-mod=vendor ${GO_BUILD_FLAGS}" ${MAKE_FLAGS}
+RUN make build-fbsql-cgo GO_BUILD_FLAGS="-mod=vendor ${GO_BUILD_FLAGS}" ${MAKE_FLAGS}
 
 FROM ubuntu:20.04 as runner
 

Makefile

@@ -358,18 +358,22 @@ BUILD_CGO ?= 0
 
 # Build fbsql
 build-fbsql:
-	@echo GOOS=$(GOOS) GOARCH=$(GOARCH) uname -p=$(UNAME_P) build_cgo=$(BUILD_CGO)
-ifeq ($(BUILD_CGO), 0)
-	make build-fbsql-non-cgo
-endif
-ifeq ($(BUILD_CGO), 1)
-	make build-fbsql-cgo
-endif
-
-build-fbsql-non-cgo:
-	CGO_ENABLED=0 $(GO) build -ldflags $(LDFLAGS) $(GO_BUILD_FLAGS) -o fbsql ./cmd/fbsql
+	CGO_ENABLED=1 $(GO) build -ldflags $(LDFLAGS) $(GO_BUILD_FLAGS) -o fbsql ./cmd/fbsql
 
+# build-fbsql-cgo target is used in the Dockerfile-fbsql to build fbsql in CI.
 build-fbsql-cgo:
+ifeq ($(GOOS), linux)
+	$(MAKE) build-fbsql-cgo-linux
+endif
+ifeq ($(GOOS), darwin)
+	$(MAKE) build-fbsql-cgo-darwin
+endif
+
+build-fbsql-cgo-darwin:
+	@echo GOOS=$(GOOS) GOARCH=$(GOARCH) uname -p=$(UNAME_P) build_cgo=$(BUILD_CGO)
+	CGO_ENABLED=1 $(GO) build -ldflags $(LDFLAGS_STATIC) $(GO_BUILD_FLAGS) -o fbsql ./cmd/fbsql
+
+build-fbsql-cgo-linux:
 ifeq ($(GOARCH), arm64)
 	CGO_ENABLED=1 $(GO) build -tags dynamic $(GO_BUILD_FLAGS) -o fbsql ./cmd/fbsql
 endif

cli/kafka/runner.go

@@ -7,7 +7,8 @@ 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_sasl"
 	"github.com/featurebasedb/featurebase/v3/logger"
 )
 
@@ -44,9 +45,9 @@ func NewRunner(cfg ConfigForIDK, batcher fbbatch.Batcher, logWriter io.Writer) *
 	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 := kafka_sasl.NewSource()
 		source.Group = kr.Group
 		source.Topics = kr.Topics
 		source.Log = kr.Main.Log()
@@ -57,10 +58,20 @@ func NewRunner(cfg ConfigForIDK, batcher fbbatch.Batcher, logWriter io.Writer) *
 		// source.S3Region = m.S3Region
 		// source.AllowMissingFields = m.AllowMissingFields
 
-		err := source.Open()
-		if err != nil {
+		// Set up ConfluentCommand configuration.
+		confluentCommand := &idk.ConfluentCommand{
+			KafkaBootstrapServers: kr.KafkaHosts,
+		}
+		if cfg, err := common.SetupConfluent(confluentCommand); err != nil {
+			return nil, errors.Wrap(err, "setting up confluent command")
+		} else {
+			source.ConfigMap = cfg
+		}
+
+		if err := source.Open(); err != nil {
 			return nil, errors.Wrap(err, "opening source")
 		}
+
 		return source, nil
 	}
 	return kr

idk/kafka_sasl/source.go

@@ -21,12 +21,20 @@ import (
 // achieve concurrency, create multiple Sources.
 type Source struct {
 	idk.ConfluentCommand
-	Topics             []string
-	Group              string
-	Log                logger.Logger
-	Timeout            time.Duration
-	SkipOld            bool
-	Header             string
+	Topics  []string
+	Group   string
+	Log     logger.Logger
+	Timeout time.Duration
+	SkipOld 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
+
 	AllowMissingFields bool
 
 	schema []idk.Field
@@ -108,7 +116,21 @@ type recordWithError struct {
 }
 
 func (s *Source) fetch() recordWithError {
-	return <-s.recordChannel
+	ctx := context.Background()
+	if s.Timeout != 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, s.Timeout)
+		defer cancel()
+	}
+
+	var rec recordWithError
+	select {
+	case rec = <-s.recordChannel:
+	case <-ctx.Done():
+		rec.Err = ctx.Err()
+	}
+
+	return rec
 }
 
 func (s *Source) decodeMessage(buf []byte) ([]interface{}, error) {
@@ -193,14 +215,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 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")