featurebase/client/ingest_api_batch.go

package client

import (
	"time"

	"github.com/featurebasedb/featurebase/v3/logger"
	"github.com/pkg/errors"
)

// NewIngestAPIBatch creates an alternate implementation of
// RecordBatch which exists to aid in testing the new Ingest API and
// is likely far slower than the Batch.
func NewIngestAPIBatch(client *Client, size int, logger logger.Logger, fields []*Field) *ingestAPIBatch {
	if len(fields) == 0 {
		return nil
	}

	return &ingestAPIBatch{
		client:    client,
		log:       logger,
		fields:    fields,
		keyed:     fields[0].index.Opts().Keys(),
		index:     fields[0].index.Name(),
		batchSize: size,

		recordsK: make(map[string]map[string]interface{}),
		records:  make(map[uint64]map[string]interface{}),
	}
}

type ingestAPIBatch struct {
	client    *Client
	log       logger.Logger
	batchSize int

	fields []*Field
	keyed  bool
	index  string

	// map[recordKey][fieldName]value
	recordsK map[string]map[string]interface{}
	records  map[uint64]map[string]interface{}
}

func (b *ingestAPIBatch) Add(row Row) error {
	if len(row.Clears) > 0 {
		return errors.New("ingest api batch does not support clears")
	}
	values := make(map[string]interface{})
	for i, val := range row.Values {
		field := b.fields[i]
		// val can be string, uint64, int64, []string, []uint64, nil
		// TODO timestamp field might need special handling
		// TODO check that the Row.Clears field is only used for packed bools, and then issue a warning/error (in IDK) if the ingest API mode is used in conjunction w/ packed bools.
		if val == nil {
			continue
		}
		zero := QuantizedTime{}
		if field.Options().Type() == FieldTypeTime && row.Time != zero {
			timeq, err := row.Time.Time()
			if err != nil {
				return errors.Wrap(err, "parsing row time")
			}
			values[field.Name()] = map[string]interface{}{"time": timeq.Format(time.RFC3339), "values": val}
		} else {
			values[field.Name()] = val
		}
	}

	if b.keyed {
		switch rowID := row.ID.(type) {
		case string:
			b.recordsK[rowID] = values
		case []byte:
			b.recordsK[string(rowID)] = values
		default:
			return errors.Errorf("unsupported rowID %v of type %[1]T, must be string, or []byte for keyed index", rowID)
		}
		if len(b.recordsK) >= b.batchSize {
			return ErrBatchNowFull
		}
	} else {
		rowID, ok := row.ID.(uint64)
		if !ok {
			return errors.Errorf("unsupported rowID %v of type %[1]T, must be uint64 for unkeyed index", row.ID)
		}
		b.records[rowID] = values
		if len(b.records) >= b.batchSize {
			return ErrBatchNowFull
		}
	}
	return nil
}

func (b *ingestAPIBatch) Import() error {
	if b.keyed {
		return b.importKeyed()
	}
	return b.importUnkeyed()
}

func (b *ingestAPIBatch) importKeyed() error {
	req := []map[string]interface{}{
		{
			"action":  "set",
			"records": b.recordsK,
		},
	}
	bod, err := b.client.IngestData(b.index, req)
	if err != nil {
		return errors.Wrapf(err, "importKeyed, body: %s", bod)
	}

	for k := range b.recordsK {
		delete(b.recordsK, k)
	}
	return nil
}

func (b *ingestAPIBatch) importUnkeyed() error {
	req := []map[string]interface{}{
		{
			"action":  "set",
			"records": b.records,
		},
	}
	bod, err := b.client.IngestData(b.index, req)
	if err != nil {
		return errors.Wrapf(err, "importKeyed, body: %s", bod)
	}

	for v := range b.records {
		delete(b.records, v)
	}
	return nil
}

func (b *ingestAPIBatch) Len() int {
	if b.keyed {
		return len(b.recordsK)
	}
	return len(b.records)
}

func (b *ingestAPIBatch) Flush() error { return nil }