mirror of https://github.com/grafana/loki
feat: Introduce special header that tells Loki not to modify query results (#12327)
Travis Patterson
2 years ago
committed by
GitHub
parent
ab17a50b90
commit
c8d2c32002
@ -0,0 +1,61 @@ |
||||
package server |
||||
|
||||
import ( |
||||
"context" |
||||
|
||||
"google.golang.org/grpc" |
||||
"google.golang.org/grpc/metadata" |
||||
|
||||
"github.com/grafana/loki/pkg/util/httpreq" |
||||
) |
||||
|
||||
func injectHTTPHeadersIntoGRPCRequest(ctx context.Context) context.Context { |
||||
header := httpreq.ExtractHeader(ctx, httpreq.LokiDisablePipelineWrappersHeader) |
||||
if header == "" { |
||||
return ctx |
||||
} |
||||
|
||||
// inject into GRPC metadata
|
||||
md, ok := metadata.FromOutgoingContext(ctx) |
||||
if !ok { |
||||
md = metadata.New(map[string]string{}) |
||||
} |
||||
md = md.Copy() |
||||
md.Set(httpreq.LokiDisablePipelineWrappersHeader, header) |
||||
|
||||
return metadata.NewOutgoingContext(ctx, md) |
||||
} |
||||
|
||||
func extractHTTPHeadersFromGRPCRequest(ctx context.Context) context.Context { |
||||
md, ok := metadata.FromIncomingContext(ctx) |
||||
if !ok { |
||||
// No metadata, just return as is
|
||||
return ctx |
||||
} |
||||
|
||||
headerValues := md.Get(httpreq.LokiDisablePipelineWrappersHeader) |
||||
if len(headerValues) == 0 { |
||||
return ctx |
||||
} |
||||
|
||||
return httpreq.InjectHeader(ctx, httpreq.LokiDisablePipelineWrappersHeader, headerValues[0]) |
||||
} |
||||
|
||||
func UnaryClientHTTPHeadersInterceptor(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error { |
||||
return invoker(injectHTTPHeadersIntoGRPCRequest(ctx), method, req, reply, cc, opts...) |
||||
} |
||||
|
||||
func StreamClientHTTPHeadersInterceptor(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) { |
||||
return streamer(injectHTTPHeadersIntoGRPCRequest(ctx), desc, cc, method, opts...) |
||||
} |
||||
|
||||
func UnaryServerHTTPHeadersnIterceptor(ctx context.Context, req interface{}, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) { |
||||
return handler(extractHTTPHeadersFromGRPCRequest(ctx), req) |
||||
} |
||||
|
||||
func StreamServerHTTPHeadersInterceptor(srv interface{}, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error { |
||||
return handler(srv, serverStream{ |
||||
ctx: extractHTTPHeadersFromGRPCRequest(ss.Context()), |
||||
ServerStream: ss, |
||||
}) |
||||
} |
||||
@ -0,0 +1,81 @@ |
||||
package server |
||||
|
||||
import ( |
||||
"context" |
||||
"testing" |
||||
|
||||
"github.com/stretchr/testify/require" |
||||
"google.golang.org/grpc/metadata" |
||||
|
||||
"github.com/grafana/loki/pkg/util/httpreq" |
||||
) |
||||
|
||||
func TestInjectHTTPHeaderIntoGRPCRequest(t *testing.T) { |
||||
for _, tt := range []struct { |
||||
name, header string |
||||
md, expectMetadata metadata.MD |
||||
}{ |
||||
{ |
||||
name: "creates new metadata and sets header", |
||||
header: "true", |
||||
expectMetadata: metadata.New(map[string]string{httpreq.LokiDisablePipelineWrappersHeader: "true"}), |
||||
}, |
||||
{ |
||||
name: "sets header on existing metadata", |
||||
header: "true", |
||||
md: metadata.New(map[string]string{"x-foo": "bar"}), |
||||
expectMetadata: metadata.New(map[string]string{"x-foo": "bar", httpreq.LokiDisablePipelineWrappersHeader: "true"}), |
||||
}, |
||||
{ |
||||
name: "no header, leave metadata untouched", |
||||
md: metadata.New(map[string]string{"x-foo": "bar"}), |
||||
expectMetadata: metadata.New(map[string]string{"x-foo": "bar"}), |
||||
}, |
||||
{ |
||||
name: "no header", |
||||
expectMetadata: nil, |
||||
}, |
||||
} { |
||||
t.Run(tt.name, func(t *testing.T) { |
||||
ctx := context.Background() |
||||
if tt.header != "" { |
||||
ctx = httpreq.InjectHeader(context.Background(), httpreq.LokiDisablePipelineWrappersHeader, tt.header) |
||||
} |
||||
|
||||
if tt.md != nil { |
||||
ctx = metadata.NewOutgoingContext(ctx, tt.md) |
||||
} |
||||
|
||||
ctx = injectHTTPHeadersIntoGRPCRequest(ctx) |
||||
md, _ := metadata.FromOutgoingContext(ctx) |
||||
require.EqualValues(t, tt.expectMetadata, md) |
||||
}) |
||||
} |
||||
} |
||||
|
||||
func TestExtractHTTPHeaderFromGRPCRequest(t *testing.T) { |
||||
for _, tt := range []struct { |
||||
name string |
||||
md metadata.MD |
||||
expectedResp string |
||||
}{ |
||||
{ |
||||
name: "extracts header from metadata", |
||||
md: metadata.New(map[string]string{httpreq.LokiDisablePipelineWrappersHeader: "true"}), |
||||
expectedResp: "true", |
||||
}, |
||||
{ |
||||
name: "non-nil metadata without header", |
||||
md: metadata.New(map[string]string{"x-foo": "bar"}), |
||||
}, |
||||
{ |
||||
name: "nil metadata", |
||||
}, |
||||
} { |
||||
t.Run(tt.name, func(t *testing.T) { |
||||
ctx := metadata.NewIncomingContext(context.Background(), tt.md) |
||||
ctx = extractHTTPHeadersFromGRPCRequest(ctx) |
||||
require.Equal(t, tt.expectedResp, httpreq.ExtractHeader(ctx, httpreq.LokiDisablePipelineWrappersHeader)) |
||||
}) |
||||
} |
||||
} |
||||
Loading…
Reference in new issue