QueryCaching: Use CachingServiceClient for query caching (#112128)

* Integrate mt querier with query caching

* typo

* let the caller set cache status response header

* fix TestQueryAPI

* make gen-go

* handle CachingServiceClient being nil and make gen-go

* include namespace in cache key

* set signed in user namespace in query_test.go

* fix test

* remove commented out code

* undo services/query/query.go changes

* make gen-go

* remove namespace requirement

* fix tests

* fix test

* remove namespace from SignedInUser in tests

* make gen-go
This commit is contained in:
Bruno
2025-10-28 11:41:46 -03:00
committed by GitHub
parent 3131a69f04
commit 437dcc875c
12 changed files with 345 additions and 305 deletions
+7 -5
View File
@@ -10,19 +10,20 @@ import (
type FakeOSSCachingService struct {
calls map[string]int
ReturnStatus CacheStatus
ReturnHit bool
ReturnResourceResponse CachedResourceDataResponse
ReturnQueryResponse CachedQueryDataResponse
}
func (f *FakeOSSCachingService) HandleQueryRequest(ctx context.Context, req *backend.QueryDataRequest) (bool, CachedQueryDataResponse) {
func (f *FakeOSSCachingService) HandleQueryRequest(ctx context.Context, req *backend.QueryDataRequest) (bool, CachedQueryDataResponse, CacheStatus) {
f.calls["HandleQueryRequest"]++
return f.ReturnHit, f.ReturnQueryResponse
return f.ReturnHit, f.ReturnQueryResponse, f.ReturnStatus
}
func (f *FakeOSSCachingService) HandleResourceRequest(ctx context.Context, req *backend.CallResourceRequest) (bool, CachedResourceDataResponse) {
func (f *FakeOSSCachingService) HandleResourceRequest(ctx context.Context, req *backend.CallResourceRequest) (bool, CachedResourceDataResponse, CacheStatus) {
f.calls["HandleResourceRequest"]++
return f.ReturnHit, f.ReturnResourceResponse
return f.ReturnHit, f.ReturnResourceResponse, f.ReturnStatus
}
func (f *FakeOSSCachingService) AssertCalls(t *testing.T, fn string, times int) {
@@ -35,7 +36,8 @@ func (f *FakeOSSCachingService) Reset() {
func NewFakeOSSCachingService() *FakeOSSCachingService {
fake := &FakeOSSCachingService{
calls: map[string]int{},
calls: map[string]int{},
ReturnStatus: "unset",
}
return fake
+43
View File
@@ -0,0 +1,43 @@
package caching
import (
"github.com/grafana/grafana/pkg/infra/metrics"
contextmodel "github.com/grafana/grafana/pkg/services/contexthandler/model"
"github.com/prometheus/client_golang/prometheus"
)
const (
QueryPubdash = "pubdash"
QueryDashboard = "dashboard"
)
var QueryCachingRequestHistogram = prometheus.NewHistogramVec(prometheus.HistogramOpts{
Namespace: metrics.ExporterName,
Subsystem: "caching",
Name: "query_caching_request_duration_seconds",
Help: "histogram of grafana query endpoint requests in seconds",
Buckets: []float64{.005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5, 10, 25, 50, 100},
}, []string{"datasource_type", "cache", "query_type"})
var ShouldCacheQueryHistogram = prometheus.NewHistogramVec(prometheus.HistogramOpts{
Namespace: metrics.ExporterName,
Subsystem: "caching",
Name: "should_cache_query_request_duration_seconds",
Help: "histogram of grafana query endpoint requests in seconds",
Buckets: []float64{.005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5, 10, 25, 50, 100},
}, []string{"datasource_type", "cache", "shouldCache", "query_type"})
var ResourceCachingRequestHistogram = prometheus.NewHistogramVec(prometheus.HistogramOpts{
Namespace: metrics.ExporterName,
Subsystem: "caching",
Name: "resource_caching_request_duration_seconds",
Help: "histogram of grafana resource endpoint requests in seconds",
Buckets: []float64{.005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5, 10, 25, 50, 100},
}, []string{"plugin_id", "cache"})
func getQueryType(req *contextmodel.ReqContext) string {
if req.IsPublicDashboardView() {
return QueryPubdash
}
return QueryDashboard
}
+157 -12
View File
@@ -7,20 +7,34 @@ import (
"encoding/hex"
"encoding/json"
"io"
"strconv"
"strings"
"time"
"github.com/grafana/grafana-aws-sdk/pkg/awsds"
"github.com/grafana/grafana/pkg/infra/log"
"github.com/grafana/grafana-plugin-sdk-go/backend"
"github.com/grafana/grafana/pkg/services/contexthandler"
"github.com/grafana/grafana/pkg/services/featuremgmt"
"github.com/prometheus/client_golang/prometheus"
)
type CacheStatus string
const (
XCacheHeader = "X-Cache"
StatusHit = "HIT"
StatusMiss = "MISS"
StatusBypass = "BYPASS"
StatusError = "ERROR"
StatusDisabled = "DISABLED"
XCacheHeader = "X-Cache"
StatusHit CacheStatus = "HIT"
StatusMiss CacheStatus = "MISS"
StatusBypass CacheStatus = "BYPASS"
StatusError CacheStatus = "ERROR"
StatusDisabled CacheStatus = "DISABLED"
)
// needed to mock the function for testing
var ShouldCacheQuery = awsds.ShouldCacheQuery
type CacheQueryResponseFn func(context.Context, *backend.QueryDataResponse)
type CacheResourceResponseFn func(context.Context, *backend.CallResourceResponse)
@@ -49,22 +63,22 @@ type CachingService interface {
// HandleQueryRequest uses a QueryDataRequest to check the cache for any existing results for that query.
// If none are found, it should return false and a CachedQueryDataResponse with an UpdateCacheFn which can be used to update the results cache after the fact.
// This function may populate any response headers (accessible through the context) with the cache status using the X-Cache header.
HandleQueryRequest(context.Context, *backend.QueryDataRequest) (bool, CachedQueryDataResponse)
HandleQueryRequest(ctx context.Context, req *backend.QueryDataRequest) (bool, CachedQueryDataResponse, CacheStatus)
// HandleResourceRequest uses a CallResourceRequest to check the cache for any existing results for that request. If none are found, it should return false.
// This function may populate any response headers (accessible through the context) with the cache status using the X-Cache header.
HandleResourceRequest(context.Context, *backend.CallResourceRequest) (bool, CachedResourceDataResponse)
HandleResourceRequest(ctx context.Context, req *backend.CallResourceRequest) (bool, CachedResourceDataResponse, CacheStatus)
}
// Implementation of interface - does nothing
type OSSCachingService struct {
}
func (s *OSSCachingService) HandleQueryRequest(ctx context.Context, req *backend.QueryDataRequest) (bool, CachedQueryDataResponse) {
return false, CachedQueryDataResponse{}
func (s *OSSCachingService) HandleQueryRequest(ctx context.Context, req *backend.QueryDataRequest) (bool, CachedQueryDataResponse, CacheStatus) {
return false, CachedQueryDataResponse{}, ""
}
func (s *OSSCachingService) HandleResourceRequest(ctx context.Context, req *backend.CallResourceRequest) (bool, CachedResourceDataResponse) {
return false, CachedResourceDataResponse{}
func (s *OSSCachingService) HandleResourceRequest(ctx context.Context, req *backend.CallResourceRequest) (bool, CachedResourceDataResponse, CacheStatus) {
return false, CachedResourceDataResponse{}, ""
}
var _ CachingService = &OSSCachingService{}
@@ -133,3 +147,134 @@ func (e *JSONEncoder) Encode(w io.Writer, v interface{}) error {
func (e *JSONEncoder) Decode(r io.Reader, v interface{}) error {
return json.NewDecoder(r).Decode(v)
}
// A service that provides methods to cache requests.
// It can be used to cache requests using `caching.CachingService` without reimplementing
// the caching logic at every call site.
type CachingServiceClient struct {
cachingService CachingService
features featuremgmt.FeatureToggles
}
func ProvideCachingServiceClient(cachingService CachingService, features featuremgmt.FeatureToggles) *CachingServiceClient {
log := log.New("caching_service_client")
if err := prometheus.Register(QueryCachingRequestHistogram); err != nil {
log.Error("Error registering prometheus collector 'QueryRequestHistogram'", "error", err)
}
if err := prometheus.Register(ResourceCachingRequestHistogram); err != nil {
log.Error("Error registering prometheus collector 'ResourceRequestHistogram'", "error", err)
}
return &CachingServiceClient{cachingService: cachingService, features: features}
}
// WithQueryDataCaching calls `f` and caches the returned value if `req` has not been cached already.
// Returns the cached value otherwise.
func (c *CachingServiceClient) WithQueryDataCaching(ctx context.Context, req *backend.QueryDataRequest, f func() (*backend.QueryDataResponse, error)) (*backend.QueryDataResponse, error) {
if c == nil || req == nil {
return f()
}
reqCtx := contexthandler.FromContext(ctx)
// time how long this request takes
start := time.Now()
// First look in the query cache if enabled
hit, cr, status := c.cachingService.HandleQueryRequest(ctx, req)
// record request duration if caching was used
if reqCtx != nil {
reqCtx.Resp.Header().Set(XCacheHeader, string(status))
defer func() {
QueryCachingRequestHistogram.With(prometheus.Labels{
"datasource_type": getDatasourceType(req.PluginContext),
"cache": string(status),
"query_type": getQueryType(reqCtx),
}).Observe(time.Since(start).Seconds())
}()
}
// Cache hit; return the response
if hit {
return cr.Response, nil
}
// Cache miss; do the actual queries
resp, err := f()
// Update the query cache with the result for this metrics request
if err == nil && cr.UpdateCacheFn != nil {
// If AWS async caching is not enabled, use the old code path
if c.features == nil || !c.features.IsEnabled(ctx, featuremgmt.FlagAwsAsyncQueryCaching) {
cr.UpdateCacheFn(ctx, resp)
} else if reqCtx != nil {
// time how long shouldCacheQuery takes
startShouldCacheQuery := time.Now()
shouldCache := ShouldCacheQuery(resp)
ShouldCacheQueryHistogram.With(prometheus.Labels{
"datasource_type": req.PluginContext.DataSourceInstanceSettings.Type,
"cache": string(status),
"shouldCache": strconv.FormatBool(shouldCache),
"query_type": getQueryType(reqCtx),
}).Observe(time.Since(startShouldCacheQuery).Seconds())
// If AWS async caching is enabled and resp is for a running async query, don't cache it
if shouldCache {
cr.UpdateCacheFn(ctx, resp)
}
}
}
return resp, err
}
// WithCallResourceCaching calls `f` and caches the returned value if `req` has not been cached already.
// Returns the cached value otherwise.
func (c *CachingServiceClient) WithCallResourceCaching(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender, f func(backend.CallResourceResponseSender) error) error {
if c == nil || req == nil {
return f(sender)
}
reqCtx := contexthandler.FromContext(ctx)
// time how long this request takes
start := time.Now()
// First look in the resource cache if enabled
hit, cr, status := c.cachingService.HandleResourceRequest(ctx, req)
if reqCtx != nil {
reqCtx.Resp.Header().Set(XCacheHeader, string(status))
}
// record request duration if caching was used
defer func() {
ResourceCachingRequestHistogram.With(prometheus.Labels{
"plugin_id": req.PluginContext.PluginID,
"cache": string(status),
}).Observe(time.Since(start).Seconds())
}()
// Cache hit; send the response and return
if hit {
return sender.Send(cr.Response)
}
// Cache miss; do the actual request
// If there is no update cache func, just pass in the original sender
if cr.UpdateCacheFn == nil {
return f(sender)
}
// Otherwise, intercept the responses in a wrapped sender so we can cache them first
cacheSender := backend.CallResourceResponseSenderFunc(func(res *backend.CallResourceResponse) error {
cr.UpdateCacheFn(ctx, res)
return sender.Send(res)
})
return f(cacheSender)
}
func getDatasourceType(pluginCtx backend.PluginContext) string {
if pluginCtx.DataSourceInstanceSettings == nil {
return "unknown"
}
return pluginCtx.DataSourceInstanceSettings.Name
}
+130
View File
@@ -0,0 +1,130 @@
package caching
import (
"context"
"errors"
"net/http/httptest"
"testing"
"github.com/grafana/grafana-plugin-sdk-go/backend"
"github.com/grafana/grafana/pkg/services/contexthandler/ctxkey"
contextmodel "github.com/grafana/grafana/pkg/services/contexthandler/model"
"github.com/grafana/grafana/pkg/web"
"github.com/stretchr/testify/require"
)
func TestWithQueryDataCaching(t *testing.T) {
t.Run("caching is a no-op when service is nil", func(t *testing.T) {
var s *CachingServiceClient
req := backend.QueryDataRequest{}
fakeResponse := &backend.QueryDataResponse{}
response, err := s.WithQueryDataCaching(t.Context(), &req, func() (*backend.QueryDataResponse, error) {
return fakeResponse, nil
})
require.NoError(t, err)
require.Equal(t, fakeResponse, response)
})
t.Run("cache status is included in the response if a request context is available", func(t *testing.T) {
fakeCachingService := NewFakeOSSCachingService()
fakeCachingService.ReturnStatus = StatusMiss
client := ProvideCachingServiceClient(fakeCachingService, nil)
req := backend.QueryDataRequest{}
reqCtx := &contextmodel.ReqContext{
Context: &web.Context{
Resp: web.NewResponseWriter("", httptest.NewRecorder()),
},
}
ctx := context.WithValue(t.Context(), ctxkey.Key{}, reqCtx)
fakeResponse := &backend.QueryDataResponse{}
response, err := client.WithQueryDataCaching(ctx, &req, func() (*backend.QueryDataResponse, error) {
return fakeResponse, nil
})
require.NoError(t, err)
require.Equal(t, fakeResponse, response)
require.EqualValues(t, StatusMiss, reqCtx.Resp.Header().Get(XCacheHeader))
})
t.Run("caching can be used without a request context", func(t *testing.T) {
fakeCachingService := NewFakeOSSCachingService()
fakeCachingService.ReturnStatus = StatusMiss
client := ProvideCachingServiceClient(fakeCachingService, nil)
req := backend.QueryDataRequest{}
fakeResponse := &backend.QueryDataResponse{}
// Using the default test context, no request context.
response, err := client.WithQueryDataCaching(t.Context(), &req, func() (*backend.QueryDataResponse, error) {
return fakeResponse, nil
})
require.NoError(t, err)
require.Equal(t, fakeResponse, response)
})
}
func TestWithCallResourceCaching(t *testing.T) {
t.Run("caching is a no-op when service is nil", func(t *testing.T) {
var s *CachingServiceClient
req := backend.CallResourceRequest{}
fakeErr := errors.New("oops")
err := s.WithCallResourceCaching(t.Context(), &req, nil, func(backend.CallResourceResponseSender) error {
return fakeErr
})
require.ErrorIs(t, err, fakeErr)
})
t.Run("cache status is included in the response if a request context is available", func(t *testing.T) {
fakeCachingService := NewFakeOSSCachingService()
fakeCachingService.ReturnStatus = StatusMiss
client := ProvideCachingServiceClient(fakeCachingService, nil)
req := backend.CallResourceRequest{}
reqCtx := &contextmodel.ReqContext{
Context: &web.Context{
Resp: web.NewResponseWriter("", httptest.NewRecorder()),
},
}
ctx := context.WithValue(t.Context(), ctxkey.Key{}, reqCtx)
sender := func(*backend.CallResourceResponse) error {
return nil
}
var fakeErr = errors.New("oops")
err := client.WithCallResourceCaching(ctx, &req, backend.CallResourceResponseSenderFunc(sender), func(backend.CallResourceResponseSender) error {
return fakeErr
})
require.ErrorIs(t, err, fakeErr)
require.EqualValues(t, StatusMiss, reqCtx.Resp.Header().Get(XCacheHeader))
})
t.Run("caching can be used without a request context", func(t *testing.T) {
fakeCachingService := NewFakeOSSCachingService()
fakeCachingService.ReturnStatus = StatusMiss
client := ProvideCachingServiceClient(fakeCachingService, nil)
req := backend.CallResourceRequest{}
sender := func(*backend.CallResourceResponse) error {
return nil
}
var fakeErr = errors.New("oops")
// Using the default test context, no request context.
err := client.WithCallResourceCaching(t.Context(), &req, backend.CallResourceResponseSenderFunc(sender), func(_ backend.CallResourceResponseSender) error {
return fakeErr
})
require.ErrorIs(t, err, fakeErr)
})
}
func TestGetDatasourceType(t *testing.T) {
t.Parallel()
require.Equal(t, "unknown", getDatasourceType(backend.PluginContext{}))
require.Equal(t, "name", getDatasourceType(backend.PluginContext{
DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{
Name: "name",
},
}))
}