Browse Source

feat(metrics): 渠道状态自定义 Collector

master
fengsilin 3 weeks ago
parent
commit
87628f5469
2 changed files with 157 additions and 0 deletions
  1. +64
    -0
      common/metrics/collector.go
  2. +93
    -0
      common/metrics/collector_test.go

+ 64
- 0
common/metrics/collector.go View File

@@ -0,0 +1,64 @@
package metrics

import (
"fmt"

"github.com/prometheus/client_golang/prometheus"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
)

// ChannelCollector 从数据库查询渠道状态并暴露为 Prometheus Gauge
type ChannelCollector struct {
channelStatus *prometheus.Desc
}

// NewChannelCollector 创建渠道状态 Collector
func NewChannelCollector() *ChannelCollector {
return &ChannelCollector{
channelStatus: prometheus.NewDesc(
"newapi_channel_status",
"Channel enabled status (1=enabled, 0=disabled)",
[]string{"channel", "channel_type"},
nil,
),
}
}

func (c *ChannelCollector) Describe(ch chan<- *prometheus.Desc) {
ch <- c.channelStatus
}

func (c *ChannelCollector) Collect(ch chan<- prometheus.Metric) {
if model.DB == nil {
return
}

var channels []struct {
Id int
Type int
Status int
}

err := model.DB.Table("channels").Select("id, type, status").Find(&channels).Error
if err != nil {
return
}

for _, ch2 := range channels {
value := float64(0)
if ch2.Status == common.ChannelStatusEnabled {
value = 1
}
m, err := prometheus.NewConstMetric(
c.channelStatus,
prometheus.GaugeValue,
value,
fmt.Sprintf("%d", ch2.Id),
fmt.Sprintf("%d", ch2.Type),
)
if err == nil {
ch <- m
}
}
}

+ 93
- 0
common/metrics/collector_test.go View File

@@ -0,0 +1,93 @@
package metrics

import (
"testing"

"github.com/glebarez/sqlite"
"github.com/prometheus/client_golang/prometheus"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
)

func setupCollectorTestDB(t *testing.T) *gorm.DB {
t.Helper()
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
sqlDB, _ := db.DB()
sqlDB.SetMaxOpenConns(1)

origDB := model.DB
model.DB = db
common.UsingSQLite = true
common.RedisEnabled = false

require.NoError(t, db.AutoMigrate(&model.Channel{}))

t.Cleanup(func() {
model.DB = origDB
sqlDB.Close()
})

return db
}

func TestChannelCollector_CollectsChannelStatus(t *testing.T) {
db := setupCollectorTestDB(t)

ch1 := model.Channel{Id: 1, Name: "openai-1", Type: 1, Status: common.ChannelStatusEnabled}
ch2 := model.Channel{Id: 2, Name: "claude-1", Type: 14, Status: common.ChannelStatusEnabled}
ch3 := model.Channel{Id: 3, Name: "gemini-1", Type: 24, Status: 3}
require.NoError(t, db.Create(&ch1).Error)
require.NoError(t, db.Create(&ch2).Error)
require.NoError(t, db.Create(&ch3).Error)

collector := NewChannelCollector()
reg := prometheus.NewRegistry()
reg.MustRegister(collector)

mfs, err := reg.Gather()
require.NoError(t, err)

var found bool
for _, mf := range mfs {
if mf.GetName() == "newapi_channel_status" {
found = true
metrics := mf.GetMetric()
assert.Len(t, metrics, 3)

statusMap := map[string]float64{}
for _, m := range metrics {
chId := ""
for _, label := range m.GetLabel() {
if label.GetName() == "channel" {
chId = label.GetValue()
}
}
statusMap[chId] = m.GetGauge().GetValue()
}
assert.Equal(t, float64(1), statusMap["1"])
assert.Equal(t, float64(1), statusMap["2"])
assert.Equal(t, float64(0), statusMap["3"])
}
}
assert.True(t, found, "newapi_channel_status metric not found")
}

func TestChannelCollector_EmptyDB(t *testing.T) {
setupCollectorTestDB(t)

collector := NewChannelCollector()
reg := prometheus.NewRegistry()
reg.MustRegister(collector)

mfs, err := reg.Gather()
require.NoError(t, err)

for _, mf := range mfs {
assert.NotEqual(t, "newapi_channel_status", mf.GetName(), "should not have channel status with empty DB")
}
}

Loading…
Cancel
Save