cost-model/modules/pricing/public/generator.go

package public

import (
	"fmt"
	"os"

	"github.com/opencost/opencost/core/pkg/log"
	"github.com/opencost/opencost/core/pkg/pricing"
	"github.com/opencost/opencost/core/pkg/unit"
	"github.com/opencost/opencost/modules/pricing/public/aws"
	"github.com/opencost/opencost/modules/pricing/public/azure"
	"github.com/opencost/opencost/modules/pricing/public/gcp"
)

// GenerateAWSPricing fetches AWS pricing data in the specified currency
func GenerateAWSPricing(currency unit.Currency) (*pricing.PricingSet, error) {
	log.Infof("Generating AWS pricing for currency: %s", currency)

	source := aws.NewAWSPricingSource(aws.AWSPricingSourceConfig{
		CurrencyCode: string(currency),
	})

	pricingSet, err := source.GetPricing()
	if err != nil {
		return nil, fmt.Errorf("failed to get AWS pricing: %w", err)
	}

	// Sort to ensure deterministic output for checksums
	pricingSet.Sort()

	log.Infof("Generated %d AWS node pricing entries", len(pricingSet.NodePricing))
	return pricingSet, nil
}

// GenerateAzurePricing fetches Azure pricing data in the specified currency
func GenerateAzurePricing(currency unit.Currency) (*pricing.PricingSet, error) {
	log.Infof("Generating Azure pricing for currency: %s", currency)

	source := azure.NewAzurePricingSource(azure.AzurePricingSourceConfig{
		CurrencyCode: string(currency),
	})

	pricingSet, err := source.GetPricing()
	if err != nil {
		return nil, fmt.Errorf("failed to get Azure pricing: %w", err)
	}

	// Sort to ensure deterministic output for checksums
	pricingSet.Sort()

	log.Infof("Generated %d Azure node pricing entries", len(pricingSet.NodePricing))
	return pricingSet, nil
}

// GenerateGCPPricing fetches GCP pricing data in the specified currency
func GenerateGCPPricing(currency unit.Currency) (*pricing.PricingSet, error) {
	log.Infof("Generating GCP pricing for currency: %s", currency)

	source := gcp.NewGCPPricingSource(gcp.GCPPricingSourceConfig{
		CurrencyCode: string(currency),
		APIKey:       os.Getenv("GCP_API_KEY"),
	})

	pricingSet, err := source.GetPricing()
	if err != nil {
		return nil, fmt.Errorf("failed to get GCP pricing: %w", err)
	}

	// Sort to ensure deterministic output for checksums
	pricingSet.Sort()

	log.Infof("Generated %d GCP node pricing entries", len(pricingSet.NodePricing))
	return pricingSet, nil
}

// GeneratePricing fetches pricing data for all supported providers
// and combines them into a single PricingSet
func GeneratePricing(currency unit.Currency) (*pricing.PricingSet, error) {
	log.Infof("Generating pricing for providers in currency: %s", currency)

	// Create a combined pricing set
	combinedSet := &pricing.PricingSet{
		NodePricing:             []*pricing.NodePricing{},
		PersistentVolumePricing: []*pricing.PersistentVolumePricing{},
		ServicePricing:          []*pricing.ServicePricing{},
	}

	// Fetch AWS pricing
	awsSet, err := GenerateAWSPricing(currency)
	if err != nil {
		return nil, fmt.Errorf("failed to get AWS pricing: %w", err)
	}
	combinedSet.NodePricing = append(combinedSet.NodePricing, awsSet.NodePricing...)
	combinedSet.PersistentVolumePricing = append(combinedSet.PersistentVolumePricing, awsSet.PersistentVolumePricing...)
	combinedSet.ServicePricing = append(combinedSet.ServicePricing, awsSet.ServicePricing...)
	log.Infof("Added %d AWS node pricing entries, %d service pricing entries", len(awsSet.NodePricing), len(awsSet.ServicePricing))

	// Fetch Azure pricing
	azureSet, err := GenerateAzurePricing(currency)
	if err != nil {
		return nil, fmt.Errorf("failed to get Azure pricing: %w", err)
	}
	combinedSet.NodePricing = append(combinedSet.NodePricing, azureSet.NodePricing...)
	combinedSet.PersistentVolumePricing = append(combinedSet.PersistentVolumePricing, azureSet.PersistentVolumePricing...)
	combinedSet.ServicePricing = append(combinedSet.ServicePricing, azureSet.ServicePricing...)
	log.Infof("Added %d Azure node pricing entries, %d service pricing entries", len(azureSet.NodePricing), len(azureSet.ServicePricing))

	// GCP does NOT support CNY
	if currency != "CNY" {
		gcpSet, err := GenerateGCPPricing(currency)
		if err != nil {
			return nil, fmt.Errorf("failed to get GCP pricing: %w", err)
		}
		combinedSet.NodePricing = append(combinedSet.NodePricing, gcpSet.NodePricing...)
		combinedSet.PersistentVolumePricing = append(combinedSet.PersistentVolumePricing, gcpSet.PersistentVolumePricing...)
		combinedSet.ServicePricing = append(combinedSet.ServicePricing, gcpSet.ServicePricing...)
		log.Infof("Added %d GCP node pricing entries, %d service pricing entries", len(gcpSet.NodePricing), len(gcpSet.ServicePricing))
	}

	// Sort the combined set to ensure deterministic output
	combinedSet.Sort()

	log.Infof("Generated combined pricing set with %d total node entries, %d volume entries, %d service entries",
		len(combinedSet.NodePricing), len(combinedSet.PersistentVolumePricing), len(combinedSet.ServicePricing))

	return combinedSet, nil
}