deprecate thrift.WithProtocol

Summary: deprecate thrift.WithProtocol

Reviewed By: codyohl

Differential Revision: D59577298

fbshipit-source-id: 671b7c90aa984eea0af7d94640a6a4552cf5d6ba

Author: awalterschulze

thrift/lib/go/thrift/context.go

@@ -78,17 +78,12 @@ func ConnInfoFromContext(ctx context.Context) (ConnInfo, bool) {
 	return v, ok
 }
 
-// WithProtocol attaches thrift protocol to a context
+// Deprecated: use thrift.WithHeaders(ctx, proto.GetResponseHeaders())
 func WithProtocol(ctx context.Context, proto Protocol) context.Context {
 	return context.WithValue(ctx, protocolKey, proto)
 }
 
-// HeadersFromContext extracts headers for this message, both per-message
-// and persistent headers. When both a per-message header and a persistent
-// header exist with the same name, the persistent header is returned. This is
-// also the behaviour of the C++ implementation.
-// This function returns nil when the underlying transport/protocol do not
-// support headers.
+// Deprecated: Use thrift.GetHeaders
 func HeadersFromContext(ctx context.Context) map[string]string {
 	t, ok := ctx.Value(protocolKey).(Protocol)
 	if !ok {

thrift/lib/go/thrift/context_headers.go

@@ -43,6 +43,30 @@ func AddHeader(ctx context.Context, key string, value string) (context.Context,
 	return ctx, nil
 }
 
+// WithHeaders attaches thrift headers to a ctx
+func WithHeaders(ctx context.Context, headers map[string]string) context.Context {
+	return context.WithValue(ctx, headersKey, headers)
+}
+
+// GetHeaders gets thrift headers from ctx.
+// WithHeaders will set this, and if nothing is set it
+// will defer to thrift.HeadersFromContext
+// can return nil if there is no headers.
+func GetHeaders(ctx context.Context) map[string]string {
+	// check for headersKey
+	v, ok := ctx.Value(headersKey).(map[string]string)
+	if ok {
+		return v
+	}
+	// deprecated
+	t, ok := ctx.Value(protocolKey).(Protocol)
+	if !ok {
+		// A nil map behaves like an empty map for reading.
+		return nil
+	}
+	return t.GetResponseHeaders()
+}
+
 // Deprecated: RequestHeaders will eventually be private.
 type RequestHeaders interface {
 	SetRequestHeader(key, value string)