featurebase/idk/cmd/ingester/main.go

package main

import (
	"fmt"
	"log"
	"net/http"
	"runtime/debug"

	pilosaclient "github.com/featurebasedb/featurebase/v3/client"
	"github.com/featurebasedb/featurebase/v3/idk"
	"github.com/featurebasedb/featurebase/v3/idk/api"
	"github.com/gorilla/mux"
	"github.com/jaffee/commandeer/pflag"
	"github.com/pkg/errors"
)

func main() {
	i := &ingester{
		PilosaHosts: []string{"localhost:10101"},
		HttpAddr:    "localhost:8080",
	}

	if err := pflag.LoadEnv(i, "INGESTER_", nil); err != nil {
		log.Fatal(err)
	}
	if i.DryRun {
		log.Printf("%+v\n", i)
		return
	}

	if err := i.Run(); err != nil {
		log.Fatal(err)
	}
}

// ingester is a HTTP handler which lets ingest data to pilosa over "programmatic ingest API".
// https://github.com/molecula/docs/blob/master/docs/internal/proposals/programmatic-ingest-api.md
type ingester struct {
	DryRun          bool     `help:"Dry run - just flag parsing."`
	PilosaHosts     []string `short:"" help:"Comma separated list of host:port pairs for FeatureBase."`
	PilosaGrpcHosts []string `short:"" help:"Comma separated list of host:port pairs for FeatureBase's gRPC endpoint."`
	BatchSize       uint     `short:"" help:"Batch size to use for ingest operations."`
	HttpAddr        string   `short:"" help:"HTTP address for Ingester."`

	client   *pilosaclient.Client
	router   *mux.Router
	tmplMain idk.Main

	AuthToken string `flag:"auth-token" help:"Authentication Token for FeatureBase"`
}

func (i *ingester) Run() error {
	client, err := pilosaclient.NewClient(i.PilosaHosts,
		pilosaclient.OptClientRetries(2),
		pilosaclient.OptClientTotalPoolSize(1000),
		pilosaclient.OptClientPoolSizePerRoute(400),
	)
	if err != nil {
		return errors.Wrap(err, "getting featurebase client")
	}
	defer client.Close()

	i.client = client

	i.router = mux.NewRouter()
	i.router.HandleFunc("/schema", i.postSchema).Methods("POST").Name("PostSchema")
	i.router.HandleFunc("/index/{index}", i.postIndex).Methods("POST").Name("PostIndex")

	m := *idk.NewMain()
	m.PilosaHosts = i.PilosaHosts
	m.PilosaGRPCHosts = i.PilosaGrpcHosts
	m.Pprof = ""
	m.Stats = ""
	m.PackBools = ""
	m.BatchSize = int(i.BatchSize)
	m.AuthToken = i.AuthToken
	i.tmplMain = m

	s := &http.Server{Addr: i.HttpAddr, Handler: i}
	return s.ListenAndServe()
}

// ServeHTTP handles an HTTP request.
func (i *ingester) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	defer func() {
		if err := recover(); err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			stack := debug.Stack()
			log.Printf("PANIC: %s\n%s", err, stack)
		}
	}()

	i.router.ServeHTTP(w, r)
}

// POST /schema
func (i *ingester) postSchema(w http.ResponseWriter, r *http.Request) {
	body := r.Body
	defer body.Close()

	schema, ifNotExists, err := api.DecodeSchema(body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	for idxName, idx := range schema.Indexes() {
		if err := i.client.CreateIndex(idx); err != nil {
			if errors.Is(err, pilosaclient.ErrIndexExists) {
				if ifNotExists[idxName] {
					continue
				}

				http.Error(w, err.Error(), http.StatusConflict)
				return
			}

			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		for fldName, fld := range idx.Fields() {
			if fldName == "_exists" {
				continue
			}

			if err := i.client.EnsureField(fld); err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
		}
	}

	// Following code is not really needed,
	// but it's kind of informative if we want to check what's the current pilosa schema.
	// If we decide to keep it, we'll need to improve "pilosa/client",
	// because with current implementation we can serialize just map of indices.
	schema, err = i.client.Schema()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
	w.Write([]byte(schema.String())) //nolint: errcheck
}

// POST /index/{index}
func (i *ingester) postIndex(w http.ResponseWriter, r *http.Request) {
	index := mux.Vars(r)["index"]

	schema, err := i.client.Schema()
	if err != nil {
		http.Error(w, fmt.Sprintf("fetching schema: %v", err), http.StatusInternalServerError)
		return
	}

	idx, ok := schema.Indexes()[index]
	if !ok {
		http.Error(w, fmt.Sprintf("index %q not found", index), http.StatusNotFound)
		return
	}

	switch ctype := r.Header.Get("Content-Type"); ctype {
	case "":
		http.Error(w, "missing Content-Type header", http.StatusUnsupportedMediaType)
		return

	case "application/json":
		err = api.IngestJSON(idx, i.tmplMain, r.Body)
		if err != nil {
			errCode := http.StatusInternalServerError
			if errors.As(err, &api.TypeError{}) || errors.As(err, &api.ErrDuplicateElement{}) {
				errCode = http.StatusBadRequest
			}
			http.Error(w, err.Error(), errCode)
			return
		}

	default:
		http.Error(w, fmt.Sprintf("unsupported Content-Type %q", ctype), http.StatusUnsupportedMediaType)
		return
	}

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("OK")) //nolint: errcheck
}