aws_cost_service.go

// Note: This code is based on parts of github.com/toricls/acos

package aws_api

import (
 "context"
 "fmt"
 "log"
 "strconv"
 "time"

 "github.com/(...)/pkg/models"
 "github.com/aws/aws-sdk-go-v2/aws"
 "github.com/aws/aws-sdk-go-v2/config"
 "github.com/aws/aws-sdk-go-v2/service/costexplorer"
 "github.com/aws/aws-sdk-go-v2/service/costexplorer/types"
 costTypes "github.com/aws/aws-sdk-go-v2/service/costexplorer/types"
 "github.com/aws/aws-sdk-go-v2/service/organizations"
 organizationTypes "github.com/aws/aws-sdk-go-v2/service/organizations/types"
)

const (
 ceDataGranularity    = "DAILY"
 ceCostMetric         = "UnblendedCost"
 ceForcastCostMetric  = "UNBLENDED_COST"
 ceCostAccountGroupBy = "LINKED_ACCOUNT"
 ceCostServiceGroupBy = "SERVICE"
 ceCostAvZoneGroupBy  = "AZ"
)

type AwsCostService struct {
 cfg                 aws.Config
 organizationsClient *organizations.Client
 ceClient            *costexplorer.Client
}

type Account organizationTypes.Account
type Accounts map[string]Account

// AccountIds returns a list of account IDs.
func (a *Accounts) accountIds() []string {
 accountIds := make([]string, len(*a))
 i := 0
 for key := range *a {
  accountIds[i] = key
  i++
 }
 return accountIds
}

type Costs map[string]*models.AwsCost // map[accountId]Cost
type Group costTypes.Group

func (g *Group) getAccountId() string {
 return g.Keys[0]
}

func (g *Group) getServiceName() string {
 return g.Keys[1]
}

func (g *Group) getRegionName() string {
 return g.Keys[1]
}

func (g *Group) getAmount() float64 {
 if f, err := strconv.ParseFloat(*g.Metrics[ceCostMetric].Amount, 32); err == nil {
  return f
 }
 // TODO: debug log the error
 return 0
}

type GetCostsOption struct {
 ExcludeCredit  bool
 ExcludeUpfront bool
 ExcludeRefund  bool
 ExcludeSupport bool

 // requires the following dates to show - THIS_MONTH, vs YESTERDAY, vs LAST_WEEK, and LAST_MONTH
 dates struct {
  asOf                string
  yesterday           string
  oneWeekAgo          string
  oneMonthAgo         string
  firstDayOfLastMonth string
  firstDayOfThisMonth string
  lastDayOfThisMonth  string
 }
}

func NewGetCostsOption(asOfInUTC time.Time) GetCostsOption {
 opt := GetCostsOption{
  ExcludeCredit:  true,
  ExcludeUpfront: true,
  ExcludeRefund:  false,
  ExcludeSupport: false,
 }

 yesterday := asOfInUTC.Add(time.Duration(-1) * 24 * time.Hour)
 oneWeekAgo := asOfInUTC.Add(time.Duration(-7) * 24 * time.Hour)
 oneMonthAgo := asOfInUTC.Add(time.Duration(-30) * 24 * time.Hour)
 year, month, _ := asOfInUTC.Date()
 firstDayOfThisMonth := time.Date(year, month, 1, 0, 0, 0, 0, asOfInUTC.Location())
 lastDayOfThisMonth := time.Date(year, month+1, 1, 0, 0, 0, 0, asOfInUTC.Location()).Add(-1 * 24 * time.Hour)
 firstDayOfLastMonth := time.Date(year, month-1, 1, 0, 0, 0, 0, asOfInUTC.Location())

 dateFmt := "2006-01-02"
 opt.dates.asOf = asOfInUTC.Format(dateFmt)
 opt.dates.yesterday = yesterday.Format(dateFmt)
 opt.dates.oneWeekAgo = oneWeekAgo.Format(dateFmt)
 opt.dates.oneMonthAgo = oneMonthAgo.Format(dateFmt)
 opt.dates.firstDayOfThisMonth = firstDayOfThisMonth.Format(dateFmt)
 opt.dates.lastDayOfThisMonth = lastDayOfThisMonth.Format(dateFmt)
 opt.dates.firstDayOfLastMonth = firstDayOfLastMonth.Format(dateFmt)
 return opt
}

