cost-model/pkg/customcost/types.go

package customcost

import (
	"fmt"
	"strings"
	"time"

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

type CostTotalRequest struct {
	Start       time.Time
	End         time.Time
	AggregateBy []string
	Step        time.Duration
	Filter      filter.Filter
}

type CostTimeseriesRequest struct {
	Start       time.Time
	End         time.Time
	AggregateBy []string
	Step        time.Duration
	Filter      filter.Filter
}

type CostResponse struct {
	Window          opencost.Window `json:"window"`
	TotalBilledCost float32         `json:"totalBilledCost"`
	TotalListCost   float32         `json:"totalListCost"`
	CustomCosts     []*CustomCost   `json:"customCosts"`
}

type CustomCost struct {
	Id             string  `json:"id"`
	Zone           string  `json:"zone"`
	AccountName    string  `json:"account_name"`
	ChargeCategory string  `json:"charge_category"`
	Description    string  `json:"description"`
	ResourceName   string  `json:"resource_name"`
	ResourceType   string  `json:"resource_type"`
	ProviderId     string  `json:"provider_id"`
	BilledCost     float32 `json:"billedCost"`
	ListCost       float32 `json:"listCost"`
	ListUnitPrice  float32 `json:"list_unit_price"`
	UsageQuantity  float32 `json:"usage_quantity"`
	UsageUnit      string  `json:"usage_unit"`
	Domain         string  `json:"domain"`
	Aggregate      string  `json:"aggregate"`
}

type CostTimeseriesResponse struct {
	Window     opencost.Window `json:"window"`
	Timeseries []*CostResponse `json:"timeseries"`
}

// --------------

func NewCostResponse(ccs *CustomCostSet) *CostResponse {
	costResponse := &CostResponse{
		Window:      ccs.Window,
		CustomCosts: []*CustomCost{},
	}

	for _, cc := range ccs.CustomCosts {
		costResponse.TotalBilledCost += cc.BilledCost
		costResponse.TotalListCost += cc.ListCost
		costResponse.CustomCosts = append(costResponse.CustomCosts, cc)
	}

	return costResponse
}

// --------------

func ParseCustomCostResponse(ccResponse *pb.CustomCostResponse) []*CustomCost {
	costs := ccResponse.GetCosts()

	customCosts := make([]*CustomCost, len(costs))
	for i, cost := range costs {
		customCosts[i] = &CustomCost{
			Id:             cost.GetId(),
			Zone:           cost.GetZone(),
			AccountName:    cost.GetAccountName(),
			ChargeCategory: cost.GetChargeCategory(),
			Description:    cost.GetDescription(),
			ResourceName:   cost.GetResourceName(),
			ResourceType:   cost.GetResourceType(),
			ProviderId:     cost.GetProviderId(),
			BilledCost:     cost.GetBilledCost(),
			ListCost:       cost.GetListCost(),
			ListUnitPrice:  cost.GetListUnitPrice(),
			UsageQuantity:  cost.GetUsageQuantity(),
			UsageUnit:      cost.GetUsageUnit(),
			Domain:         ccResponse.GetDomain(),
		}
	}

	return customCosts
}

func (cc *CustomCost) Add(other *CustomCost) {
	cc.BilledCost += other.BilledCost
	cc.ListCost += other.ListCost
	cc.ListUnitPrice += other.ListUnitPrice
	cc.UsageQuantity += other.UsageQuantity
}

type CustomCostSet struct {
	CustomCosts []*CustomCost
	Window      opencost.Window
}

func NewCustomCostSet(window opencost.Window) *CustomCostSet {
	return &CustomCostSet{
		CustomCosts: []*CustomCost{},
		Window:      window,
	}
}

func (ccs *CustomCostSet) Add(customCosts []*CustomCost) {
	ccs.CustomCosts = append(ccs.CustomCosts, customCosts...)
}

func (ccs *CustomCostSet) Aggregate(aggregateBy []string) error {
	if len(aggregateBy) == 0 {
		return fmt.Errorf("found empty aggregateBy")
	}

	aggMap := make(map[string]*CustomCost)
	for _, cc := range ccs.CustomCosts {
		aggKey, err := generateAggKey(cc, aggregateBy)
		if err != nil {
			return fmt.Errorf("failed to aggregate CustomCostSet: %w", err)
		}
		cc.Aggregate = aggKey

		if existing, ok := aggMap[aggKey]; ok {
			existing.Add(cc)
		} else {
			aggMap[aggKey] = cc
		}
	}

	var newCustomCosts []*CustomCost
	for _, customCost := range aggMap {
		newCustomCosts = append(newCustomCosts, customCost)
	}
	ccs.CustomCosts = newCustomCosts

	return nil
}

func generateAggKey(cc *CustomCost, aggregateBy []string) (string, error) {
	var aggKeys []string
	for _, agg := range aggregateBy {
		// TODO only domain is supported currently
		if agg == string(CustomCostDomainProp) {
			aggKeys = append(aggKeys, cc.Domain)
		} else {
			return "", fmt.Errorf("unsupported aggregation type: %s", agg)
		}
	}
	aggKey := strings.Join(aggKeys, "/")

	return aggKey, nil
}