Files
NianAIGC/backend/internal/usage/usage.go

100 lines
3.4 KiB
Go

package usage
import (
"errors"
"fmt"
"sort"
"time"
)
type Preset string
const (
PresetToday Preset = "today"
Preset7Days Preset = "7d"
Preset30Days Preset = "30d"
PresetMonth Preset = "month"
)
var ErrForbidden = errors.New("usage scope forbidden")
var shanghai = time.FixedZone("Asia/Shanghai", 8*60*60)
type DateRange struct {
From, To, StartDate, EndDate string
DayCount int
}
func PresetRange(preset Preset, now time.Time) (DateRange, error) {
today := now.In(shanghai)
start := time.Date(today.Year(), today.Month(), today.Day(), 0, 0, 0, 0, shanghai)
switch preset {
case PresetToday:
case Preset7Days:
start = start.AddDate(0, 0, -6)
case Preset30Days:
start = start.AddDate(0, 0, -29)
case PresetMonth:
start = time.Date(today.Year(), today.Month(), 1, 0, 0, 0, 0, shanghai)
default:
return DateRange{}, fmt.Errorf("unknown usage preset %q", preset)
}
end := time.Date(today.Year(), today.Month(), today.Day(), 0, 0, 0, 0, shanghai)
return DateRange{From: start.Format(time.RFC3339), To: end.AddDate(0, 0, 1).Format(time.RFC3339), StartDate: start.Format("2006-01-02"), EndDate: end.Format("2006-01-02"), DayCount: int(end.Sub(start).Hours()/24) + 1}, nil
}
type Event struct {
ID, OwnerID, JobID, Source, Capability, Provider, ReqKey, AccountUsername, AccountDisplayName, TenantID, OrganizationID, OrganizationName, EstimatedUnit, Currency, CreatedAt string
Quantity int
ChargedAmountFen *int64
}
type Filters struct{ OwnerID, OrganizationID, Capability, Provider, From, To string }
type Requester struct{ AccountID, OrganizationID, Role string }
type Repository interface {
Insert(Event) (Event, bool, error)
List(Filters) ([]Event, error)
}
type Service struct{ Repository Repository }
func (s Service) Record(event Event) (*Event, error) {
if event.Source == "api" || event.Provider == "mock" {
return nil, nil
}
saved, _, err := s.Repository.Insert(event)
if err != nil {
return nil, err
}
return &saved, nil
}
func (s Service) Report(who Requester, filters Filters) ([]Event, error) {
switch who.Role {
case "super_admin":
case "organization_admin":
if filters.OrganizationID != "" && filters.OrganizationID != who.OrganizationID {
return nil, ErrForbidden
}
filters.OrganizationID = who.OrganizationID
case "user":
filters.OwnerID = who.AccountID
default:
filters.OwnerID = who.AccountID
}
events, err := s.Repository.List(filters)
if err != nil {
return nil, err
}
seen := map[string]bool{}
out := make([]Event, 0, len(events))
for _, e := range events {
if e.Source == "api" || e.Provider == "mock" || seen[e.JobID] || filters.OwnerID != "" && e.OwnerID != filters.OwnerID || filters.OrganizationID != "" && e.OrganizationID != filters.OrganizationID || filters.Capability != "" && e.Capability != filters.Capability || filters.Provider != "" && e.Provider != filters.Provider {
continue
}
seen[e.JobID] = true
if who.Role == "organization_admin" {
e.AccountUsername = ""
}
out = append(out, e)
}
sort.SliceStable(out, func(i, j int) bool { return out[i].CreatedAt > out[j].CreatedAt })
return out, nil
}