func NewAwsCostService() (*AwsCostService, error) {
 var err error
 cfg, err := config.LoadDefaultConfig(context.TODO())
 if err != nil {
  log.Fatalf("unable to load AWS SDK config, %v", err)
 }

 return &AwsCostService{
  cfg:                 cfg,
  organizationsClient: organizations.NewFromConfig(cfg),
  ceClient:            costexplorer.NewFromConfig(cfg),
 }, nil
}

// ListAccounts returns a list of AWS accounts within an AWS Organization organization.
func (a AwsCostService) ListAccounts(ctx context.Context) (Accounts, error) {
 var nextToken *string
 accnts := make(map[string]Account)
 for {
  out, err := a.organizationsClient.ListAccounts(
   ctx,
   &organizations.ListAccountsInput{
    NextToken: nextToken,
   },
  )
  if err != nil {
   return nil, err
  }
  for _, acc := range out.Accounts {
   accnts[*acc.Id] = Account(acc)
  }
  nextToken = out.NextToken
  if nextToken == nil {
   break
  }
 }
 return accnts, nil
}

// GetCosts returns the costs for given accounts.
// It raises an error when the `accounts` arg doesn't contain any account.
func (a AwsCostService) GetCosts(ctx context.Context, accounts Accounts, opt GetCostsOption) (Costs, error) {
 accountIds := accounts.accountIds()
 if len(accountIds) == 0 {
  return nil, fmt.Errorf("error no account to retrieve cost: GetCosts requires at least one account in Accounts")
 }
 ceOpt := getCostsOptionOptToCostExplorerOpt(opt, accountIds)
 costs := make(Costs)

 for _, acc := range accounts {
  a.addCostEntry(*acc.Id, *acc.Name, costs)
 }

 var nextToken *string
 costPerServiceMap := map[string]*models.AwsDailyCost{}
 for {
  ceOpt.NextPageToken = nextToken

  out, err := a.ceClient.GetCostAndUsage(ctx, &ceOpt)
  if err != nil {
   return nil, err
  }

  a.processResultsByTime(costs, costPerServiceMap, opt, out)

  nextToken = out.NextPageToken
  if nextToken == nil {
   break
  }
 }
 if err := a.updateByRegionCosts(ctx, ceOpt, opt, costs); err != nil {
  return nil, fmt.Errorf("error updating region costs: %s", err.Error())
 }

 for _, account := range accountIds {
  if err := a.updateByForecastedCosts(ctx, getCostsOptionOptToForecastCostExplorerOpt(opt, []string{account}), costs[account]); err != nil {
   return nil, fmt.Errorf("error updating forecasted costs: %s", err.Error())
  }
 }

 return costs, nil
}

func (a AwsCostService) processResultsByTime(costs Costs, costPerServiceMap map[string]*models.AwsDailyCost, opt GetCostsOption, out *costexplorer.GetCostAndUsageOutput) Costs {
 for _, r := range out.ResultsByTime {
  thisMonth := *r.TimePeriod.Start >= opt.dates.firstDayOfThisMonth
  lastWeek := *r.TimePeriod.Start >= opt.dates.oneWeekAgo
  lastMonth := *r.TimePeriod.Start >= opt.dates.oneMonthAgo

  for _, g := range r.Groups {
   grp := Group(g)
   accntId := grp.getAccountId()
   amount := grp.getAmount()
   service := grp.getServiceName()

   c := costs[accntId]
   dailyCosts := &c.PreviousDailyCosts

   if thisMonth {
    dailyCosts = &c.CurrentDailyCosts
    if *r.TimePeriod.Start > opt.dates.asOf {
     dailyCosts = &c.ForecastedCosts
    }
    a.updateServiceCosts(c, costPerServiceMap, fmt.Sprintf("%s%s", accntId, service), service, amount)
   }

   if *r.TimePeriod.Start == opt.dates.asOf {
    c.Totals.TodaysCost += amount
   }

   if *r.TimePeriod.Start == opt.dates.yesterday {
    c.Totals.YesterdayCost += amount
   }

   if lastWeek {
    c.Totals.LastSevenDaysCost += amount
   }

   if lastMonth {
    c.Totals.LastThirtyDaysCost += amount
   }

   if exists, val := a.checkIfLastDailyCostExists(*r.TimePeriod.Start, *dailyCosts); exists {
    val.Cost += amount
   } else {
    *dailyCosts = append(*dailyCosts, &models.AwsDailyCost{
     Key:  *r.TimePeriod.Start,
     Cost: amount,
    })
   }
  }
 }
 return costs
}

