internal/budget/factory.go

Rama main · réplica de lectura del repositorio original.

package budget

import (
	"context"
	"errors"
	"fmt"
	"sync"

	"go.mongodb.org/mongo-driver/v2/mongo"

	"github.com/enterpilot/gomodel/config"
	"github.com/enterpilot/gomodel/internal/storage"
	"github.com/enterpilot/gomodel/internal/storage/sqlx"
)

type Result struct {
	Service *Service
	Store   Store

	closeOnce sync.Once
	closeErr  error
}

func (r *Result) Close() error {
	if r == nil {
		return nil
	}
	r.closeOnce.Do(func() {
		var errs []error
		if r.Store != nil {
			if err := r.Store.Close(); err != nil {
				errs = append(errs, fmt.Errorf("store close: %w", err))
			}
		}
		if len(errs) > 0 {
			r.closeErr = fmt.Errorf("close errors: %w", errors.Join(errs...))
		}
	})
	return r.closeErr
}

// New builds the budget service. quotaTemplatesEnabled and billedBudgetsEnabled
// report the ext capabilities; budgets on billed cost also need
// pricing.markups_enabled, which New reads from cfg.
func New(ctx context.Context, cfg *config.Config, shared storage.Storage, quotaTemplatesEnabled, billedBudgetsEnabled bool) (*Result, error) {
	if cfg == nil {
		return nil, fmt.Errorf("config is required")
	}
	if !cfg.Budgets.Enabled {
		return &Result{}, nil
	}
	if shared == nil {
		return nil, fmt.Errorf("shared storage is required")
	}
	return newResult(ctx, cfg, shared, quotaTemplatesEnabled, billedBudgetsEnabled)
}

func newResult(ctx context.Context, cfg *config.Config, storeConn storage.Storage, quotaTemplatesEnabled, billedBudgetsEnabled bool) (*Result, error) {
	store, err := createStore(ctx, storeConn)
	if err != nil {
		return nil, err
	}
	caps := capabilities{
		quotaTemplates: quotaTemplatesEnabled,
		billedBudgets:  billedBudgetsEnabled && cfg.Pricing.MarkupsEnabled,
	}
	// Apply the configured budgets before the service loads the stored set, so
	// a config change that drops or switches a budget takes effect even when
	// the stored row it replaces needs a capability that is gone.
	if err := seedConfiguredBudgets(ctx, store, cfg.Budgets, caps); err != nil {
		return nil, err
	}
	service, err := NewService(ctx, store,
		WithQuotaTemplates(caps.quotaTemplates),
		WithBilledBudgets(caps.billedBudgets),
		WithBilledCacheHits(cfg.Pricing.MarkupsEnabled && cfg.Pricing.BillCacheHits),
		WithProviderBilledCacheHits(cfg.Budgets.ProviderBilledCacheHits),
	)
	if err != nil {
		return nil, err
	}
	return &Result{Service: service, Store: store}, nil
}

func createStore(ctx context.Context, store storage.Storage) (Store, error) {
	return storage.ResolveSQLBackend[Store](
		ctx,
		store,
		func(db sqlx.DB) (Store, error) { return NewSQLStore(ctx, db) },
		func(db *mongo.Database) (Store, error) { return NewMongoDBStore(ctx, db) },
	)
}

// seedConfiguredBudgets replaces the config-sourced budgets in store with the
// ones cfg declares, refusing any that need a capability caps lacks.
func seedConfiguredBudgets(ctx context.Context, store Store, cfg config.BudgetsConfig, caps capabilities) error {
	budgets := make([]Budget, 0, len(cfg.UserPaths)+len(cfg.Labels)+len(cfg.Providers))
	for _, entry := range cfg.UserPaths {
		seeded, err := seedBudgets(ScopeUserPath, entry.Path, entry.PerChild, entry.Limits)
		if err != nil {
			return err
		}
		budgets = append(budgets, seeded...)
	}
	for _, entry := range cfg.Labels {
		seeded, err := seedBudgets(ScopeLabel, entry.Label, false, entry.Limits)
		if err != nil {
			return err
		}
		budgets = append(budgets, seeded...)
	}
	for _, entry := range cfg.Providers {
		seeded, err := seedBudgets(ScopeProvider, entry.Name, false, entry.Limits)
		if err != nil {
			return err
		}
		budgets = append(budgets, seeded...)
	}
	if err := caps.validate(budgets); err != nil {
		return err
	}
	return store.ReplaceConfigBudgets(ctx, budgets)
}

func seedBudgets(scope Scope, rawSubject string, perChild bool, limits []config.BudgetLimitConfig) ([]Budget, error) {
	subject, err := NormalizeSubject(scope, rawSubject)
	if err != nil {
		return nil, fmt.Errorf("invalid budget %s %q: %w", scope, rawSubject, err)
	}
	budgets := make([]Budget, 0, len(limits))
	for limitIdx, limit := range limits {
		seconds := limit.PeriodSeconds
		if seconds <= 0 {
			parsed, ok := PeriodSeconds(limit.Period)
			if !ok {
				return nil, fmt.Errorf("invalid budget period for %s %q limit %d: %q", scope, subject, limitIdx, limit.Period)
			}
			seconds = parsed
		}
		costBasis, err := NormalizeCostBasis(limit.CostBasis)
		if err != nil {
			return nil, fmt.Errorf("invalid budget cost basis for %s %q limit %d: %w", scope, subject, limitIdx, err)
		}
		budgets = append(budgets, Budget{
			Scope:         scope,
			Subject:       subject,
			PerChild:      perChild || limit.PerChild,
			CostBasis:     costBasis,
			PeriodSeconds: seconds,
			Amount:        limit.Amount,
			Source:        SourceConfig,
		})
	}
	return budgets, nil
}