From 87628f54697cbfd3eaf2f0f83ba6093025b69b1f Mon Sep 17 00:00:00 2001 From: fengsilin Date: Wed, 20 May 2026 18:12:29 +0800 Subject: [PATCH] =?UTF-8?q?feat(metrics):=20=E6=B8=A0=E9=81=93=E7=8A=B6?= =?UTF-8?q?=E6=80=81=E8=87=AA=E5=AE=9A=E4=B9=89=20Collector?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- common/metrics/collector.go | 64 ++++++++++++++++++++++ common/metrics/collector_test.go | 93 ++++++++++++++++++++++++++++++++ 2 files changed, 157 insertions(+) create mode 100644 common/metrics/collector.go create mode 100644 common/metrics/collector_test.go diff --git a/common/metrics/collector.go b/common/metrics/collector.go new file mode 100644 index 0000000..772552c --- /dev/null +++ b/common/metrics/collector.go @@ -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 + } + } +} diff --git a/common/metrics/collector_test.go b/common/metrics/collector_test.go new file mode 100644 index 0000000..44a920d --- /dev/null +++ b/common/metrics/collector_test.go @@ -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") + } +}