Caching: Consolidate resource cache checking and updating in plugin middleware (#67002)

* Update the HandleResourceRequest function to mimic the HandleQueryRequest function

* Remove CacheResourceResponse function from interface

* revert additional thing I missed

pkg/services/pluginsintegration/clientmiddleware/caching_middleware.go

@@ -100,7 +100,7 @@ func (m *CachingMiddleware) CallResource(ctx context.Context, req *backend.CallR
 	start := time.Now()
 
 	// First look in the resource cache if enabled
-	hit, resp := m.caching.HandleResourceRequest(ctx, req)
+	hit, cr := m.caching.HandleResourceRequest(ctx, req)
 
 	defer func() {
 		// record request duration if caching was used
@@ -114,13 +114,21 @@ func (m *CachingMiddleware) CallResource(ctx context.Context, req *backend.CallR
 
 	// Cache hit; send the response and return
 	if hit {
-		return sender.Send(resp)
+		return sender.Send(cr.Response)
 	}
 
 	// Cache miss; do the actual request
-	// The call to update the cache happens in /pkg/api/plugin_resource.go in the flushStream() func
-	// TODO: Implement updating the cache from this method
-	return m.next.CallResource(ctx, req, sender)
+	// If there is no update cache func, just pass in the original sender
+	if cr.UpdateCacheFn == nil {
+		return m.next.CallResource(ctx, req, sender)
+	}
+	// Otherwise, intercept the responses in a wrapped sender so we can cache them first
+	cacheSender := cachedSenderFunc(func(res *backend.CallResourceResponse) error {
+		cr.UpdateCacheFn(ctx, res)
+		return sender.Send(res)
+	})
+
+	return m.next.CallResource(ctx, req, cacheSender)
 }
 
 func (m *CachingMiddleware) CheckHealth(ctx context.Context, req *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {
@@ -142,3 +150,9 @@ func (m *CachingMiddleware) PublishStream(ctx context.Context, req *backend.Publ
 func (m *CachingMiddleware) RunStream(ctx context.Context, req *backend.RunStreamRequest, sender *backend.StreamSender) error {
 	return m.next.RunStream(ctx, req, sender)
 }
+
+type cachedSenderFunc func(res *backend.CallResourceResponse) error
+
+func (fn cachedSenderFunc) Send(res *backend.CallResourceResponse) error {
+	return fn(res)
+}

pkg/services/pluginsintegration/clientmiddleware/caching_middleware_test.go

@@ -97,10 +97,30 @@ func TestCachingMiddleware(t *testing.T) {
 		req, err := http.NewRequest(http.MethodGet, "/resource/blah", nil)
 		require.NoError(t, err)
 
+		// This is the response returned by the HandleResourceRequest call
+		// Track whether the update cache fn was called, depending on what the response headers are in the cache request
+		var updateCacheCalled bool
+		dataResponse := caching.CachedResourceDataResponse{
+			Response: &backend.CallResourceResponse{
+				Status: 200,
+				Body:   []byte("bogus"),
+			},
+			UpdateCacheFn: func(ctx context.Context, rdr *backend.CallResourceResponse) {
+				updateCacheCalled = true
+			},
+		}
+
+		// This is the response sent via the passed-in sender when there is a cache miss
+		simulatedPluginResponse := &backend.CallResourceResponse{
+			Status: 201,
+			Body:   []byte("bogus"),
+		}
+
 		cs := caching.NewFakeOSSCachingService()
 		cdt := clienttest.NewClientDecoratorTest(t,
 			clienttest.WithReqContext(req, &user.SignedInUser{}),
 			clienttest.WithMiddlewares(NewCachingMiddleware(cs)),
+			clienttest.WithResourceResponses([]*backend.CallResourceResponse{simulatedPluginResponse}),
 		)
 
 		jsonDataMap := map[string]interface{}{}
@@ -121,11 +141,6 @@ func TestCachingMiddleware(t *testing.T) {
 			PluginContext: pluginCtx,
 		}
 
-		resourceResponse := &backend.CallResourceResponse{
-			Status: 200,
-			Body:   []byte("bogus"),
-		}
-
 		var sentResponse *backend.CallResourceResponse
 		var storeOneResponseCallResourceSender = callResourceResponseSenderFunc(func(res *backend.CallResourceResponse) error {
 			sentResponse = res
@@ -139,32 +154,37 @@ func TestCachingMiddleware(t *testing.T) {
 			})
 
 			cs.ReturnHit = true
-			cs.ReturnResourceResponse = resourceResponse
+			cs.ReturnResourceResponse = dataResponse
 
 			err := cdt.Decorator.CallResource(req.Context(), crr, storeOneResponseCallResourceSender)
 			assert.NoError(t, err)
 			// Cache service is called once
 			cs.AssertCalls(t, "HandleResourceRequest", 1)
-			// Equals the mocked response was sent
+			// The mocked cached response was sent
 			assert.NotNil(t, sentResponse)
-			assert.Equal(t, resourceResponse, sentResponse)
+			assert.Equal(t, dataResponse.Response, sentResponse)
+			// Cache was not updated by the middleware
+			assert.False(t, updateCacheCalled)
 		})
 
-		t.Run("If cache returns a miss, resource call is issued", func(t *testing.T) {
+		t.Run("If cache returns a miss, resource call is issued and the update cache function is called", func(t *testing.T) {
 			t.Cleanup(func() {
 				sentResponse = nil
 				cs.Reset()
 			})
 
 			cs.ReturnHit = false
-			cs.ReturnResourceResponse = resourceResponse
+			cs.ReturnResourceResponse = dataResponse
 
 			err := cdt.Decorator.CallResource(req.Context(), crr, storeOneResponseCallResourceSender)
 			assert.NoError(t, err)
 			// Cache service is called once
 			cs.AssertCalls(t, "HandleResourceRequest", 1)
-			// Nil response was sent
-			assert.Nil(t, sentResponse)
+			// Simulated plugin response was sent
+			assert.NotNil(t, sentResponse)
+			assert.Equal(t, simulatedPluginResponse, sentResponse)
+			// Since it was a miss, the middleware called the update func
+			assert.True(t, updateCacheCalled)
 		})
 	})