Add access token to internal context

* Extract access token from header
* Inject access token in ctx
* Refactor/split constants in separate file

Author: kalil-pelissier

pkg/trino/constants/keys.go

@@ -0,0 +1,6 @@
+package constants
+
+const (
+	AccessTokenKey  = "accessToken"
+	TrinoUserHeader = "X-Trino-User"
+)

pkg/trino/datasource-context.go

@@ -3,12 +3,17 @@ package trino
 import (
 	"context"
 	"fmt"
+	"strings"
+
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
 	"github.com/grafana/sqlds/v2"
+	"github.com/trinodb/grafana-trino/pkg/trino/constants"
 	"github.com/trinodb/grafana-trino/pkg/trino/models"
 )
 
+const bearerPrefix = "Bearer "
+
 type SQLDatasourceWithTrinoUserContext struct {
 	sqlds.SQLDatasource
 }
@@ -21,13 +26,15 @@ func (ds *SQLDatasourceWithTrinoUserContext) QueryData(ctx context.Context, req
 		return nil, fmt.Errorf("error reading settings: %s", err.Error())
 	}
 
+	ctx = injectAccessToken(ctx, req)
+
 	if settings.EnableImpersonation {
 		user := req.PluginContext.User
 		if user == nil {
 			return nil, fmt.Errorf("user can't be nil if impersonation is enabled")
 		}
 
-		ctx = context.WithValue(ctx, "X-Trino-User", user)
+		ctx = context.WithValue(ctx, constants.TrinoUserHeader, user)
 	}
 
 	return ds.SQLDatasource.QueryData(ctx, req)
@@ -45,3 +52,21 @@ func NewDatasource(c sqlds.Driver) *SQLDatasourceWithTrinoUserContext {
 	base := sqlds.NewDatasource(c)
 	return &SQLDatasourceWithTrinoUserContext{*base}
 }
+
+func extractBearerToken(header string) (string, bool) {
+	if strings.HasPrefix(header, bearerPrefix) {
+		return strings.TrimPrefix(header, bearerPrefix), true
+	}
+	return "", false
+}
+
+func injectAccessToken(ctx context.Context, req *backend.QueryDataRequest) context.Context {
+	header := req.GetHTTPHeader(backend.OAuthIdentityTokenHeaderName)
+
+	token, ok := extractBearerToken(header)
+	if !ok {
+		return ctx
+	}
+
+	return context.WithValue(ctx, constants.AccessTokenKey, token)
+}