objectstorage.go

package sos

import (
	"context"
	"fmt"
	"github.com/exoscale/egoscale/v2/oapi"
	"time"

	pipeline "github.com/ccremer/go-command-pipeline"
	egoscale "github.com/exoscale/egoscale/v2"
	db "github.com/vshn/exoscale-metrics-collector/pkg/database"
	exoscalev1 "github.com/vshn/provider-exoscale/apis/exoscale/v1"
	ctrl "sigs.k8s.io/controller-runtime"
	k8s "sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	organizationLabel = "appuio.io/organization"
	namespaceLabel    = "crossplane.io/claim-namespace"
	// Exoscale gathers metrics of its bucket at 6AM UTC
	exoscaleTimeZone    = "UTC"
	exoscaleBillingHour = 6
)

type ObjectStorage struct {
	k8sClient         *k8s.Client
	exoscaleClient    *egoscale.Client
	database          *db.Database
	bucketDetails     []BucketDetail
	aggregatedBuckets map[string]db.AggregatedBucket
	billingDate       time.Time
}

type BucketDetail struct {
	Organization, BucketName, Namespace string
}

func NewObjectStorage(exoscaleClient *egoscale.Client, k8sClient *k8s.Client, databaseURL string) ObjectStorage {
	return ObjectStorage{
		exoscaleClient: exoscaleClient,
		k8sClient:      k8sClient,
		database:       &db.Database{Url: databaseURL},
	}
}

func (o *ObjectStorage) Execute(ctx context.Context) error {
	log := ctrl.LoggerFrom(ctx)
	log.Info("Running metrics collector by step")

	p := pipeline.NewPipeline[context.Context]()
	p.WithSteps(
		p.NewStep("Fetch managed buckets", o.fetchManagedBuckets),
		p.NewStep("Get bucket usage", o.getBucketUsage),
		p.NewStep("Get billing date", o.getBillingDate),
		p.NewStep("Save to database", o.saveToDatabase),
	)
	return p.RunWithContext(ctx)
}

func (o *ObjectStorage) getBucketUsage(ctx context.Context) error {
	log := ctrl.LoggerFrom(ctx)
	log.Info("Fetching bucket usage from Exoscale")
	resp, err := o.exoscaleClient.ListSosBucketsUsageWithResponse(ctx)
	if err != nil {
		return err
	}

	aggregatedBuckets := *getAggregatedBuckets(ctx, resp.JSON200.SosBucketsUsage, &o.bucketDetails)
	o.aggregatedBuckets = aggregatedBuckets
	return nil
}

func (o *ObjectStorage) fetchManagedBuckets(ctx context.Context) error {
	log := ctrl.LoggerFrom(ctx)
	log.Info("Fetching buckets from cluster")

	buckets := exoscalev1.BucketList{}
	client := *o.k8sClient
	log.V(1).Info("Listing buckets from cluster")
	err := client.List(ctx, &buckets)
	if err != nil {
		return fmt.Errorf("cannot list buckets: %w", err)
	}
	o.bucketDetails = *addOrgAndNamespaceToBucket(ctx, &buckets)
	return nil
}

func (o *ObjectStorage) saveToDatabase(ctx context.Context) error {
	log := ctrl.LoggerFrom(ctx)
	log.Info("Creating a database connection")

	log.V(1).Info("Opening database connection")
	err := o.database.OpenConnection()
	if err != nil {
		return err
	}
	defer o.database.CloseConnection()

	log.V(1).Info("Ensuring initial database configuration")
	err = o.database.EnsureInitConfiguration(ctx)
	if err != nil {
		return err
	}

	log.V(1).Info("Saving buckets information usage to database")
	for namespace, aggregatedBucket := range o.aggregatedBuckets {
		err = o.database.EnsureBucketUsage(ctx, namespace, aggregatedBucket, o.billingDate)
		if err != nil {
			log.Info("WARNING: Buckets usage have not been saved to the database", "namespace", namespace, "error", err.Error())
		}
	}

	return nil
}

