QueryTypes: Add feature toggle to show query types in datasource apiservers (#88213)

* initial attempt * show query types * show query types * with formatting * with formatting * more cleanup * add feature toggle * fix build
2024-05-23 18:46:28 +02:00
parent 006d0021e3
commit 42b0f802de
17 changed files with 590 additions and 244 deletions
@@ -2,13 +2,11 @@ package datasource

 import (
 	"context"
+	"encoding/json"
 	"fmt"
-	"net/http"
 	"time"

 	"github.com/grafana/grafana-plugin-sdk-go/backend"
-	data "github.com/grafana/grafana-plugin-sdk-go/experimental/apis/data/v0alpha1"
-	"github.com/grafana/grafana-plugin-sdk-go/experimental/schemabuilder"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
@@ -19,7 +17,6 @@ import (
 	genericapiserver "k8s.io/apiserver/pkg/server"
 	openapi "k8s.io/kube-openapi/pkg/common"
 	"k8s.io/kube-openapi/pkg/spec3"
-	"k8s.io/kube-openapi/pkg/validation/spec"
 	"k8s.io/utils/strings/slices"

 	common "github.com/grafana/grafana/pkg/apimachinery/apis/common/v0alpha1"
@@ -27,15 +24,15 @@ import (
 	query "github.com/grafana/grafana/pkg/apis/query/v0alpha1"
 	"github.com/grafana/grafana/pkg/apiserver/builder"
 	"github.com/grafana/grafana/pkg/plugins"
+	"github.com/grafana/grafana/pkg/promlib/models"
+	"github.com/grafana/grafana/pkg/registry/apis/query/queryschema"
 	"github.com/grafana/grafana/pkg/services/accesscontrol"
 	"github.com/grafana/grafana/pkg/services/apiserver/utils"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/pluginsintegration/pluginstore"
+	"github.com/grafana/grafana/pkg/tsdb/grafana-testdata-datasource/kinds"
 )

-const QueryRequestSchemaKey = "QueryRequestSchema"
-const QueryPayloadSchemaKey = "QueryPayloadSchema"
-
 var _ builder.APIGroupBuilder = (*DataSourceAPIBuilder)(nil)

 // DataSourceAPIBuilder is used just so wire has something unique to return
@@ -83,6 +80,7 @@ func RegisterAPIService(
 			datasources.GetDatasourceProvider(ds.JSONData),
 			contextProvider,
 			accessControl,
+			features.IsEnabledGlobally(featuremgmt.FlagDatasourceQueryTypes),
 		)
 		if err != nil {
 			return nil, err
@@ -105,20 +103,48 @@ func NewDataSourceAPIBuilder(
 	client PluginClient,
 	datasources PluginDatasourceProvider,
 	contextProvider PluginContextWrapper,
-	accessControl accesscontrol.AccessControl) (*DataSourceAPIBuilder, error) {
+	accessControl accesscontrol.AccessControl,
+	loadQueryTypes bool,
+) (*DataSourceAPIBuilder, error) {
 	ri, err := resourceFromPluginID(plugin.ID)
 	if err != nil {
 		return nil, err
 	}

-	return &DataSourceAPIBuilder{
+	builder := &DataSourceAPIBuilder{
 		connectionResourceInfo: ri,
 		pluginJSON:             plugin,
 		client:                 client,
 		datasources:            datasources,
 		contextProvider:        contextProvider,
 		accessControl:          accessControl,
-	}, nil
+	}
+	if loadQueryTypes {
+		// In the future, this will somehow come from the plugin
+		builder.queryTypes, err = getHardcodedQueryTypes(ri.GroupResource().Group)
+	}
+	return builder, err
+}
+
+// TODO -- somehow get the list from the plugin -- not hardcoded
+func getHardcodedQueryTypes(group string) (*query.QueryTypeDefinitionList, error) {
+	var err error
+	var raw json.RawMessage
+	switch group {
+	case "testdata.datasource.grafana.app":
+		raw, err = kinds.QueryTypeDefinitionListJSON()
+	case "prometheus.datasource.grafana.app":
+		raw, err = models.QueryTypeDefinitionListJSON()
+	}
+	if err != nil {
+		return nil, err
+	}
+	if raw != nil {
+		types := &query.QueryTypeDefinitionList{}
+		err = json.Unmarshal(raw, types)
+		return types, err
+	}
+	return nil, err
 }

 func (b *DataSourceAPIBuilder) GetGroupVersion() schema.GroupVersion {
@@ -134,6 +160,8 @@ func addKnownTypes(scheme *runtime.Scheme, gv schema.GroupVersion) {
 		// Query handler
 		&query.QueryDataRequest{},
 		&query.QueryDataResponse{},
+		&query.QueryTypeDefinition{},
+		&query.QueryTypeDefinitionList{},
 		&metav1.Status{},
 	)
 }
@@ -211,12 +239,16 @@ func (b *DataSourceAPIBuilder) GetAPIGroupInfo(
 		storage[conn.StoragePath("proxy")] = &subProxyREST{pluginJSON: b.pluginJSON}
 	}

+	// Register hardcoded query schemas
+	err := queryschema.RegisterQueryTypes(b.queryTypes, storage)
+
+	// Create the group info
 	apiGroupInfo := genericapiserver.NewDefaultAPIGroupInfo(
 		conn.GroupResource().Group, scheme,
 		metav1.ParameterCodec, codecs)

 	apiGroupInfo.VersionedResourcesStorageMap[conn.GroupVersion().Version] = storage
-	return &apiGroupInfo, nil
+	return &apiGroupInfo, err
 }

 func (b *DataSourceAPIBuilder) getPluginContext(ctx context.Context, uid string) (backend.PluginContext, error) {
@@ -247,66 +279,19 @@ func (b *DataSourceAPIBuilder) PostProcessOpenAPI(oas *spec3.OpenAPI) (*spec3.Op
 	// Hide the ability to list all connections across tenants
 	delete(oas.Paths.Paths, root+b.connectionResourceInfo.GroupResource().Resource)

-	var err error
-	opts := schemabuilder.QuerySchemaOptions{
-		PluginID:   []string{b.pluginJSON.ID},
-		QueryTypes: []data.QueryTypeDefinition{},
-		Mode:       schemabuilder.SchemaTypeQueryPayload,
-	}
-	if b.pluginJSON.AliasIDs != nil {
-		opts.PluginID = append(opts.PluginID, b.pluginJSON.AliasIDs...)
-	}
-	if b.queryTypes != nil {
-		for _, qt := range b.queryTypes.Items {
-			// The SDK type and api type are not the same so we recreate it here
-			opts.QueryTypes = append(opts.QueryTypes, data.QueryTypeDefinition{
-				ObjectMeta: data.ObjectMeta{
-					Name: qt.Name,
-				},
-				Spec: qt.Spec,
-			})
-		}
-	}
-	oas.Components.Schemas[QueryPayloadSchemaKey], err = schemabuilder.GetQuerySchema(opts)
-	if err != nil {
-		return oas, err
-	}
-	opts.Mode = schemabuilder.SchemaTypeQueryRequest
-	oas.Components.Schemas[QueryRequestSchemaKey], err = schemabuilder.GetQuerySchema(opts)
-	if err != nil {
-		return oas, err
-	}
-
-	// Update the request object
-	sub := oas.Paths.Paths[root+"namespaces/{namespace}/connections/{name}/query"]
-	if sub != nil && sub.Post != nil {
-		sub.Post.Description = "Execute queries"
-		sub.Post.RequestBody = &spec3.RequestBody{
-			RequestBodyProps: spec3.RequestBodyProps{
-				Required: true,
-				Content: map[string]*spec3.MediaType{
-					"application/json": {
-						MediaTypeProps: spec3.MediaTypeProps{
-							Schema:   spec.RefSchema("#/components/schemas/" + QueryRequestSchemaKey),
-							Examples: getExamples(b.queryTypes),
-						},
-					},
-				},
-			},
-		}
-		okrsp, ok := sub.Post.Responses.StatusCodeResponses[200]
-		if ok {
-			sub.Post.Responses.StatusCodeResponses[http.StatusMultiStatus] = &spec3.Response{
-				ResponseProps: spec3.ResponseProps{
-					Description: "Query executed, but errors may exist in the datasource.  See the payload for more details.",
-					Content:     okrsp.Content,
-				},
-			}
-		}
-	}
+	// Add queries to the request properties
+	// Add queries to the request properties
+	err := queryschema.AddQueriesToOpenAPI(queryschema.OASQueryOptions{
+		Swagger:          oas,
+		PluginJSON:       &b.pluginJSON,
+		QueryTypes:       b.queryTypes,
+		Root:             root,
+		QueryPath:        "namespaces/{namespace}/connections/{name}/query",
+		QueryDescription: fmt.Sprintf("Query the %s datasources", b.pluginJSON.Name),
+	})

 	// The root API discovery list
-	sub = oas.Paths.Paths[root]
+	sub := oas.Paths.Paths[root]
 	if sub != nil && sub.Get != nil {
 		sub.Get.Tags = []string{"API Discovery"} // sorts first in the list
 	}
@@ -317,38 +302,3 @@ func (b *DataSourceAPIBuilder) PostProcessOpenAPI(oas *spec3.OpenAPI) (*spec3.Op
 func (b *DataSourceAPIBuilder) GetAPIRoutes() *builder.APIRoutes {
 	return nil
 }
-
-func getExamples(queryTypes *query.QueryTypeDefinitionList) map[string]*spec3.Example {
-	if queryTypes == nil {
-		return nil
-	}
-
-	tr := data.TimeRange{From: "now-1h", To: "now"}
-	examples := map[string]*spec3.Example{}
-	for _, queryType := range queryTypes.Items {
-		for idx, example := range queryType.Spec.Examples {
-			q := data.NewDataQuery(example.SaveModel.Object)
-			q.RefID = "A"
-			for _, dis := range queryType.Spec.Discriminators {
-				_ = q.Set(dis.Field, dis.Value)
-			}
-			if q.MaxDataPoints < 1 {
-				q.MaxDataPoints = 1000
-			}
-			if q.IntervalMS < 1 {
-				q.IntervalMS = 5000 // 5s
-			}
-			examples[fmt.Sprintf("%s-%d", example.Name, idx)] = &spec3.Example{
-				ExampleProps: spec3.ExampleProps{
-					Summary:     example.Name,
-					Description: example.Description,
-					Value: data.QueryDataRequest{
-						TimeRange: tr,
-						Queries:   []data.DataQuery{q},
-					},
-				},
-			}
-		}
-	}
-	return examples
-}