elasticstream/worker.go

package elasticstream

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	"github.com/elastic/go-elasticsearch/esapi"
	"github.com/elastic/go-elasticsearch/v8"
)

type Worker struct {
	client *elasticsearch.Client
	index  string // name of the indexes worker pulls data from
	from   int    // from where to start read data
	size   int    // batch size
	buffer chan Data
}

func NewWorker(client *elasticsearch.Client, index string, from, size int, buffer chan Data) *Worker {
	w := &Worker{
		client: client,
		index:  index,
		from:   from,
		size:   size,
		buffer: buffer,
	}

	go w.start()

	return w
}

func (w *Worker) start() {

}

// search is calling Elastic Search search API
func search(client *elasticsearch.Client, from, size int, index string) ([]Data, error) {
	query := fmt.Sprintf(`{
		"query": {
			"match_all": {}
		}
	}`)

	// Create the search request
	req := esapi.SearchRequest{
		Index: []string{index},
		Body:  strings.NewReader(query),
		From:  from,
		Size:  size,
	}

	// Perform the request
	res, err := req.Do(context.Background(), client)
	if err != nil {
		return nil, fmt.Errorf("error getting response: %s", err)
	}
	defer res.Body.Close()

	if res.IsError() {
		return nil, fmt.Errorf("error: %s", res.String())
	}

	// Parse the response
	var result struct {
		Hits struct {
			Hits []struct {
				Source map[string]interface{} `json:"_source"`
			} `json:"hits"`
		} `json:"hits"`
	}

	if err := json.NewDecoder(res.Body).Decode(&result); err != nil {
		return nil, fmt.Errorf("error parsing the response body: %s", err)
	}

	// Collect the records
	newRecords := make([]map[string]interface{}, len(result.Hits.Hits))
	for i, hit := range result.Hits.Hits {
		newRecords[i] = hit.Source
	}

	header := map[string]string{"index": index}

	var records []Data
	for _, v := range newRecords {
		data := Data{
			Header:  header,
			Payload: v,
		}
		records = append(records, data)
	}

	return records, nil
}