esctl/pkg/es/shard.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"
	"fmt"
	"log/slog"

	"codeberg.org/scip/esctl/pkg/cfg"
	"codeberg.org/scip/esctl/pkg/printer"
	"github.com/elastic/go-elasticsearch/v9/typedapi/cat/shards"
)

func colorzizeShard(conf *cfg.Config, state, name string) string {
	color := "red" // in case of RELOCATING and UNASSIGNED

	switch state {
	case "STARTED":
		color = "green"
	case "INITIALIZING":
		color = "yellow"
	}

	return printer.Colorize(conf, color, name)
}

func resolvePrirep(state string) string {
	switch state {
	case "p":
		return "primary"
	}

	return "replica"
}

func filterShards(conf *cfg.Config, shardlist shards.Response) shards.Response {
	filtered := shards.Response{}

	size := len(shardlist)

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

	for idx, shard := range shardlist {
		if conf.Failed {
			if *shard.State == "STARTED" {
				continue
			}
		}

		if conf.Primary {
			if *shard.Prirep != "p" {
				continue
			}
		}

		if idx == size-1 {
			break
		}

		filtered = append(filtered, shard)
	}

	return filtered
}

func ShardList(conf *cfg.Config) error {
	res, err := conf.DefaultCluster.ES().Cat.Shards().
		Do(context.Background())
	if err != nil {
		return fmt.Errorf("failed to get shards: %w", esErrorString(err))
	}

	shardlist := filterShards(conf, res)
	slog.Debug("ES result", "shards", res)

	return printShards(conf, shardlist)
}

func printShards(conf *cfg.Config, shardlist shards.Response) error {
	headers := []string{"index", "shard", "is primary", "store", "dataset", "docs"}
	if conf.Verbose {
		headers = append(headers, "node", "ip")
	}

	table := printer.NewTable(conf, len(headers), len(shardlist))
	table.Addheaders(headers...)

	for idx, shard := range shardlist {
		name := colorzizeShard(conf, *shard.State, *shard.Index)

		table.Entries[idx] = []any{
			name,
			*shard.Shard,
			resolvePrirep(*shard.Prirep),
			*shard.Store,
			*shard.Dataset,
			*shard.Docs,
		}

		if conf.Verbose {
			table.Entries[idx] = append(table.Entries[idx],
				*shard.Node,
				*shard.Ip,
			)
		}
	}

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

	return nil
}

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

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

	conf.Verbose = true

	return printShards(conf, res)
}

func ShardAllocation(conf *cfg.Config, index string) error {
	explain := conf.DefaultCluster.ES().Cluster.AllocationExplain().
		Index(index).
		Primary(conf.Primary).
		Shard(conf.Shards)

	if conf.FromNode != "" {
		explain.CurrentNode(conf.FromNode)
	}

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

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

	currentNode := res.CurrentNode

	table := printer.NewTable(conf, 2, 10)
	table.Addheaders("shard allocation setting", "value")

	roles := make([]string, len(currentNode.Roles))
	for idx, role := range currentNode.Roles {
		roles[idx] = role.Name
	}

	table.Entries = [][]any{
		{"Index", index},
		{"Current state", res.CurrentState},
		{"Current node", currentNode.Name},
		{"Current k8s node", currentNode.Attributes["k8s_node_name"]},
		{"Current node address", currentNode.TransportAddress},
		{"Current node id", currentNode.Id},
		{"Current node weight", currentNode.WeightRanking},
		{"Current node roles", roles},
		{"Can rebalance cluster", res.CanRebalanceCluster.Name},
		{"Can rebalance to another node", res.CanRebalanceToOtherNode.Name},
		{"Can remain on current node", res.CanRemainOnCurrentNode.Name},
	}

	if res.CurrentState == "unassigned" {
		table.AddRow("Unassignment reason", res.UnassignedInfo.Reason.String()+" at "+res.UnassignedInfo.At.(string))
	}

	for _, nodeDecision := range res.NodeAllocationDecisions {
		for _, decider := range nodeDecision.Deciders {
			table.AddRow("Allocation decider", decider.Decider)
			table.AddRow("    -> decision", decider.Decision.String())
			table.AddRow("    -> explanation", decider.Explanation)
		}
	}

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

	return nil
}