Convert adapter utils to middlewares

backend/admission_adapter.go

@@ -18,15 +18,8 @@ func newAdmissionSDKAdapter(handler AdmissionHandler) *admissionSDKAdapter {
 }
 
 func (a *admissionSDKAdapter) ValidateAdmission(ctx context.Context, req *pluginv2.AdmissionRequest) (*pluginv2.ValidationResponse, error) {
-	ctx = setupContext(ctx, EndpointValidateAdmission)
 	parsedReq := FromProto().AdmissionRequest(req)
-
-	var resp *ValidationResponse
-	err := wrapHandler(ctx, parsedReq.PluginContext, func(ctx context.Context) (RequestStatus, error) {
-		var innerErr error
-		resp, innerErr = a.handler.ValidateAdmission(ctx, parsedReq)
-		return RequestStatusFromError(innerErr), innerErr
-	})
+	resp, err := a.handler.ValidateAdmission(ctx, parsedReq)
 	if err != nil {
 		return nil, err
 	}
@@ -35,15 +28,8 @@ func (a *admissionSDKAdapter) ValidateAdmission(ctx context.Context, req *plugin
 }
 
 func (a *admissionSDKAdapter) MutateAdmission(ctx context.Context, req *pluginv2.AdmissionRequest) (*pluginv2.MutationResponse, error) {
-	ctx = setupContext(ctx, EndpointMutateAdmission)
 	parsedReq := FromProto().AdmissionRequest(req)
-
-	var resp *MutationResponse
-	err := wrapHandler(ctx, parsedReq.PluginContext, func(ctx context.Context) (RequestStatus, error) {
-		var innerErr error
-		resp, innerErr = a.handler.MutateAdmission(ctx, parsedReq)
-		return RequestStatusFromError(innerErr), innerErr
-	})
+	resp, err := a.handler.MutateAdmission(ctx, parsedReq)
 	if err != nil {
 		return nil, err
 	}

backend/common.go

@@ -297,13 +297,6 @@ func SecureJSONDataFromHTTPClientOptions(opts httpclient.Options) (res map[strin
 	return secureJSONData
 }
 
-func propagateTenantIDIfPresent(ctx context.Context) context.Context {
-	if tid, exists := tenant.IDFromIncomingGRPCContext(ctx); exists {
-		ctx = tenant.WithTenant(ctx, tid)
-	}
-	return ctx
-}
-
 func (s *DataSourceInstanceSettings) ProxyOptionsFromContext(ctx context.Context) (*proxy.Options, error) {
 	cfg := GrafanaConfigFromContext(ctx)
 	p, err := cfg.proxy()
@@ -383,3 +376,76 @@ func (s *DataSourceInstanceSettings) ProxyClient(ctx context.Context) (proxy.Cli
 func WithTenant(ctx context.Context, tenantID string) context.Context {
 	return tenant.WithTenant(ctx, tenantID)
 }
+
+type handlerWrapperFunc func(ctx context.Context) (RequestStatus, error)
+
+// newTenantIDMiddleware creates a new handler middleware that extract tenant ID from the incoming gRPC context, if available.
+func newTenantIDMiddleware() HandlerMiddleware {
+	return HandlerMiddlewareFunc(func(next Handler) Handler {
+		return &tenantIDMiddleware{
+			BaseHandler: NewBaseHandler(next),
+		}
+	})
+}
+
+// tenantIDMiddleware a handler middleware that extract tenant ID from the incoming gRPC context, if available.
+type tenantIDMiddleware struct {
+	BaseHandler
+}
+
+func (m *tenantIDMiddleware) setup(ctx context.Context) context.Context {
+	if tid, exists := tenant.IDFromIncomingGRPCContext(ctx); exists {
+		ctx = tenant.WithTenant(ctx, tid)
+	}
+	return ctx
+}
+
+func (m *tenantIDMiddleware) QueryData(ctx context.Context, req *QueryDataRequest) (*QueryDataResponse, error) {
+	ctx = m.setup(ctx)
+	return m.BaseHandler.QueryData(ctx, req)
+}
+
+func (m *tenantIDMiddleware) CallResource(ctx context.Context, req *CallResourceRequest, sender CallResourceResponseSender) error {
+	ctx = m.setup(ctx)
+	return m.BaseHandler.CallResource(ctx, req, sender)
+}
+
+func (m *tenantIDMiddleware) CheckHealth(ctx context.Context, req *CheckHealthRequest) (*CheckHealthResult, error) {
+	ctx = m.setup(ctx)
+	return m.BaseHandler.CheckHealth(ctx, req)
+}
+
+func (m *tenantIDMiddleware) CollectMetrics(ctx context.Context, req *CollectMetricsRequest) (*CollectMetricsResult, error) {
+	ctx = m.setup(ctx)
+	return m.BaseHandler.CollectMetrics(ctx, req)
+}
+
+func (m *tenantIDMiddleware) SubscribeStream(ctx context.Context, req *SubscribeStreamRequest) (*SubscribeStreamResponse, error) {
+	ctx = m.setup(ctx)
+	return m.BaseHandler.SubscribeStream(ctx, req)
+}
+
+func (m *tenantIDMiddleware) PublishStream(ctx context.Context, req *PublishStreamRequest) (*PublishStreamResponse, error) {
+	ctx = m.setup(ctx)
+	return m.BaseHandler.PublishStream(ctx, req)
+}
+
+func (m *tenantIDMiddleware) RunStream(ctx context.Context, req *RunStreamRequest, sender *StreamSender) error {
+	ctx = m.setup(ctx)
+	return m.BaseHandler.RunStream(ctx, req, sender)
+}
+
+func (m *tenantIDMiddleware) ValidateAdmission(ctx context.Context, req *AdmissionRequest) (*ValidationResponse, error) {
+	ctx = m.setup(ctx)
+	return m.BaseHandler.ValidateAdmission(ctx, req)
+}
+
+func (m *tenantIDMiddleware) MutateAdmission(ctx context.Context, req *AdmissionRequest) (*MutationResponse, error) {
+	ctx = m.setup(ctx)
+	return m.BaseHandler.MutateAdmission(ctx, req)
+}
+
+func (m *tenantIDMiddleware) ConvertObjects(ctx context.Context, req *ConversionRequest) (*ConversionResponse, error) {
+	ctx = m.setup(ctx)
+	return m.BaseHandler.ConvertObjects(ctx, req)
+}

backend/conversion_adapter.go

@@ -64,24 +64,23 @@ func (a *conversionSDKAdapter) convertQueryDataRequest(ctx context.Context, requ
 }
 
 func (a *conversionSDKAdapter) ConvertObjects(ctx context.Context, req *pluginv2.ConversionRequest) (*pluginv2.ConversionResponse, error) {
-	ctx = setupContext(ctx, EndpointConvertObjects)
 	parsedReq := FromProto().ConversionRequest(req)
 
-	resp := &ConversionResponse{}
-	err := wrapHandler(ctx, parsedReq.PluginContext, func(ctx context.Context) (RequestStatus, error) {
-		var innerErr error
-		if a.queryConversionHandler != nil {
-			// Try to parse it as a query data request
-			reqs, err := parseAsQueryRequest(parsedReq)
-			if err == nil {
-				resp, innerErr = a.convertQueryDataRequest(ctx, reqs)
-				return RequestStatusFromError(innerErr), innerErr
-			}
+	var resp *ConversionResponse
+	var err error
+	if a.queryConversionHandler != nil {
+		// Try to parse it as a query data request
+		var reqs []*QueryDataRequest
+		reqs, err = parseAsQueryRequest(parsedReq)
+		if err == nil {
+			resp, err = a.convertQueryDataRequest(ctx, reqs)
+		} else {
 			// The object cannot be parsed as a query data request, so we will try to convert it as a generic object
+			resp, err = a.handler.ConvertObjects(ctx, parsedReq)
 		}
-		resp, innerErr = a.handler.ConvertObjects(ctx, parsedReq)
-		return RequestStatusFromError(innerErr), innerErr
-	})
+	} else {
+		resp, err = a.handler.ConvertObjects(ctx, parsedReq)
+	}
 	if err != nil {
 		return nil, err
 	}