Merge pull request #32 from dsrhub/zz/add-dsrhub-grpc-header

Add grpc header support
open-privacy · Oct 6, 2020 · cf44a5d · cf44a5d
2 parents 9935d7f + faa3196
commit cf44a5d
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 11 deletions.
diff --git a/plugins/dsrhub_grpc/dsrhub_grpc.go b/plugins/dsrhub_grpc/dsrhub_grpc.go
@@ -5,33 +5,36 @@ import (
 	"fmt"
 	"time"
 
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+
 	"github.com/dsrhub/dsrhub/idl_dsrhub"
 	"github.com/ovh/utask/pkg/plugins/taskplugin"
-	"google.golang.org/grpc"
 )
 
 var (
 	// Plugin is the dsrhub_grpc plugin that we can use as a customized plugin
 	// nolint
 	Plugin = taskplugin.New(
 		"dsrhub_grpc", "1.0", exec,
-		taskplugin.WithConfig(validConfig, DsrHubGRPCConfig{}),
+		taskplugin.WithConfig(validConfig, DSRHubGRPCConfig{}),
 	)
 )
 
 const (
 	defaultTimeoutSeconds = 10
 )
 
-// DsrHubGRPCConfig is the configuration needed to perform an gRPC client side call
-type DsrHubGRPCConfig struct {
+// DSRHubGRPCConfig is the configuration needed to perform an gRPC client side call
+type DSRHubGRPCConfig struct {
 	URL     string                      `json:"url"`
 	Request idl_dsrhub.CreateDSRRequest `json:"request"`
 	Timeout int                         `json:"timeout,omitempty"` // timeout in seconds
+	Header  map[string]string           `json:"header,omitempty"`
 }
 
 func validConfig(config interface{}) error {
-	cfg := config.(*DsrHubGRPCConfig)
+	cfg := config.(*DSRHubGRPCConfig)
 	if cfg.URL == "" {
 		return fmt.Errorf("invalid dsrhub_grpc config url: empty url")
 	}
@@ -50,8 +53,8 @@ func validConfig(config interface{}) error {
 	return nil
 }
 
-func exec(stepName string, config interface{}, execCtx interface{}) (output interface{}, metadata interface{}, err error) {
-	cfg := config.(*DsrHubGRPCConfig)
+func exec(stepName string, config interface{}, execCtx interface{}) (output interface{}, _ interface{}, err error) {
+	cfg := config.(*DSRHubGRPCConfig)
 
 	// TODO: support secure connection with configuration
 	conn, err := grpc.Dial(cfg.URL, grpc.WithInsecure())
@@ -69,7 +72,8 @@ func exec(stepName string, config interface{}, execCtx interface{}) (output inte
 	)
 	defer cancel()
 
-	res, err := idl_dsrhub.NewDSRHubServiceClient(conn).CreateDSR(ctx, &cfg.Request)
+	header := metadata.New(cfg.Header)
+	res, err := idl_dsrhub.NewDSRHubServiceClient(conn).CreateDSR(ctx, &cfg.Request, grpc.Header(&header))
 	if err != nil {
 		return nil, nil, err
 	}

diff --git a/plugins/dsrhub_grpc/dsrhub_grpc_test.go b/plugins/dsrhub_grpc/dsrhub_grpc_test.go
@@ -8,12 +8,12 @@ import (
 )
 
 func TestValidConfig(t *testing.T) {
-	var cfg DsrHubGRPCConfig
+	var cfg DSRHubGRPCConfig
 
-	cfg = DsrHubGRPCConfig{}
+	cfg = DSRHubGRPCConfig{}
 	assert.Error(t, validConfig(&cfg))
 
-	cfg = DsrHubGRPCConfig{
+	cfg = DSRHubGRPCConfig{
 		URL: "localhost:50051",
 		Request: idl_dsrhub.CreateDSRRequest{
 			Regulation:         "gdpr",