feat: capture logs drilldown flag in query samples (#19242)

Author: trevorwhitney

pkg/goldfish/migrations/20250918000001_add_logs_drilldown_column.sql

@@ -0,0 +1,17 @@
+-- +goose Up
+-- Add is_logs_drilldown column to track queries from Logs Drilldown app
+
+ALTER TABLE sampled_queries 
+ADD COLUMN is_logs_drilldown BOOLEAN DEFAULT FALSE AFTER user;
+
+-- Update existing rows to have FALSE instead of NULL
+UPDATE sampled_queries SET is_logs_drilldown = FALSE WHERE is_logs_drilldown IS NULL;
+
+-- Add index for efficient filtering by logs drilldown status
+CREATE INDEX idx_sampled_queries_logs_drilldown ON sampled_queries(is_logs_drilldown);
+
+-- +goose Down
+-- Remove is_logs_drilldown column and index
+
+DROP INDEX idx_sampled_queries_logs_drilldown ON sampled_queries;
+ALTER TABLE sampled_queries DROP COLUMN is_logs_drilldown;

pkg/goldfish/storage_mysql.go

@@ -78,7 +78,7 @@ func NewMySQLStorage(config StorageConfig, logger log.Logger) (*MySQLStorage, er
 func (s *MySQLStorage) StoreQuerySample(ctx context.Context, sample *QuerySample) error {
 	query := `
 		INSERT INTO sampled_queries (
-			correlation_id, tenant_id, user, query, query_type,
+			correlation_id, tenant_id, user, is_logs_drilldown, query, query_type,
 			start_time, end_time, step_duration,
 			cell_a_exec_time_ms, cell_b_exec_time_ms,
 			cell_a_queue_time_ms, cell_b_queue_time_ms,
@@ -96,7 +96,7 @@ func (s *MySQLStorage) StoreQuerySample(ctx context.Context, sample *QuerySample
 			cell_a_span_id, cell_b_span_id,
 			cell_a_used_new_engine, cell_b_used_new_engine,
 			sampled_at
-		) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+		) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
 	`
 
 	// Convert empty span IDs to NULL for database storage
@@ -112,6 +112,7 @@ func (s *MySQLStorage) StoreQuerySample(ctx context.Context, sample *QuerySample
 		sample.CorrelationID,
 		sample.TenantID,
 		sample.User,
+		sample.IsLogsDrilldown,
 		sample.Query,
 		sample.QueryType,
 		sample.StartTime,
@@ -329,6 +330,12 @@ func buildWhereClause(filter QueryFilter) (string, []any) {
 		args = append(args, filter.User)
 	}
 
+	// Add logs drilldown filter
+	if filter.IsLogsDrilldown != nil {
+		conditions = append(conditions, "is_logs_drilldown = ?")
+		args = append(args, *filter.IsLogsDrilldown)
+	}
+
 	// Add new engine filter
 	if filter.UsedNewEngine != nil {
 		if *filter.UsedNewEngine {

pkg/goldfish/storage_mysql_test.go

@@ -61,15 +61,32 @@ func TestQueryFilter_BuildWhereClause(t *testing.T) {
 			expectedWhere: "WHERE cell_a_used_new_engine = 0 AND cell_b_used_new_engine = 0",
 			expectedArgs:  nil,
 		},
+		{
+			name: "logs drilldown filter true",
+			filter: QueryFilter{
+				IsLogsDrilldown: boolPtr(true),
+			},
+			expectedWhere: "WHERE is_logs_drilldown = ?",
+			expectedArgs:  []any{true},
+		},
+		{
+			name: "logs drilldown filter false",
+			filter: QueryFilter{
+				IsLogsDrilldown: boolPtr(false),
+			},
+			expectedWhere: "WHERE is_logs_drilldown = ?",
+			expectedArgs:  []any{false},
+		},
 		{
 			name: "all filters combined",
 			filter: QueryFilter{
-				Tenant:        "tenant-b",
-				User:          "user123",
-				UsedNewEngine: boolPtr(true),
+				Tenant:          "tenant-b",
+				User:            "user123",
+				IsLogsDrilldown: boolPtr(true),
+				UsedNewEngine:   boolPtr(true),
 			},
-			expectedWhere: "WHERE tenant_id = ? AND user = ? AND (cell_a_used_new_engine = 1 OR cell_b_used_new_engine = 1)",
-			expectedArgs:  []any{"tenant-b", "user123"},
+			expectedWhere: "WHERE tenant_id = ? AND user = ? AND is_logs_drilldown = ? AND (cell_a_used_new_engine = 1 OR cell_b_used_new_engine = 1)",
+			expectedArgs:  []any{"tenant-b", "user123", true},
 		},
 		{
 			name: "with time range specified",

pkg/goldfish/types.go

@@ -6,14 +6,15 @@ import (
 
 // QuerySample represents a sampled query with performance stats from both cells
 type QuerySample struct {
-	CorrelationID string        `json:"correlationId"`
-	TenantID      string        `json:"tenantId"`
-	User          string        `json:"user"`
-	Query         string        `json:"query"`
-	QueryType     string        `json:"queryType"`
-	StartTime     time.Time     `json:"startTime"`
-	EndTime       time.Time     `json:"endTime"`
-	Step          time.Duration `json:"step"`
+	CorrelationID   string        `json:"correlationId"`
+	TenantID        string        `json:"tenantId"`
+	User            string        `json:"user"`
+	IsLogsDrilldown bool          `json:"isLogsDrilldown"`
+	Query           string        `json:"query"`
+	QueryType       string        `json:"queryType"`
+	StartTime       time.Time     `json:"startTime"`
+	EndTime         time.Time     `json:"endTime"`
+	Step            time.Duration `json:"step"`
 
 	// Performance statistics instead of raw responses
 	CellAStats QueryStats `json:"cellAStats"`
@@ -82,8 +83,9 @@ type PerformanceMetrics struct {
 
 // QueryFilter contains filters for querying sampled queries
 type QueryFilter struct {
-	Tenant        string
-	User          string
-	UsedNewEngine *bool // pointer to handle true/false/nil states
-	From, To      time.Time
+	Tenant          string
+	User            string
+	IsLogsDrilldown *bool // pointer to handle true/false/nil states
+	UsedNewEngine   *bool // pointer to handle true/false/nil states
+	From, To        time.Time
 }

tools/querytee/goldfish/manager.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"net/http"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/go-kit/log"
@@ -16,6 +17,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 
+	"github.com/grafana/loki/v3/pkg/util/constants"
 	"github.com/grafana/loki/v3/pkg/util/httpreq"
 )
 
@@ -124,6 +126,7 @@ func (m *Manager) ProcessQueryPair(ctx context.Context, req *http.Request, cellA
 		CorrelationID:      correlationID,
 		TenantID:           extractTenant(req),
 		User:               extractUserFromQueryTags(req, m.logger),
+		IsLogsDrilldown:    isLogsDrilldownRequest(req),
 		Query:              req.URL.Query().Get("query"),
 		QueryType:          getQueryType(req.URL.Path),
 		StartTime:          startTime,
@@ -410,3 +413,21 @@ func extractUserFromQueryTags(req *http.Request, logger log.Logger) string {
 	level.Debug(logger).Log("goldfish", "user-extraction", "result", unknownUser)
 	return unknownUser
 }
+
+// isLogsDrilldownRequest checks if the request comes from Logs Drilldown by examining the X-Query-Tags header
+func isLogsDrilldownRequest(req *http.Request) bool {
+	tags := httpreq.ExtractQueryTagsFromHTTP(req)
+	kvs := httpreq.TagsToKeyValues(tags)
+
+	// Iterate through key-value pairs (keys at even indices, values at odd)
+	for i := 0; i < len(kvs); i += 2 {
+		if i+1 < len(kvs) {
+			key, keyOK := kvs[i].(string)
+			value, valueOK := kvs[i+1].(string)
+			if keyOK && valueOK && key == "source" && strings.EqualFold(value, constants.LogsDrilldownAppName) {
+				return true
+			}
+		}
+	}
+	return false
+}

tools/querytee/goldfish/manager_test.go

@@ -401,3 +401,89 @@ func TestExtractUserFromQueryTags(t *testing.T) {
 		})
 	}
 }
+
+func TestProcessQueryPair_CapturesLogsDrilldown(t *testing.T) {
+	tests := []struct {
+		name              string
+		queryTags         string
+		expectedDrilldown bool
+	}{
+		{
+			name:              "query with logs drilldown source first",
+			queryTags:         "source=grafana-lokiexplore-app,user=john.doe",
+			expectedDrilldown: true,
+		},
+		{
+			name:              "query with logs drilldown source last",
+			queryTags:         "user=john.doe,source=grafana-lokiexplore-app",
+			expectedDrilldown: true,
+		},
+		{
+			name:              "query with different source",
+			queryTags:         "source=grafana,user=john.doe",
+			expectedDrilldown: false,
+		},
+		{
+			name:              "query without source",
+			queryTags:         "user=john.doe",
+			expectedDrilldown: false,
+		},
+		{
+			name:              "query without query tags",
+			queryTags:         "",
+			expectedDrilldown: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			config := Config{
+				Enabled: true,
+				SamplingConfig: SamplingConfig{
+					DefaultRate: 1.0,
+				},
+			}
+
+			storage := &mockStorage{}
+			manager, err := NewManager(config, storage, log.NewNopLogger(), prometheus.NewRegistry())
+			require.NoError(t, err)
+
+			req, _ := http.NewRequest("GET", "/loki/api/v1/query_range?query=count_over_time({job=\"test\"}[5m])&start=1700000000&end=1700001000&step=60s", nil)
+			req.Header.Set("X-Scope-OrgID", "tenant1")
+			if tt.queryTags != "" {
+				req.Header.Set("X-Query-Tags", tt.queryTags)
+			}
+
+			cellAResp := &ResponseData{
+				Body:          []byte(`{"status":"success","data":{"resultType":"matrix","result":[]}}`),
+				StatusCode:    200,
+				Duration:      100 * time.Millisecond,
+				Stats:         goldfish.QueryStats{ExecTimeMs: 100},
+				Hash:          "hash123",
+				Size:          150,
+				UsedNewEngine: false,
+			}
+
+			cellBResp := &ResponseData{
+				Body:          []byte(`{"status":"success","data":{"resultType":"matrix","result":[]}}`),
+				StatusCode:    200,
+				Duration:      120 * time.Millisecond,
+				Stats:         goldfish.QueryStats{ExecTimeMs: 120},
+				Hash:          "hash123",
+				Size:          155,
+				UsedNewEngine: false,
+			}
+
+			ctx := context.Background()
+			manager.ProcessQueryPair(ctx, req, cellAResp, cellBResp)
+
+			// Give async processing time to complete
+			time.Sleep(100 * time.Millisecond)
+
+			// Verify sample was stored with correct logs drilldown flag
+			require.Len(t, storage.samples, 1)
+			sample := storage.samples[0]
+			assert.Equal(t, tt.expectedDrilldown, sample.IsLogsDrilldown, "IsLogsDrilldown field should be captured from X-Query-Tags header")
+		})
+	}
+}