Skip to content

Commit d3157f5

Browse files
divyansh-chhabria-dbdivyansh-chhabria-db
authored
Cherrypick for release (#234) (#238)
2 parents d10ba97 + a6d21ab commit d3157f5

3 files changed

Lines changed: 77 additions & 36 deletions

File tree

‎pkg/queryfrontend/query_logger.go‎

Lines changed: 59 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ type UserInfo struct {
3636
Tenant string
3737
ForwardedFor string
3838
UserAgent string
39+
Groups string
40+
Email string
3941
}
4042

4143
// ResponseStats holds statistics extracted from query response.
@@ -91,6 +93,10 @@ func ExtractUserInfoFromHeaders(headers []*RequestHeader) UserInfo {
9193
if userInfo.Source == "" {
9294
userInfo.Source = headerValue
9395
}
96+
case "x-auth-request-groups":
97+
userInfo.Groups = headerValue
98+
case "x-auth-request-email":
99+
userInfo.Email = headerValue
94100
}
95101
}
96102

@@ -102,29 +108,6 @@ func ExtractUserInfoFromHeaders(headers []*RequestHeader) UserInfo {
102108
return userInfo
103109
}
104110

105-
// ExtractEmailFromResponse extracts the email from response headers (works for both range and instant queries).
106-
func ExtractEmailFromResponse(resp queryrange.Response) string {
107-
if resp == nil {
108-
return ""
109-
}
110-
111-
// Check both response types using OR condition
112-
var headers []*queryrange.PrometheusResponseHeader
113-
if promResp, ok := resp.(*queryrange.PrometheusResponse); ok {
114-
headers = promResp.GetHeaders()
115-
} else if promResp, ok := resp.(*queryrange.PrometheusInstantQueryResponse); ok {
116-
headers = promResp.GetHeaders()
117-
}
118-
119-
for _, header := range headers {
120-
if strings.ToLower(header.Name) == "x-auth-request-email" && len(header.Values) > 0 {
121-
return header.Values[0]
122-
}
123-
}
124-
125-
return ""
126-
}
127-
128111
// ConvertStoreMatchers converts internal store matchers to logging format.
129112
func ConvertStoreMatchers(storeMatchers [][]*labels.Matcher) []StoreMatcherSet {
130113
if len(storeMatchers) == 0 {
@@ -174,6 +157,59 @@ func GetResponseStats(resp queryrange.Response) ResponseStats {
174157
return stats
175158
}
176159

160+
// ExtractMetricNames extracts all unique __name__ labels from query response (works for both range and instant queries).
161+
func ExtractMetricNames(resp queryrange.Response) []string {
162+
if resp == nil {
163+
return nil
164+
}
165+
166+
metricNamesMap := make(map[string]struct{})
167+
168+
// Handle range query response (resultType: matrix)
169+
if r, ok := resp.(*queryrange.PrometheusResponse); ok {
170+
for _, stream := range r.Data.Result {
171+
for _, label := range stream.Labels {
172+
if label.Name == "__name__" {
173+
metricNamesMap[label.Value] = struct{}{}
174+
break
175+
}
176+
}
177+
}
178+
} else if r, ok := resp.(*queryrange.PrometheusInstantQueryResponse); ok {
179+
// Handle instant query response - check all result types
180+
if vector := r.Data.Result.GetVector(); vector != nil {
181+
// resultType: vector
182+
for _, sample := range vector.Samples {
183+
for _, label := range sample.Labels {
184+
if label.Name == "__name__" {
185+
metricNamesMap[label.Value] = struct{}{}
186+
break
187+
}
188+
}
189+
}
190+
} else if matrix := r.Data.Result.GetMatrix(); matrix != nil {
191+
// resultType: matrix (subqueries in instant queries)
192+
for _, stream := range matrix.SampleStreams {
193+
for _, label := range stream.Labels {
194+
if label.Name == "__name__" {
195+
metricNamesMap[label.Value] = struct{}{}
196+
break
197+
}
198+
}
199+
}
200+
}
201+
// Scalar and StringSample don't have __name__ labels
202+
}
203+
204+
// Convert map to slice
205+
metricNames := make([]string, 0, len(metricNamesMap))
206+
for name := range metricNamesMap {
207+
metricNames = append(metricNames, name)
208+
}
209+
210+
return metricNames
211+
}
212+
177213
// WriteJSONLogToFile writes query logs to file in JSON format.
178214
func WriteJSONLogToFile(logger log.Logger, writer interface{}, queryLog interface{}, queryType string) error {
179215
if writer == nil {

‎pkg/queryfrontend/queryinstant_logger.go‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ type MetricsInstantQueryLogging struct {
3737
ForwardedFor string `json:"forwarded_for"`
3838
UserAgent string `json:"user_agent"`
3939
EmailId string `json:"email_id"`
40+
Groups string `json:"groups"`
4041
// Query-related fields (instant query specific)
4142
QueryTimestampMs int64 `json:"query_timestamp_ms"` // Query timestamp for instant queries
4243
Path string `json:"path"`
@@ -50,6 +51,8 @@ type MetricsInstantQueryLogging struct {
5051
Analyze bool `json:"analyze"` // Whether query analysis is enabled
5152
Engine string `json:"engine"` // Query engine being used
5253
Stats string `json:"stats"` // Query statistics information
54+
MetricNames []string `json:"metric_names"` // Unique metric names (__name__ labels) in response
55+
Shard string `json:"shard"` // Pantheon shard name
5356
// Store-matcher details
5457
StoreMatchers []StoreMatcherSet `json:"store_matchers"`
5558
}
@@ -134,18 +137,12 @@ func (m *instantQueryLoggingMiddleware) logInstantQuery(req *ThanosQueryInstantR
134137
success := err == nil
135138
userInfo := ExtractUserInfoFromHeaders(req.Headers)
136139

137-
// Extract email from response headers
138-
email := ExtractEmailFromResponse(resp)
139-
140-
// This is to avoid logging queries that come from rule manager.
141-
if userInfo.UserAgent == "Databricks-RuleManager/1.0" {
142-
return
143-
}
144-
145140
// Calculate stats (only for successful queries).
146141
var stats ResponseStats
142+
var metricNames []string
147143
if success && resp != nil {
148144
stats = GetResponseStats(resp)
145+
metricNames = ExtractMetricNames(resp)
149146
}
150147

151148
// Create the instant query log entry.
@@ -166,7 +163,8 @@ func (m *instantQueryLoggingMiddleware) logInstantQuery(req *ThanosQueryInstantR
166163
Tenant: userInfo.Tenant,
167164
ForwardedFor: userInfo.ForwardedFor,
168165
UserAgent: userInfo.UserAgent,
169-
EmailId: email,
166+
EmailId: userInfo.Email,
167+
Groups: userInfo.Groups,
170168
// Query-related fields (instant query specific)
171169
QueryTimestampMs: req.Time,
172170
Path: req.Path,
@@ -180,6 +178,8 @@ func (m *instantQueryLoggingMiddleware) logInstantQuery(req *ThanosQueryInstantR
180178
Analyze: req.Analyze,
181179
Engine: req.Engine,
182180
Stats: req.Stats,
181+
MetricNames: metricNames,
182+
Shard: os.Getenv("PANTHEON_SHARDNAME"),
183183
// Store-matcher details
184184
StoreMatchers: ConvertStoreMatchers(req.StoreMatchers),
185185
}

‎pkg/queryfrontend/queryrange_logger.go‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ type MetricsRangeQueryLogging struct {
3737
ForwardedFor string `json:"forwardedFor"`
3838
UserAgent string `json:"userAgent"`
3939
EmailId string `json:"emailId"`
40+
Groups string `json:"groups"`
4041
// Query-related fields
4142
StartTimestampMs int64 `json:"startTimestampMs"`
4243
EndTimestampMs int64 `json:"endTimestampMs"`
@@ -53,6 +54,8 @@ type MetricsRangeQueryLogging struct {
5354
Engine string `json:"engine"` // Query engine being used
5455
SplitIntervalMs int64 `json:"splitIntervalMs"` // Query splitting interval in milliseconds
5556
Stats string `json:"stats"` // Query statistics information
57+
MetricNames []string `json:"metricNames"` // Unique metric names (__name__ labels) in response
58+
Shard string `json:"shard"` // Pantheon shard name
5659
// Store-matcher details
5760
StoreMatchers []StoreMatcherSet `json:"storeMatchers"`
5861
}
@@ -137,13 +140,12 @@ func (m *rangeQueryLoggingMiddleware) logRangeQuery(req *ThanosQueryRangeRequest
137140
success := err == nil
138141
userInfo := ExtractUserInfoFromHeaders(req.Headers)
139142

140-
// Extract email from response headers
141-
email := ExtractEmailFromResponse(resp)
142-
143143
// Calculate stats (only for successful queries).
144144
var stats ResponseStats
145+
var metricNames []string
145146
if success && resp != nil {
146147
stats = GetResponseStats(resp)
148+
metricNames = ExtractMetricNames(resp)
147149
}
148150

149151
// Create the range query log entry.
@@ -164,7 +166,8 @@ func (m *rangeQueryLoggingMiddleware) logRangeQuery(req *ThanosQueryRangeRequest
164166
Tenant: userInfo.Tenant,
165167
ForwardedFor: userInfo.ForwardedFor,
166168
UserAgent: userInfo.UserAgent,
167-
EmailId: email,
169+
EmailId: userInfo.Email,
170+
Groups: userInfo.Groups,
168171
// Query-related fields
169172
StartTimestampMs: req.Start,
170173
EndTimestampMs: req.End,
@@ -181,6 +184,8 @@ func (m *rangeQueryLoggingMiddleware) logRangeQuery(req *ThanosQueryRangeRequest
181184
Engine: req.Engine,
182185
SplitIntervalMs: req.SplitInterval.Milliseconds(),
183186
Stats: req.Stats,
187+
MetricNames: metricNames,
188+
Shard: os.Getenv("PANTHEON_SHARDNAME"),
184189
// Store-matcher details
185190
StoreMatchers: ConvertStoreMatchers(req.StoreMatchers),
186191
}

0 commit comments

Comments
 (0)