feat(log): wire log db, metrics, and body toggle

internal/api/stats_handler.go

@@ -3,6 +3,7 @@ package api
 import (
 	"fmt"
 	"net/http"
+	"sort"
 	"strings"
 	"time"
 
@@ -59,9 +60,25 @@ func (h *MasterHandler) GetSelfStats(c *gin.Context) {
 		return
 	}
 
-	base := h.db.Model(&model.LogRecord{}).
-		Joins("JOIN keys ON keys.id = log_records.key_id").
-		Where("keys.master_id = ?", m.ID)
+	logDB := h.logDBConn()
+	base := logDB.Model(&model.LogRecord{})
+	if logDB == h.db {
+		base = base.Joins("JOIN keys ON keys.id = log_records.key_id").
+			Where("keys.master_id = ?", m.ID)
+	} else {
+		var keyIDs []uint
+		if err := h.db.Model(&model.Key{}).
+			Where("master_id = ?", m.ID).
+			Pluck("id", &keyIDs).Error; err != nil {
+			c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load master keys", "details": err.Error()})
+			return
+		}
+		if len(keyIDs) == 0 {
+			base = base.Where("1 = 0")
+		} else {
+			base = base.Where("log_records.key_id IN ?", keyIDs)
+		}
+	}
 	base = applyStatsRange(base, rng)
 
 	totalRequests, totalTokens, err := aggregateTotals(base)
@@ -152,7 +169,8 @@ func (h *AdminHandler) GetAdminStats(c *gin.Context) {
 		return
 	}
 
-	base := h.db.Model(&model.LogRecord{})
+	logDB := h.logDBConn()
+	base := logDB.Model(&model.LogRecord{})
 	base = applyStatsRange(base, rng)
 
 	totalRequests, totalTokens, err := aggregateTotals(base)
@@ -162,13 +180,22 @@ func (h *AdminHandler) GetAdminStats(c *gin.Context) {
 	}
 
 	var byMaster []MasterUsageAgg
-	if err := base.Session(&gorm.Session{}).
-		Joins("JOIN keys ON keys.id = log_records.key_id").
-		Select("keys.master_id as master_id, COUNT(*) as requests, COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
-		Group("keys.master_id").
-		Scan(&byMaster).Error; err != nil {
-		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to group by master", "details": err.Error()})
-		return
+	if logDB == h.db {
+		if err := base.Session(&gorm.Session{}).
+			Joins("JOIN keys ON keys.id = log_records.key_id").
+			Select("keys.master_id as master_id, COUNT(*) as requests, COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
+			Group("keys.master_id").
+			Scan(&byMaster).Error; err != nil {
+			c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to group by master", "details": err.Error()})
+			return
+		}
+	} else {
+		var err error
+		byMaster, err = aggregateByMasterFromLogs(base, h.db)
+		if err != nil {
+			c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to group by master", "details": err.Error()})
+			return
+		}
 	}
 
 	var byProvider []ProviderUsageAgg
@@ -268,3 +295,69 @@ func aggregateTotals(q *gorm.DB) (int64, int64, error) {
 	}
 	return totalRequests, t.Tokens, nil
 }
+
+func aggregateByMasterFromLogs(base *gorm.DB, mainDB *gorm.DB) ([]MasterUsageAgg, error) {
+	if base == nil || mainDB == nil {
+		return nil, nil
+	}
+	var byKey []KeyUsageStat
+	if err := base.Session(&gorm.Session{}).
+		Select("log_records.key_id as key_id, COUNT(*) as requests, COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
+		Group("log_records.key_id").
+		Scan(&byKey).Error; err != nil {
+		return nil, err
+	}
+	if len(byKey) == 0 {
+		return []MasterUsageAgg{}, nil
+	}
+
+	keyIDs := make([]uint, 0, len(byKey))
+	for _, row := range byKey {
+		if row.KeyID > 0 {
+			keyIDs = append(keyIDs, row.KeyID)
+		}
+	}
+	if len(keyIDs) == 0 {
+		return []MasterUsageAgg{}, nil
+	}
+
+	type keyMaster struct {
+		ID       uint
+		MasterID uint
+	}
+	var keyMasters []keyMaster
+	if err := mainDB.Model(&model.Key{}).
+		Select("id, master_id").
+		Where("id IN ?", keyIDs).
+		Scan(&keyMasters).Error; err != nil {
+		return nil, err
+	}
+	keyToMaster := make(map[uint]uint, len(keyMasters))
+	for _, row := range keyMasters {
+		keyToMaster[row.ID] = row.MasterID
+	}
+
+	agg := make(map[uint]*MasterUsageAgg)
+	for _, row := range byKey {
+		masterID := keyToMaster[row.KeyID]
+		if masterID == 0 {
+			continue
+		}
+		entry, ok := agg[masterID]
+		if !ok {
+			entry = &MasterUsageAgg{MasterID: masterID}
+			agg[masterID] = entry
+		}
+		entry.Requests += row.Requests
+		entry.Tokens += row.Tokens
+	}
+
+	out := make([]MasterUsageAgg, 0, len(agg))
+	for _, row := range agg {
+		out = append(out, *row)
+	}
+	sort.Slice(out, func(i, j int) bool {
+		return out[i].MasterID < out[j].MasterID
+	})
+	return out, nil
+}