func (a AwsCostService) updateServiceCosts(cost *models.AwsCost, costPerServiceMap map[string]*models.AwsDailyCost, accountServiceKey string, service string, amount float64) {
 if val, ok := costPerServiceMap[accountServiceKey]; ok {
  val.Cost += amount
 } else {
  dc := &models.AwsDailyCost{
   Key:  service,
   Cost: amount,
  }
  cost.CostsPerService = append(cost.CostsPerService, dc)
  costPerServiceMap[accountServiceKey] = dc
 }
}

func (a AwsCostService) checkIfLastDailyCostExists(day string, dailyCostList []*models.AwsDailyCost) (bool, *models.AwsDailyCost) {
 if len(dailyCostList) == 0 {
  return false, nil
 }
 ret := dailyCostList[len(dailyCostList)-1]
 if day != ret.Key {
  return false, nil
 }
 return true, dailyCostList[len(dailyCostList)-1]
}

func (a AwsCostService) addCostEntry(Id string, Name string, costs Costs) {
 costs[Id] = &models.AwsCost{
  AccountID:   Id,
  AccountName: Name,
  Totals: models.AwsCostTotals{
   TodaysCost:         0.0,
   YesterdayCost:      0.0,
   LastSevenDaysCost:  0.0,
   LastThirtyDaysCost: 0.0,
  },
  PreviousDailyCosts: []*models.AwsDailyCost{},
  CurrentDailyCosts:  []*models.AwsDailyCost{},
  ForecastedCosts:    []*models.AwsDailyCost{},
  CostsPerService:    []*models.AwsDailyCost{},
  CostsPerLocation:   []*models.AwsDailyCost{},
 }
}

func (a AwsCostService) updateByForecastedCosts(ctx context.Context, ceOpt costexplorer.GetCostForecastInput, cost *models.AwsCost) error {
 out, err := a.ceClient.GetCostForecast(ctx, &ceOpt)
 if err != nil {
  return err
 }
 for _, r := range out.ForecastResultsByTime {
  amount, err := strconv.ParseFloat(*r.MeanValue, 64)
  if err != nil {
   return err
  }
  cost.ForecastedCosts = append(cost.ForecastedCosts, &models.AwsDailyCost{
   Key:  *r.TimePeriod.Start,
   Cost: amount,
  })
 }
 return err
}

func (a AwsCostService) updateByRegionCosts(ctx context.Context, ceOpt costexplorer.GetCostAndUsageInput, opt GetCostsOption, costs Costs) error {
 ceOpt.GroupBy = ceOpt.GroupBy[:len(ceOpt.GroupBy)-1]
 ceOpt.GroupBy = append(ceOpt.GroupBy, types.GroupDefinition{
  Type: types.GroupDefinitionTypeDimension,
  Key:  aws.String(ceCostAvZoneGroupBy),
 })
 ceOpt.TimePeriod.Start = &opt.dates.firstDayOfThisMonth
 costPerLocationMap := map[string]*models.AwsDailyCost{}

 var nextToken *string
 for {
  ceOpt.NextPageToken = nextToken

  out, err := a.ceClient.GetCostAndUsage(ctx, &ceOpt)
  if err != nil {
   return err
  }
  for _, r := range out.ResultsByTime {
   for _, g := range r.Groups {
    grp := Group(g)
    accntId := grp.getAccountId()
    amount := grp.getAmount()
    region := grp.getRegionName()
    accntServiceKey := fmt.Sprintf("%s%s", accntId, region)

    c := costs[accntId]

    if val, ok := costPerLocationMap[accntServiceKey]; ok {
     val.Cost += amount
    } else {
     dc := &models.AwsDailyCost{
      Key:  region,
      Cost: amount,
     }
     c.CostsPerLocation = append(c.CostsPerLocation, dc)
     costPerLocationMap[accntServiceKey] = dc
    }
   }
  }

  nextToken = out.NextPageToken
  if nextToken == nil {
   break
  }
 }
 return nil
}