func (o *ObjectStorage) getBillingDate(_ context.Context) error {
	location, err := time.LoadLocation(exoscaleTimeZone)
	if err != nil {
		return fmt.Errorf("cannot initialize location from time zone %s: %w", location, err)
	}
	now := time.Now().In(location)
	previousDay := now.Day() - 1
	o.billingDate = time.Date(now.Year(), now.Month(), previousDay, exoscaleBillingHour, 0, 0, 0, now.Location())
	return nil
}

func getAggregatedBuckets(ctx context.Context, sosBucketsUsage *[]oapi.SosBucketUsage, bucketDetails *[]BucketDetail) *map[string]db.AggregatedBucket {
	log := ctrl.LoggerFrom(ctx)
	log.Info("Aggregating buckets by namespace")
	aggregatedBuckets := make(map[string]db.AggregatedBucket)
	for _, bucketDetail := range *bucketDetails {
		log.V(1).Info("Checking bucket", "bucket", bucketDetail.BucketName)
		storageUsed := int64(0)
		usageBucketExists := false

		// Match managed kubernetes resource bucket to the exoscale usage bucket to retrieve the amount of storage used
		for _, bucketUsage := range *sosBucketsUsage {
			if bucketDetail.BucketName == *bucketUsage.Name {
				log.V(1).Info("Found exoscale bucket usage", "bucket", bucketUsage.Name, "bucket size", bucketUsage.Name)
				storageUsed = *bucketUsage.Size
				usageBucketExists = true
				break
			}
		}

		// Create AggregatedBucket object only if the object exists in exoscale
		if usageBucketExists {
			aggregatedBucket, exists := aggregatedBuckets[bucketDetail.Namespace]
			if exists {
				log.V(1).Info("AggregatedBucket exists, adding its used storage to the current AggregatedBucket",
					"bucket", bucketDetail.BucketName, "storage to be added", storageUsed)
				aggregatedBucket.StorageUsed += storageUsed
			} else {
				log.V(1).Info("Adding new AggregatedBucket object", "bucket", bucketDetail.BucketName, "bucket storage", storageUsed)
				aggregatedBuckets[bucketDetail.Namespace] = db.AggregatedBucket{
					Organization: bucketDetail.Organization,
					StorageUsed:  storageUsed,
				}
			}
		} else {
			log.Info("Could not find any bucket on exoscale", "bucket", bucketDetail.BucketName)
		}
	}
	return &aggregatedBuckets
}

func addOrgAndNamespaceToBucket(ctx context.Context, buckets *exoscalev1.BucketList) *[]BucketDetail {
	log := ctrl.LoggerFrom(ctx)
	log.V(1).Info("Gathering more information from buckets")

	var bucketDetails []BucketDetail
	for _, bucket := range buckets.Items {
		bucketDetail := BucketDetail{
			BucketName: bucket.Spec.ForProvider.BucketName,
		}
		if organization, exist := bucket.ObjectMeta.Labels[organizationLabel]; exist {
			bucketDetail.Organization = organization
		} else {
			// cannot get organization from bucket
			log.Info("Organization label is missing in bucket, skipping...",
				"label", organizationLabel,
				"bucket", bucket.Name)
			continue
		}
		if namespace, exist := bucket.ObjectMeta.Labels[namespaceLabel]; exist {
			bucketDetail.Namespace = namespace
		} else {
			// cannot get namespace from bucket
			log.Info("Namespace label is missing in bucket, skipping...",
				"label", namespaceLabel,
				"bucket", bucket.Name)
			continue
		}
		log.V(1).Info("Added namespace and organization to bucket",
			"bucket", bucket.Name,
			"namespace", bucketDetail.Namespace,
			"organization", bucketDetail.Organization)
		bucketDetails = append(bucketDetails, bucketDetail)
	}
	return &bucketDetails
}