esgo2dump/xes/es7/read.go

package es7

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"time"

	elastic "github.com/elastic/go-elasticsearch/v7"
	"github.com/elastic/go-elasticsearch/v7/esapi"
	"github.com/loveuer/esgo2dump/internal/tool"
	"github.com/loveuer/esgo2dump/pkg/model"
	"github.com/samber/lo"
)

type streamer struct {
	ctx    context.Context
	client *elastic.Client
	index  string
	scroll string
}

// ReadData implements model.IO.
func (s *streamer) ReadData(limit int, query map[string]any, fields []string, sort []string) ([]map[string]any, error) {
	var (
		err    error
		qs     []func(*esapi.SearchRequest)
		resp   *esapi.Response
		result = new(model.ESResponseV7[map[string]any])
	)

	if s.scroll != "" {
		if resp, err = s.client.Scroll(
			s.client.Scroll.WithContext(tool.TimeoutCtx(s.ctx)),
			s.client.Scroll.WithScrollID(s.scroll),
			s.client.Scroll.WithScroll(35*time.Second),
		); err != nil {
			return nil, err
		}

		goto HandleResp
	}

	qs = []func(*esapi.SearchRequest){
		s.client.Search.WithContext(tool.TimeoutCtx(s.ctx)),
		s.client.Search.WithIndex(s.index),
		s.client.Search.WithSize(limit),
		s.client.Search.WithScroll(35 * time.Second),
	}

	if len(fields) > 0 {
		qs = append(qs, s.client.Search.WithSourceIncludes(fields...))
	}

	if len(sort) > 0 {
		qs = append(qs, s.client.Search.WithSort(sort...))
	}

	if len(query) > 0 {
		queryBs, err := json.Marshal(map[string]any{"query": query})
		if err != nil {
			return nil, err
		}

		qs = append(qs, s.client.Search.WithBody(bytes.NewReader(queryBs)))
	}

	if resp, err = s.client.Search(qs...); err != nil {
		return nil, err
	}

HandleResp:

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("resp status=%d, resp=%s", resp.StatusCode, resp.String())
	}

	if err = json.NewDecoder(resp.Body).Decode(result); err != nil {
		return nil, err
	}

	s.scroll = result.ScrollId

	return lo.Map(
		result.Hits.Hits,
		func(item *model.ESSource[map[string]any], _ int) map[string]any {
			return item.Content
		},
	), nil
}

// WriteData implements model.IO.
func (s *streamer) WriteData([]map[string]any) (int, error) {
	panic("unimplemented")
}

func NewStreamer(ctx context.Context, client *elastic.Client, index string) (model.IO[map[string]any], error) {
	s := &streamer{ctx: ctx, client: client, index: index}
	return s, nil
}