cost-model/pkg/asset/autocompletequeryservice.go

package asset

import (
	"context"
	"fmt"
	"sort"
	"strings"

	"github.com/opencost/opencost/core/pkg/filter"
	"github.com/opencost/opencost/core/pkg/opencost"
)

const DefaultAutocompleteResultLimit = 100
const MaxAutocompleteResultLimit = 1000

type AssetAutocompleteRequest struct {
	TenantID string
	Search   string
	Field    string
	Limit    int
	Window   opencost.Window
	Filter   filter.Filter
}

type AssetAutocompleteResponse struct {
	Data []string `json:"data"`
}

type AutocompleteQueryService interface {
	QueryAssetAutocomplete(AssetAutocompleteRequest, context.Context) (*AssetAutocompleteResponse, error)
}

func QueryAssetAutocompleteFromSet(assetSet *opencost.AssetSet, req AssetAutocompleteRequest) (*AssetAutocompleteResponse, error) {
	if req.TenantID == "" {
		return nil, fmt.Errorf("tenant ID is required")
	}

	field, err := validateAutocompleteField(req.Field)
	if err != nil {
		return nil, fmt.Errorf("invalid field: %w", err)
	}

	limit := req.Limit
	if limit <= 0 {
		limit = DefaultAutocompleteResultLimit
	}
	if limit > MaxAutocompleteResultLimit {
		return nil, fmt.Errorf("exceeded maxiumum autocomplete result limit of %d", MaxAutocompleteResultLimit)
	}

	var matcher opencost.AssetMatcher
	if req.Filter != nil {
		compiler := opencost.NewAssetMatchCompiler()
		matcher, err = compiler.Compile(req.Filter)
		if err != nil {
			return nil, fmt.Errorf("failed to compile filter: %w", err)
		}
	}

	search := strings.ToLower(req.Search)
	results := map[string]struct{}{}
	for _, a := range assetSet.Assets {
		if a == nil {
			continue
		}
		if matcher != nil && !matcher.Matches(a) {
			continue
		}

		values := assetAutocompleteValues(a, field)
		for _, value := range values {
			if value == "" {
				continue
			}
			if search != "" && !strings.Contains(strings.ToLower(value), search) {
				continue
			}
			results[value] = struct{}{}
		}
	}

	data := make([]string, 0, len(results))
	for value := range results {
		data = append(data, value)
	}
	sort.Strings(data)
	if len(data) > limit {
		data = data[:limit]
	}
	return &AssetAutocompleteResponse{Data: data}, nil
}

func validateAutocompleteField(field string) (string, error) {
	f := strings.ToLower(field)
	switch f {
	case "account", "cluster", "name", "provider", "providerid", "type", "category":
		return f, nil
	}
	if strings.HasPrefix(f, "label") {
		return f, nil
	}
	return "", fmt.Errorf("unrecognized field: %s", field)
}

func assetAutocompleteValues(asset opencost.Asset, field string) []string {
	props := asset.GetProperties()
	if props == nil {
		return nil
	}
	switch {
	case field == "account":
		return []string{props.Account}
	case field == "cluster":
		return []string{props.Cluster}
	case field == "name":
		return []string{props.Name}
	case field == "provider":
		return []string{props.Provider}
	case field == "providerid":
		return []string{props.ProviderID}
	case field == "type":
		return []string{asset.Type().String()}
	case field == "category":
		return []string{props.Category}
	case field == "label":
		keys := make([]string, 0, len(asset.GetLabels()))
		for key := range asset.GetLabels() {
			keys = append(keys, key)
		}
		return keys
	case strings.HasPrefix(field, "label:"):
		labelName := strings.TrimPrefix(field, "label:")
		if value, ok := asset.GetLabels()[labelName]; ok {
			return []string{value}
		}
	}
	return nil
}