@@ -3,193 +3,207 @@ package provider
33import (
44 "context"
55 "fmt"
6+ "math"
67 "net/http"
7- "net/url"
8- "strconv"
98 "time"
109 "watchAlert/internal/models"
11- utilsHttp "watchAlert/pkg/tools"
10+ "watchAlert/pkg/tools"
11+
12+ "github.com/prometheus/client_golang/api"
13+ v1 "github.com/prometheus/client_golang/api/prometheus/v1"
14+ "github.com/prometheus/common/model"
1215
1316 "github.com/zeromicro/go-zero/core/logc"
1417)
1518
1619type PrometheusProvider struct {
20+ client v1.API
1721 ExternalLabels map [string ]interface {}
1822 Address string
1923 Username string
2024 Password string
2125 Headers map [string ]string
26+ Timeout int64
27+ }
28+
29+ // authenticatedTransport 包装 http.RoundTripper 以添加认证头和额外的headers
30+ type authenticatedTransport struct {
31+ Transport http.RoundTripper
32+ Username string
33+ Password string
34+ Headers map [string ]string
35+ }
36+
37+ // RoundTrip 实现 http.RoundTripper 接口
38+ func (t * authenticatedTransport ) RoundTrip (req * http.Request ) (* http.Response , error ) {
39+ if t .Username != "" && t .Password != "" {
40+ req .SetBasicAuth (t .Username , t .Password )
41+ }
42+
43+ for key , value := range t .Headers {
44+ req .Header .Set (key , value )
45+ }
46+
47+ return t .Transport .RoundTrip (req )
2248}
2349
2450func NewPrometheusClient (ds models.AlertDataSource ) (MetricsFactoryProvider , error ) {
51+ transport := & http.Transport {
52+ Proxy : http .ProxyFromEnvironment ,
53+ MaxIdleConns : 100 ,
54+ MaxIdleConnsPerHost : 10 ,
55+ IdleConnTimeout : 90 * time .Second ,
56+ }
57+
58+ var roundTripper http.RoundTripper = transport
59+ if ds .Auth .User != "" || ds .Auth .Pass != "" || len (ds .HTTP .Headers ) > 0 {
60+ roundTripper = & authenticatedTransport {
61+ Transport : transport ,
62+ Username : ds .Auth .User ,
63+ Password : ds .Auth .Pass ,
64+ Headers : ds .HTTP .Headers ,
65+ }
66+ }
67+
68+ clientConfig := api.Config {
69+ Address : ds .HTTP .URL ,
70+ RoundTripper : roundTripper ,
71+ }
72+
73+ client , err := api .NewClient (clientConfig )
74+ if err != nil {
75+ return nil , err
76+ }
77+
2578 return PrometheusProvider {
79+ client : v1 .NewAPI (client ),
2680 Address : ds .HTTP .URL ,
2781 ExternalLabels : ds .Labels ,
2882 Username : ds .Auth .User ,
2983 Password : ds .Auth .Pass ,
3084 Headers : ds .HTTP .Headers ,
85+ Timeout : ds .HTTP .Timeout ,
3186 }, nil
3287}
3388
3489type QueryResponse struct {
35- Status string `json:"status"`
36- VMData VMData `json:"data"`
90+ Status string `json:"status"`
91+ MetricData MetricData `json:"data"`
3792}
3893
39- type VMData struct {
40- VMResult [] VMResult `json:"result"`
41- ResultType string `json:"resultType"`
94+ type MetricData struct {
95+ MetricResult [] MetricResult `json:"result"`
96+ ResultType string `json:"resultType"`
4297}
4398
44- type VMResult struct {
99+ type MetricResult struct {
45100 Metric map [string ]interface {} `json:"metric"`
46101 Value []interface {} `json:"value"`
47- Values [][]interface {} `json:"values"` // for range query
102+ Values [][]interface {} `json:"values"`
48103}
49104
50105func (v PrometheusProvider ) Query (promQL string ) ([]Metrics , error ) {
51- params := url.Values {}
52- params .Add ("query" , promQL )
53- params .Add ("time" , strconv .FormatInt (time .Now ().Unix (), 10 ))
54- fullURL := fmt .Sprintf ("%s%s?%s" , v .Address , "/api/v1/query" , params .Encode ())
55-
56- // 创建带认证的HTTP请求
57- var headers = make (map [string ]string )
58- for key , value := range v .Headers {
59- headers [key ] = value
60- }
61- for key , value := range utilsHttp .CreateBasicAuthHeader (v .Username , v .Password ) {
62- headers [key ] = value
63- }
64-
65- resp , err := utilsHttp .Get (headers , fullURL , 10 )
106+ ctx , cancel := context .WithTimeout (context .Background (), time .Duration (v .Timeout )* time .Second )
107+ defer cancel ()
108+ result , _ , err := v .client .Query (ctx , promQL , time .Now (), v1 .WithTimeout (time .Duration (v .Timeout )* time .Second ))
66109 if err != nil {
67- logc .Error (context .Background (), "Prometheus query failed" , "error" , err )
68- return nil , fmt .Errorf ("query failed: %w" , err )
110+ return nil , err
69111 }
70- defer resp .Body .Close ()
71-
72- if resp .StatusCode != http .StatusOK {
73- return nil , fmt .Errorf ("unexpected status code: %d" , resp .StatusCode )
74- }
75-
76- var vmRespBody QueryResponse
77- if err := utilsHttp .ParseReaderBody (resp .Body , & vmRespBody ); err != nil {
78- logc .Error (context .Background (), "Parse response failed" , "error" , err )
79- return nil , fmt .Errorf ("parse response failed: %w" , err )
80- }
81-
82- return Vectors (vmRespBody .VMData .VMResult ), nil
112+ return Vectors (result ), nil
83113}
84114
85115func (v PrometheusProvider ) QueryRange (promQL string , start , end time.Time , step time.Duration ) ([]Metrics , error ) {
86- params := url.Values {}
87- params .Add ("query" , promQL )
88- params .Add ("start" , strconv .FormatInt (start .Unix (), 10 ))
89- params .Add ("end" , strconv .FormatInt (end .Unix (), 10 ))
90- params .Add ("step" , fmt .Sprintf ("%.0fs" , step .Seconds ()))
91- fullURL := fmt .Sprintf ("%s%s?%s" , v .Address , "/api/v1/query_range" , params .Encode ())
116+ ctx , cancel := context .WithTimeout (context .Background (), time .Duration (v .Timeout )* time .Second )
117+ defer cancel ()
92118
93- var headers = make (map [string ]string )
94- for key , value := range v .Headers {
95- headers [key ] = value
96- }
97- for key , value := range utilsHttp .CreateBasicAuthHeader (v .Username , v .Password ) {
98- headers [key ] = value
119+ r := v1.Range {
120+ Start : start ,
121+ End : end ,
122+ Step : step ,
99123 }
100124
101- resp , err := utilsHttp . Get ( headers , fullURL , 30 )
125+ result , _ , err := v . client . QueryRange ( ctx , promQL , r , v1 . WithTimeout ( time . Duration ( v . Timeout ) * time . Second ) )
102126 if err != nil {
103- logc .Error (context .Background (), "Prometheus query_range failed" , "error" , err )
104- return nil , fmt .Errorf ("query_range failed: %w" , err )
127+ return nil , err
105128 }
106- defer resp .Body .Close ()
107129
108- if resp .StatusCode != http .StatusOK {
109- return nil , fmt .Errorf ("unexpected status code: %d" , resp .StatusCode )
110- }
111-
112- var vmRespBody QueryResponse
113- if err := utilsHttp .ParseReaderBody (resp .Body , & vmRespBody ); err != nil {
114- logc .Error (context .Background (), "Parse response failed" , "error" , err )
115- return nil , fmt .Errorf ("parse response failed: %w" , err )
116- }
117-
118- return Matrix (vmRespBody .VMData .VMResult ), nil
130+ return Matrix (result ), nil
119131}
120132
121- func Vectors (res [] VMResult ) []Metrics {
133+ func Vectors (value model. Value ) []Metrics {
122134 var vectors []Metrics
123- for _ , item := range res {
124- if len ( item . Value ) < 2 {
125- continue
126- }
135+ items , ok := value .(model. Vector )
136+ if ! ok {
137+ return [] Metrics {}
138+ }
127139
128- timestamp , ok1 := item .Value [0 ].(float64 )
129- valueStr , ok2 := item .Value [1 ].(string )
130- if ! ok1 || ! ok2 {
131- logc .Error (context .Background (), "Invalid value format" )
140+ for _ , item := range items {
141+ if math .IsNaN (float64 (item .Value )) {
142+ logc .Infof (context .Background (), "Skipping NaN or Inf value: %v" , item .Value )
132143 continue
133144 }
134145
135- valueFloat , err := strconv .ParseFloat (valueStr , 64 )
136- if err != nil {
137- logc .Error (context .Background (), "Value conversion failed" , "error" , err )
138- continue
146+ var metric = make (map [string ]interface {})
147+ for k , v := range item .Metric {
148+ metric [string (k )] = string (v )
139149 }
140150
141151 vectors = append (vectors , Metrics {
142- Metric : item . Metric ,
143- Value : valueFloat ,
144- Timestamp : timestamp ,
152+ Metric : metric ,
153+ Value : float64 ( item . Value ) ,
154+ Timestamp : float64 ( item . Timestamp ) ,
145155 })
146156 }
157+
147158 return vectors
148159}
149160
150- // vmMatrix 将 Prometheus QueryRange 结果转换为 Metrics 列表
151- func Matrix (res [] VMResult ) []Metrics {
161+ // Matrix 将 Prometheus QueryRange 结果转换为 Metrics 列表
162+ func Matrix (value model. Value ) []Metrics {
152163 var metrics []Metrics
153- for _ , item := range res {
154- // 遍历每个时间序列的所有时间点
155- for _ , value := range item .Values {
156- if len (value ) < 2 {
157- continue
158- }
164+ matrix , ok := value .(model.Matrix )
165+ if ! ok {
166+ return []Metrics {}
167+ }
159168
160- timestamp , ok1 := value [0 ].(float64 )
161- valueStr , ok2 := value [1 ].(string )
162- if ! ok1 || ! ok2 {
163- logc .Error (context .Background (), "Invalid value format" )
164- continue
165- }
169+ for _ , stream := range matrix {
170+ var metric = make (map [string ]interface {})
171+ for k , v := range stream .Metric {
172+ metric [string (k )] = string (v )
173+ }
166174
167- valueFloat , err := strconv .ParseFloat (valueStr , 64 )
168- if err != nil {
169- logc .Error (context .Background (), "Value conversion failed" , "error" , err )
175+ for _ , value := range stream .Values {
176+ if math .IsNaN (float64 (value .Value )) {
170177 continue
171178 }
172179
173180 metrics = append (metrics , Metrics {
174- Metric : item . Metric ,
175- Value : valueFloat ,
176- Timestamp : timestamp ,
181+ Timestamp : float64 ( value . Timestamp ) ,
182+ Value : float64 ( value . Value ) ,
183+ Metric : metric ,
177184 })
178185 }
179186 }
187+
180188 return metrics
181189}
182190
183191func (v PrometheusProvider ) Check () (bool , error ) {
184- res , err := utilsHttp .Get (utilsHttp .CreateBasicAuthHeader (v .Username , v .Password ), v .Address + "/api/v1/query?query=1%2B1" , 10 )
192+ var headers map [string ]string
193+ checkURL := v .Address + "/api/v1/query?query=1%2B1"
194+ if v .Username != "" && v .Password != "" {
195+ headers = tools .CreateBasicAuthHeader (v .Username , v .Password )
196+ }
197+ headers = tools .MergeHeaders (headers , v .Headers )
198+ res , err := tools .Get (headers , checkURL , int (v .Timeout ))
185199 if err != nil {
186- logc .Error (context .Background (), fmt . Errorf ( "health check failed: %w "err ) )
200+ logc .Errorf (context .Background (), "Health check failed, URL : %s, Error: %v "checkURL , err )
187201 return false , fmt .Errorf ("health check failed: %w" , err )
188202 }
189203 defer res .Body .Close ()
190204
191205 if res .StatusCode != http .StatusOK {
192- logc .Error (context .Background (), fmt . Errorf ( " unhealthy status: %d"res .StatusCode ) )
206+ logc .Errorf (context .Background (), "Health check received unhealthy status: %d, URL: %s "res .StatusCode , checkURL )
193207 return false , fmt .Errorf ("unhealthy status: %d" , res .StatusCode )
194208 }
195209 return true , nil
0 commit comments