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 }