Luminary/internal/store/sqlite/repository.go

package sqlite

import (
	"time"

	"github.com/rose_cat707/luminary/internal/auth"
	"github.com/rose_cat707/luminary/internal/model"
	"github.com/rose_cat707/luminary/internal/settings"
	"github.com/rose_cat707/luminary/internal/store/cache"
	"gorm.io/gorm"
)

type Repository struct {
	db    *gorm.DB
	cache *cache.Store
}

func NewRepository(store *Store, c *cache.Store) *Repository {
	return &Repository{db: store.DB(), cache: c}
}

func (r *Repository) BootstrapFromDB() error {
	var providers []model.Provider
	if err := r.db.Find(&providers).Error; err != nil {
		return err
	}
	r.cache.LoadProviders(providers)

	var keys []model.ProviderKey
	if err := r.db.Find(&keys).Error; err != nil {
		return err
	}
	r.cache.LoadProviderKeys(keys)

	var models []model.ModelEntry
	if err := r.db.Find(&models).Error; err != nil {
		return err
	}
	r.cache.LoadModels(models)

	var ingress []model.IngressKey
	if err := r.db.Find(&ingress).Error; err != nil {
		return err
	}
	r.cache.LoadIngressKeys(ingress)

	var rules []model.IPRule
	if err := r.db.Find(&rules).Error; err != nil {
		return err
	}
	r.cache.LoadIPRules(rules)

	var routing []model.RoutingRule
	if err := r.db.Find(&routing).Error; err != nil {
		return err
	}
	r.cache.LoadRoutingRules(routing)
	r.PruneUsageRecords(model.MaxUsageLogRecords)
	return nil
}

func (r *Repository) EnsureAdmin() error {
	var count int64
	r.db.Model(&model.AdminUser{}).Count(&count)
	if count > 0 {
		return nil
	}
	username, password := settings.LegacyAdminCredentials()
	hash, err := auth.HashPassword(password)
	if err != nil {
		return err
	}
	return r.db.Create(&model.AdminUser{Username: username, PasswordHash: hash}).Error
}

func (r *Repository) FlushUsage(deltas []cache.UsageDelta) {
	for _, d := range deltas {
		r.db.Create(&model.UsageRecord{
			IngressKeyID:  d.IngressKeyID,
			ProviderID:    d.ProviderID,
			ProviderKeyID: d.ProviderKeyID,
			ClientIP:      d.ClientIP,
			Endpoint:      d.Endpoint,
			Model:         d.Model,
			TokensIn:      d.TokensIn,
			TokensOut:     d.TokensOut,
			Cost:          d.Cost,
			LatencyMs:     d.LatencyMs,
			Status:        d.Status,
			Stream:        d.Stream,
			ErrorMsg:      d.ErrorMsg,
			RequestBody:   d.RequestBody,
			ResponseBody:  d.ResponseBody,
			CreatedAt:     time.Now(),
		})
		if d.IngressKeyID > 0 {
			r.db.Model(&model.IngressKey{}).Where("id = ?", d.IngressKeyID).Updates(map[string]any{
				"budget_used":   gorm.Expr("budget_used + ?", d.BudgetUsed),
				"request_count": gorm.Expr("request_count + ?", d.Requests),
				"token_count":   gorm.Expr("token_count + ?", d.Tokens),
			})
		}
		if d.ProviderKeyID > 0 {
			r.db.Model(&model.ProviderKey{}).Where("id = ?", d.ProviderKeyID).Updates(map[string]any{
				"requests_today": gorm.Expr("requests_today + ?", d.Requests),
				"tokens_today":   gorm.Expr("tokens_today + ?", d.Tokens),
			})
		}
	}
	r.PruneUsageRecords(model.MaxUsageLogRecords)
}

func (r *Repository) RecordRequestLog(rec model.UsageRecord) {
	if rec.CreatedAt.IsZero() {
		rec.CreatedAt = time.Now()
	}
	r.db.Create(&rec)
	r.PruneUsageRecords(model.MaxUsageLogRecords)
}

func (r *Repository) PruneUsageRecords(max int) {
	if max <= 0 {
		return
	}
	var count int64
	r.db.Model(&model.UsageRecord{}).Count(&count)
	if count <= int64(max) {
		return
	}
	var keepIDs []uint
	r.db.Model(&model.UsageRecord{}).Order("created_at desc").Limit(max).Pluck("id", &keepIDs)
	if len(keepIDs) == 0 {
		return
	}
	r.db.Where("id NOT IN ?", keepIDs).Delete(&model.UsageRecord{})
}