diff options
author | Ash McKenzie <amckenzie@gitlab.com> | 2019-10-29 00:58:31 +0000 |
---|---|---|
committer | Ash McKenzie <amckenzie@gitlab.com> | 2019-10-29 00:58:31 +0000 |
commit | 0afa8ec5fcc571d7bdfb5f52533e3df4ab78f793 (patch) | |
tree | cc0ac9e4c29b4f3b954aa9912403fe1bfc9c9867 /internal/handler | |
parent | 7d36bc3bba67d42f7d23bd16624ef5e6aab2bddc (diff) | |
parent | bee2c423858c0763cce91bcacf345eb0b45227a6 (diff) | |
download | gitlab-shell-0afa8ec5fcc571d7bdfb5f52533e3df4ab78f793.tar.gz |
Merge branch 'pass-ff-to-gitaly' into 'master'
Add support for Gitaly feature flags
See merge request gitlab-org/gitlab-shell!351
Diffstat (limited to 'internal/handler')
-rw-r--r-- | internal/handler/exec.go | 16 | ||||
-rw-r--r-- | internal/handler/exec_test.go | 44 |
2 files changed, 60 insertions, 0 deletions
diff --git a/internal/handler/exec.go b/internal/handler/exec.go index ba9a4ff..19621fa 100644 --- a/internal/handler/exec.go +++ b/internal/handler/exec.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "os" + "strings" "gitlab.com/gitlab-org/gitaly/auth" "gitlab.com/gitlab-org/gitaly/client" @@ -11,6 +12,7 @@ import ( "gitlab.com/gitlab-org/gitlab-shell/internal/config" "gitlab.com/gitlab-org/labkit/tracing" "google.golang.org/grpc" + "google.golang.org/grpc/metadata" ) // GitalyHandlerFunc implementations are responsible for making @@ -29,6 +31,7 @@ type GitalyCommand struct { ServiceName string Address string Token string + Features map[string]string } // RunGitalyCommand provides a bootstrap for Gitaly commands executed @@ -48,6 +51,18 @@ func (gc *GitalyCommand) RunGitalyCommand(handler GitalyHandlerFunc) error { return err } +func withOutgoingMetadata(ctx context.Context, features map[string]string) context.Context { + md := metadata.New(nil) + for k, v := range features { + if !strings.HasPrefix(k, "gitaly-feature-") { + continue + } + md.Append(k, v) + } + + return metadata.NewOutgoingContext(ctx, md) +} + func getConn(gc *GitalyCommand) (*GitalyConn, error) { if gc.Address == "" { return nil, fmt.Errorf("no gitaly_address given") @@ -80,6 +95,7 @@ func getConn(gc *GitalyCommand) (*GitalyConn, error) { ) ctx, finished := tracing.ExtractFromEnv(context.Background()) + ctx = withOutgoingMetadata(ctx, gc.Features) conn, err := client.Dial(gc.Address, connOpts) if err != nil { diff --git a/internal/handler/exec_test.go b/internal/handler/exec_test.go index 6c7d3f5..7c85e16 100644 --- a/internal/handler/exec_test.go +++ b/internal/handler/exec_test.go @@ -5,8 +5,10 @@ import ( "errors" "testing" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "google.golang.org/grpc" + "google.golang.org/grpc/metadata" "gitlab.com/gitlab-org/gitlab-shell/internal/config" ) @@ -40,3 +42,45 @@ func TestMissingGitalyAddress(t *testing.T) { err := cmd.RunGitalyCommand(makeHandler(t, nil)) require.EqualError(t, err, "no gitaly_address given") } + +func TestGetConnMetadata(t *testing.T) { + tests := []struct { + name string + gc *GitalyCommand + want map[string]string + }{ + { + name: "gitaly_feature_flags", + gc: &GitalyCommand{ + Config: &config.Config{}, + Address: "tcp://localhost:9999", + Features: map[string]string{ + "gitaly-feature-cache_invalidator": "true", + "other-ff": "true", + "gitaly-feature-inforef_uploadpack_cache": "false", + }, + }, + want: map[string]string{ + "gitaly-feature-cache_invalidator": "true", + "gitaly-feature-inforef_uploadpack_cache": "false", + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + conn, err := getConn(tt.gc) + require.NoError(t, err) + + md, exists := metadata.FromOutgoingContext(conn.ctx) + require.True(t, exists) + require.Equal(t, len(tt.want), md.Len()) + + for k, v := range tt.want { + values := md.Get(k) + assert.Equal(t, 1, len(values)) + assert.Equal(t, v, values[0]) + } + + }) + } +} |