102 lines
2.3 KiB
Go
102 lines
2.3 KiB
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
|
|
}
|