esctl/pkg/es/index.go

/*
Copyright © 2026 Thomas von Dein

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
package es

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"regexp"
	"slices"
	"strconv"
	"strings"
	"time"

	"codeberg.org/scip/esctl/pkg/cfg"
	"codeberg.org/scip/esctl/pkg/printer"
	"codeberg.org/scip/mapmap"
	"github.com/charmbracelet/lipgloss"
	"github.com/elastic/go-elasticsearch/v9/typedapi/cat/indices"
	"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/healthstatus"
)

// used for completion
func IndexNames(conf *cfg.Config) ([]string, error) {
	res, err := conf.DefaultCluster.ES().Cat.Indices().
		Do(context.Background())
	if err != nil {
		return nil, fmt.Errorf("failed to get indices: %w", esErrorString(err))
	}

	indices := make([]string, len(res))
	for idx, index := range res {
		indices[idx] = *index.Index
	}

	return indices, nil
}

func filterIndices(conf *cfg.Config, list indices.Response) indices.Response {
	var filter *regexp.Regexp

	if len(conf.Filter) > 0 {
		filter = regexp.MustCompile(conf.Filter[0])
	}

	return mapmap.NewSlicer(list).MapSliceValuesImmutable(func(index types.IndicesRecord) bool {
		if !conf.Partials && strings.HasPrefix(*index.Index, "partial-") {
			return false
		}

		if !conf.Hidden && strings.HasPrefix(*index.Index, ".") {
			return false
		}

		if len(conf.Filter) > 0 {
			if !filter.MatchString(*index.Index) {
				return false
			}
		}

		return true
	})
}

func IndexList(conf *cfg.Config) error {
	cat := conf.DefaultCluster.ES().Cat.Indices()

	if conf.Failed {
		cat = cat.Health(healthstatus.Red)
	}

	res, err := cat.Do(context.Background())
	if err != nil {
		return fmt.Errorf("failed to get indices: %w", esErrorString(err))
	}

	slog.Debug("ES result", "indices", res)

	list := filterIndices(conf, res)

	size := len(list)

	if conf.MaxItems > 0 {
		if size > conf.MaxItems {
			size = conf.MaxItems
		}
	}

	table := printer.NewTable(conf).WithSize(3, size)
	table.Addheaders("name", "size", "docscount")

	for idx, index := range list {
		name := printer.Colorize(conf, *index.Health, *index.Index)

		table.Entries[idx] = []any{name, *index.DatasetSize, *index.DocsCount}

		if idx == size-1 {
			break
		}
	}

	table.Sort()

	return table.Print()
}

func IndexShow(conf *cfg.Config, indexpattern string) error {
	res, err := conf.DefaultCluster.ES().Indices.Get(indexpattern).
		Do(context.Background())
	if err != nil {
		return fmt.Errorf("failed to get index: %w", esErrorString(err))
	}

	slog.Debug("index show", "index", res)

	for name, index := range res {
		fields := make([]string, len(index.Mappings.Properties))

		idx := 0
		for field := range index.Mappings.Properties {
			fields[idx] = field
			idx++
		}

		table := printer.NewTable(conf).WithSize(2, 7)
		table.Addheaders("index property", "value")

		ts, err := strconv.ParseInt(index.Settings.Index.CreationDate.(string), 10, 64)
		if err != nil {
			ts = 0
		}

		created := time.Unix(ts/1000, 0)

		table.Entries = [][]any{
			{"name", name},
			{"replicas", *index.Settings.Index.NumberOfReplicas},
			{"shards", *index.Settings.Index.NumberOfShards},
			{"created", created},
			{"uuid", *index.Settings.Index.Uuid},
			{"version", *index.Settings.Index.Version.Created},
			{"fields", fields},
		}

		if index.Settings.Index.Lifecycle != nil {
			table.Entries = append(table.Entries, [][]any{
				{"ilm policy", *index.Settings.Index.Lifecycle.Name},
				{"ilm rollover alias", *index.Settings.Index.Lifecycle.RolloverAlias},
			}...)
		}

		if err := table.Print(); err != nil {
			return err
		}

		fmt.Println()
	}

	return nil
}

func IndexCreate(conf *cfg.Config, index string, mappings []string) error {
	settings := esdsl.NewIndexSettings()

	create := conf.DefaultCluster.ES().Indices.Create(index)

	if conf.Wait {
		create.WaitForActiveShards("all")
	}

	if conf.Shards > 0 {
		settings.NumberOfShards(strconv.Itoa(conf.Shards))
	}

	if conf.Replicas > 0 {
		settings.NumberOfReplicas(strconv.Itoa(conf.Replicas))
	}

	if len(mappings) > 0 {
		maps := esdsl.NewTypeMapping()

		for _, mapping := range mappings {
			parts := strings.Split(mapping, ":")
			if len(parts) != 2 {
				return fmt.Errorf("invalid mapping %s, expect <name:type> (type: integer, text, date, keyword)", mapping)
			}

			switch parts[1] {
			case "text":
				maps.AddProperty(parts[0], esdsl.NewTextProperty())
			case "integer":
				maps.AddProperty(parts[0], esdsl.NewIntegerNumberProperty())
			case "date":
				maps.AddProperty(parts[0], esdsl.NewDateProperty())
			case "keyword":
				maps.AddProperty(parts[0], esdsl.NewKeywordProperty())
			}
		}

		create.Mappings(maps)
	}

	if conf.Policy != "" {
		settings.Lifecycle(
			esdsl.NewIndexSettingsLifecycle().
				Name(conf.Policy).
				RolloverAlias(index))
	}

	_, err := create.Settings(settings).
		Do(context.Background())
	if err != nil {
		return fmt.Errorf("failed to create index: %w", esErrorString(err))
	}

	return nil
}

func IndexDelete(conf *cfg.Config, index string) error {
	_, err := conf.DefaultCluster.ES().Indices.Delete(index).
		Do(context.Background())
	if err != nil {
		return fmt.Errorf("failed to delete index: %w", esErrorString(err))
	}

	return nil
}

func IndexClose(conf *cfg.Config, index string) error {
	_, err := conf.DefaultCluster.ES().Indices.Close(index).
		Do(context.Background())
	if err != nil {
		return fmt.Errorf("failed to close index: %w", esErrorString(err))
	}

	return nil
}

func IndexFields(conf *cfg.Config, index string) error {
	res, err := conf.DefaultCluster.ES().FieldCaps().
		Index(index).
		Fields("*").
		Do(context.Background())
	if err != nil {
		return fmt.Errorf("failed to retrieve field capabilties: %w", esErrorString(err))
	}

	table := printer.NewTable(conf).WithSize(5, 0)
	table.Addheaders("field", "type", "searchable", "aggretable", "metadata")

	idx := 0

	for name, field := range res.Fields {
		for fieldtype, caps := range field {
			// fields only have 1 type, so this one is it
			switch {
			case conf.Searchable && !caps.Searchable:
				continue
			case conf.Aggretable && !caps.Aggregatable:
				continue
			case len(conf.Filter) > 0 && !slices.Contains(conf.Filter, fieldtype):
				continue
			}

			table.AddRow(name, fieldtype,
				caps.Searchable,
				caps.Aggregatable,
				*caps.MetadataField)

			break
		}

		idx++
	}

	table.Sort()

	if err := table.Print(); err != nil {
		return err
	}

	return nil
}

type Diskusage struct {
	Total       int64 `json:"total_in_bytes"`
	Points      int64 `json:"points_in_bytes"`
	Norms       int64 `json:"norms_in_bytes"`
	TermVectors int64 `json:"term_vectors_in_bytes"`
	KnnVectors  int64 `json:"knn_vectors_in_bytes"`
	BloomFilter int64 `json:"bloom_filter_in_bytes"`
}

type IndexDiskUsage struct {
	AllFields Diskusage            `json:"all_fields"`
	Fields    map[string]Diskusage `json:"fields"`
}

type ResIndexDiskUsage map[string]IndexDiskUsage

func IndexDiskusage(conf *cfg.Config, index string) error {
	var bold = lipgloss.NewStyle().Bold(true)

	res, err := conf.DefaultCluster.ES().Indices.DiskUsage(index).
		RunExpensiveTasks(true).
		Do(context.Background())
	if err != nil {
		return fmt.Errorf("failed to retrieve index disk usage: %w", esErrorString(err))
	}

	duRes := ResIndexDiskUsage{}
	if err := json.Unmarshal(res, &duRes); err != nil {
		return fmt.Errorf("failed to unmarshal disk usage response: %w", err)
	}

	diskusage, exists := duRes[index]
	if !exists {
		return errors.New("no disk usage reported for index")
	}

	table := printer.NewTable(conf).
		WithHeaders("field", "bloom filter", "norms", "points", "term vectors", "knn vectors", "total")

	for name, field := range diskusage.Fields {
		if strings.HasPrefix(name, "_") || strings.HasSuffix(name, ".keyword") {
			continue
		}

		table.AddRow(
			name,
			printer.Bytes(field.BloomFilter),
			printer.Bytes(field.Norms),
			printer.Bytes(field.Points),
			printer.Bytes(field.TermVectors),
			printer.Bytes(field.KnnVectors),
			printer.Bytes(field.Total),
		)
	}

	all := diskusage.AllFields

	table.Sort()

	table.AddRowLate(
		bold.Render("Summary"),
		printer.Bytes(all.BloomFilter),
		printer.Bytes(all.Norms),
		printer.Bytes(all.Points),
		printer.Bytes(all.TermVectors),
		printer.Bytes(all.KnnVectors),
		printer.Bytes(all.Total),
	)

	if err := table.Print(); err != nil {
		return err
	}

	return nil
}