elasticstream/client.go

package elasticstream

import (
	"fmt"
	"log"

	"github.com/boltdb/bolt"
	"github.com/elastic/go-elasticsearch/v8"
)

type Client struct {
	es      *elasticsearch.Client
	config  *Config
	ch      chan Data
	db      *bolt.DB
	offsets map[string]int
}

func NewClient(config *Config) (*Client, error) {
	client := &Client{
		config:  config,
		offsets: make(map[string]int),
	}
	return client, nil
}

func (c *Client) Open() error {
	// Open a connection with ElasticSearch
	cfg := elasticsearch.Config{
		Addresses: []string{
			c.config.Host,
		},
	}

	var err error
	c.es, err = elasticsearch.NewClient(cfg)
	if err != nil {
		return err
	}

	// create a buffer channel
	c.ch = make(chan Data, 25)

	// open bolt db
	c.db, err = bolt.Open(c.config.DBPath, 0666, nil)
	if err != nil {
		return err
	}
	// defer db.Close()

	for _, index := range c.config.Indexes {
		offset, err := getOffset(c.db, index)
		if err != nil {
			log.Printf("getOffset(%s) err:%s\n", index, err)
			continue
		}

		c.offsets[index] = offset

		NewWorker(c, index, offset, c.config.BatchSize)
	}

	return nil
}

func (c *Client) Read() (*Data, error) {

	data, ok := <-c.ch
	if !ok {
		return nil, fmt.Errorf("error reading data from channel")
	}

	index, ok := data.Header["index"]
	if !ok {
		return nil, fmt.Errorf("error index not found in data header")
	}

	offset, ok := c.offsets[index]
	if !ok {
		return nil, fmt.Errorf("error offset index not found")
	}

	c.offsets[index] = offset + 1

	err := setOffset(c.db, index, c.offsets[index])
	if err != nil {
		return nil, err
	}

	return &data, nil
}

// func (c *Client) Ack(ctx context.Context, position int) error {
// 	return nil
// }

// close the client
func (c *Client) Teardown() error {
	return nil
}