more search enhancements (#22)

pkg/es/search.go

@@ -19,10 +19,22 @@ package es
 import (
 	"context"
 	"fmt"
+	"log"
 	"log/slog"
+	"strings"
+	"time"
 
 	"codeberg.org/scip/esctl/pkg/cfg"
-	"github.com/tidwall/gjson"
+	"codeberg.org/scip/esctl/pkg/printer"
+	"github.com/alecthomas/repr"
+	"github.com/elastic/go-elasticsearch/v9/typedapi/core/search"
+	"github.com/elastic/go-elasticsearch/v9/typedapi/esdsl"
+	"github.com/elastic/go-elasticsearch/v9/typedapi/types"
+	"github.com/elastic/go-elasticsearch/v9/typedapi/types/enums/sortorder"
+)
+
+const (
+	MAXPAGE = 5000
 )
 
 /*
@@ -35,36 +47,146 @@ func Search(conf *cfg.Config, queries []string) error {
 	search := conf.DefaultCluster.ES.Search().
 		Index(conf.Index)
 
-	if len(queries) > 0 {
-		req, err := prepareQuery(conf, queries)
-		if err != nil {
-			return err
-		}
-
-		search.Request(req)
+	req, err := prepareQuery(conf, queries)
+	if err != nil {
+		return err
 	}
 
-	res, err := search.Do(context.Background())
+	search.Request(req)
+
+	switch conf.Tail {
+	case true:
+		return searchTail(conf, search)
+	case false:
+		if conf.To > MAXPAGE {
+			return searchPit(conf, req)
+		} else {
+			return searchOnce(conf, search)
+		}
+	}
+
+	return nil
+}
+
+func Debug(conf *cfg.Config) error {
+	res, err := conf.DefaultCluster.ES.Search().
+		Index(conf.Index).
+		Size(0).
+		Aggregations(map[string]types.Aggregations{
+			"min_ts": *esdsl.NewMinAggregation().Field("@timestamp").AggregationsCaster(),
+			"max_ts": *esdsl.NewMaxAggregation().Field("@timestamp").AggregationsCaster(),
+		}).
+		Do(context.Background())
+
 	if err != nil {
+		return err
+	}
+
+	repr.Println(res)
+
+	return nil
+}
+
+func searchOnce(conf *cfg.Config, search *search.Search) error {
+	res, err := search.
+		From(conf.From).
+		Size(conf.To).
+		Do(context.Background())
+	if err != nil {
+		if strings.Contains(err.Error(), "reason: all shards failed") {
+			return nil
+		}
+
 		return fmt.Errorf("failed to run search (esdsl): %s", err)
 	}
 
 	slog.Debug("ES result", "search", res)
 
 	for _, hit := range res.Hits.Hits {
-		docjson := fmt.Sprintf(`{"id":%s, "score":%0.4f, "index":"%s", "source":%s}`,
-			*hit.Id_,
-			*hit.Score_,
-			hit.Index_,
-			hit.Source_)
+		printer.PrintDoc(conf, hit)
+	}
 
-		if conf.Path != "" {
-			value := gjson.Get(docjson, conf.Path)
-			fmt.Println(value.String())
-		} else {
-			fmt.Println(docjson)
+	return nil
+}
+
+// https://www.elastic.co/docs/reference/elasticsearch/clients/go/using-the-api/searching#_pit_search_after
+func searchPit(conf *cfg.Config, req *search.Request) error {
+	ctx := context.Background()
+	pit, err := conf.DefaultCluster.ES.OpenPointInTime(conf.Index).KeepAlive("1m").Do(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to open point-in-time request for search: %s", err)
+	}
+	defer func() {
+		_, err := conf.DefaultCluster.ES.ClosePointInTime().Id(pit.Id).Do(ctx)
+		if err != nil {
+			log.Fatalf("failed to close PIT: %s", err)
+		}
+	}()
+
+	search := conf.DefaultCluster.ES.Search().
+		Request(req).
+		Pit(esdsl.NewPointInTimeReference().
+			Id(pit.Id).
+			KeepAlive(esdsl.NewDuration().String("1m"))).
+		Sort(esdsl.NewSortOptions().
+			AddSortOption("_shard_doc", esdsl.NewFieldSort(sortorder.Asc))).
+		Size(conf.To)
+
+	for {
+		res, err := search.Do(ctx)
+		if err != nil {
+			return fmt.Errorf("failed to run search (esdsl pit): %s", err)
+		}
+
+		if len(res.Hits.Hits) == 0 {
+			break
+		}
+
+		for _, hit := range res.Hits.Hits {
+			printer.PrintDoc(conf, hit)
+		}
+
+		last := res.Hits.Hits[len(res.Hits.Hits)-1]
+		search = search.SearchAfterValues(last.Sort)
+
+		if res.PitId != nil {
+			search = search.Pit(esdsl.NewPointInTimeReference().
+				Id(*res.PitId).
+				KeepAlive(esdsl.NewDuration().String("1m")))
 		}
 	}
 
 	return nil
 }
+
+func searchTail(conf *cfg.Config, search *search.Search) error {
+	docs := map[string]int{}
+
+	fmt.Println("enter ctrl-c to abort...")
+
+	for {
+		res, err := search.Do(context.Background())
+		if err != nil {
+			if strings.Contains(err.Error(), "reason: all shards failed") {
+				return nil
+			}
+
+			return fmt.Errorf("failed to run search (esdsl): %s", err)
+		}
+
+		slog.Debug("ES result", "search", res)
+
+		for _, hit := range res.Hits.Hits {
+			_, exists := docs[*hit.Id_]
+			if exists {
+				continue
+			}
+
+			printer.PrintDoc(conf, hit)
+
+			docs[*hit.Id_] = 1
+		}
+
+		time.Sleep(100 * time.Millisecond)
+	}
+}