// acosOptToCostExplorerOpt returns the AWS Cost Explorer's GetCostAndUsageInput param built from the acos options.
func getCostsOptionOptToCostExplorerOpt(opt GetCostsOption, accountIds []string) costexplorer.GetCostAndUsageInput {
 // Base input parameter
 in := costexplorer.GetCostAndUsageInput{
  Granularity: ceDataGranularity,
  Metrics:     []string{ceCostMetric},
  TimePeriod: &types.DateInterval{
   // Get the cost for the last month and this month
   Start: aws.String(opt.dates.firstDayOfLastMonth),
   End:   aws.String(opt.dates.lastDayOfThisMonth),
  },
  GroupBy: []types.GroupDefinition{
   {
    Type: types.GroupDefinitionTypeDimension,
    Key:  aws.String(ceCostAccountGroupBy),
   },
   {
    Type: types.GroupDefinitionTypeDimension,
    Key:  aws.String(ceCostServiceGroupBy),
   },
  },
  Filter: &types.Expression{
   And: []types.Expression{
    {
     Dimensions: &types.DimensionValues{
      Key:    ceCostAccountGroupBy,
      Values: accountIds,
     },
    },
   },
  },
 }

 // Exclude options
 v := []string{}
 if opt.ExcludeCredit {
  v = append(v, "Credit")
 }
 if opt.ExcludeUpfront {
  v = append(v, "Upfront")
 }
 if opt.ExcludeRefund {
  v = append(v, "Refund")
 }
 if opt.ExcludeSupport {
  v = append(v, "Support")
 }
 if len(v) > 0 {
  in.Filter.And = append(in.Filter.And, types.Expression{
   Not: &types.Expression{
    Dimensions: &types.DimensionValues{
     Key:    "RECORD_TYPE",
     Values: v,
    },
   },
  })
 }
 return in
}

func getCostsOptionOptToForecastCostExplorerOpt(opt GetCostsOption, accountIds []string) costexplorer.GetCostForecastInput {
 // Base input parameter
 in := costexplorer.GetCostForecastInput{
  Granularity: ceDataGranularity,
  Metric:      ceForcastCostMetric,
  TimePeriod: &types.DateInterval{
   // Get the cost for the last month and this month
   Start: aws.String(opt.dates.asOf),
   End:   aws.String(opt.dates.lastDayOfThisMonth),
  },
  Filter: &types.Expression{
   And: []types.Expression{
    {
     Dimensions: &types.DimensionValues{
      Key:    ceCostAccountGroupBy,
      Values: accountIds,
     },
    },
   },
  },
 }

 // Exclude options
 v := []string{}
 if opt.ExcludeCredit {
  v = append(v, "Credit")
 }
 if opt.ExcludeUpfront {
  v = append(v, "Upfront")
 }
 if opt.ExcludeRefund {
  v = append(v, "Refund")
 }
 if opt.ExcludeSupport {
  v = append(v, "Support")
 }
 if len(v) > 0 {
  in.Filter.And = append(in.Filter.And, types.Expression{
   Not: &types.Expression{
    Dimensions: &types.DimensionValues{
     Key:    "RECORD_TYPE",
     Values: v,
    },
   },
  })
 }
 return in
}