THRIFT-4914: Add TResponseHelper
This is the third part of THRIFT-4914, which handles the server writing
part in the response (server -> client direction).
Define a new type, TResponseHelper, which only contains THeader related
functions for now, but can be extended for other functions in the
future.
In TSimpleServer, inject a TResponseHelper into the context object
passed into the handler functions. Handler function code could retrieve
the injected TResponseHelper to set headers to be written to the client.
Client: go
This closes #1923.
diff --git a/lib/go/thrift/response_helper.go b/lib/go/thrift/response_helper.go
new file mode 100644
index 0000000..d884c6a
--- /dev/null
+++ b/lib/go/thrift/response_helper.go
@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package thrift
+
+import (
+ "context"
+)
+
+// See https://godoc.org/context#WithValue on why do we need the unexported typedefs.
+type responseHelperKey struct{}
+
+// TResponseHelper defines a object with a set of helper functions that can be
+// retrieved from the context object passed into server handler functions.
+//
+// Use GetResponseHelper to retrieve the injected TResponseHelper implementation
+// from the context object.
+//
+// The zero value of TResponseHelper is valid with all helper functions being
+// no-op.
+type TResponseHelper struct {
+ // THeader related functions
+ *THeaderResponseHelper
+}
+
+// THeaderResponseHelper defines THeader related TResponseHelper functions.
+//
+// The zero value of *THeaderResponseHelper is valid with all helper functions
+// being no-op.
+type THeaderResponseHelper struct {
+ proto *THeaderProtocol
+}
+
+// NewTHeaderResponseHelper creates a new THeaderResponseHelper from the
+// underlying TProtocol.
+func NewTHeaderResponseHelper(proto TProtocol) *THeaderResponseHelper {
+ if hp, ok := proto.(*THeaderProtocol); ok {
+ return &THeaderResponseHelper{
+ proto: hp,
+ }
+ }
+ return nil
+}
+
+// SetHeader sets a response header.
+//
+// It's no-op if the underlying protocol/transport does not support THeader.
+func (h *THeaderResponseHelper) SetHeader(key, value string) {
+ if h != nil && h.proto != nil {
+ h.proto.SetWriteHeader(key, value)
+ }
+}
+
+// ClearHeaders clears all the response headers previously set.
+//
+// It's no-op if the underlying protocol/transport does not support THeader.
+func (h *THeaderResponseHelper) ClearHeaders() {
+ if h != nil && h.proto != nil {
+ h.proto.ClearWriteHeaders()
+ }
+}
+
+// GetResponseHelper retrieves the TResponseHelper implementation injected into
+// the context object.
+//
+// If no helper was found in the context object, a nop helper with ok == false
+// will be returned.
+func GetResponseHelper(ctx context.Context) (helper TResponseHelper, ok bool) {
+ if v := ctx.Value(responseHelperKey{}); v != nil {
+ helper, ok = v.(TResponseHelper)
+ }
+ return
+}
+
+// SetResponseHelper injects TResponseHelper into the context object.
+func SetResponseHelper(ctx context.Context, helper TResponseHelper) context.Context {
+ return context.WithValue(ctx, responseHelperKey{}, helper)
+}
diff --git a/lib/go/thrift/response_helper_test.go b/lib/go/thrift/response_helper_test.go
new file mode 100644
index 0000000..69f76d3
--- /dev/null
+++ b/lib/go/thrift/response_helper_test.go
@@ -0,0 +1,137 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package thrift
+
+import (
+ "context"
+ "testing"
+)
+
+func TestResponseHelperContext(t *testing.T) {
+ ctx := context.Background()
+
+ t.Run(
+ "empty-noop",
+ func(t *testing.T) {
+ helper, ok := GetResponseHelper(ctx)
+ if ok {
+ t.Error("GetResponseHelper expected ok == false")
+ }
+ // Just make sure those function calls does not panic
+ helper.SetHeader("foo", "bar")
+ helper.ClearHeaders()
+ },
+ )
+
+ t.Run(
+ "set-get",
+ func(t *testing.T) {
+ trans := NewTHeaderTransport(NewTMemoryBuffer())
+ proto := NewTHeaderProtocol(trans)
+ ctx = SetResponseHelper(
+ ctx,
+ TResponseHelper{
+ THeaderResponseHelper: NewTHeaderResponseHelper(proto),
+ },
+ )
+ helper, ok := GetResponseHelper(ctx)
+ if !ok {
+ t.Error("GetResponseHelper expected ok == true")
+ }
+ if helper.THeaderResponseHelper == nil {
+ t.Error("GetResponseHelper expected THeaderResponseHelper to be non-nil")
+ }
+ },
+ )
+}
+
+func TestHeaderHelper(t *testing.T) {
+ t.Run(
+ "THeaderProtocol",
+ func(t *testing.T) {
+ trans := NewTHeaderTransport(NewTMemoryBuffer())
+ proto := NewTHeaderProtocol(trans)
+ helper := NewTHeaderResponseHelper(proto)
+
+ const (
+ key = "key"
+ value = "value"
+ )
+ helper.SetHeader(key, value)
+ if len(trans.writeHeaders) != 1 {
+ t.Errorf(
+ "Expected THeaderTransport.writeHeaders to be with size of 1, got %+v",
+ trans.writeHeaders,
+ )
+ }
+ actual := trans.writeHeaders[key]
+ if actual != value {
+ t.Errorf(
+ "Expected THeaderTransport.writeHeaders to have %q:%q, got %+v",
+ key,
+ value,
+ trans.writeHeaders,
+ )
+ }
+ helper.ClearHeaders()
+ if len(trans.writeHeaders) != 0 {
+ t.Errorf(
+ "Expected THeaderTransport.writeHeaders to be empty after ClearHeaders call, got %+v",
+ trans.writeHeaders,
+ )
+ }
+ },
+ )
+
+ t.Run(
+ "other-protocol",
+ func(t *testing.T) {
+ trans := NewTMemoryBuffer()
+ proto := NewTCompactProtocol(trans)
+ helper := NewTHeaderResponseHelper(proto)
+
+ // We only need to make sure that functions in helper
+ // don't panic here.
+ helper.SetHeader("foo", "bar")
+ helper.ClearHeaders()
+ },
+ )
+
+ t.Run(
+ "zero-value",
+ func(t *testing.T) {
+ var helper *THeaderResponseHelper
+
+ // We only need to make sure that functions in helper
+ // don't panic here.
+ helper.SetHeader("foo", "bar")
+ helper.ClearHeaders()
+ },
+ )
+}
+
+func TestTResponseHelperZeroValue(t *testing.T) {
+ var helper THeaderResponseHelper
+
+ // We only need to make sure that functions in helper
+ // don't panic here.
+ helper.SetHeader("foo", "bar")
+ helper.ClearHeaders()
+}
diff --git a/lib/go/thrift/simple_server.go b/lib/go/thrift/simple_server.go
index 756d4cf..5a9c9c9 100644
--- a/lib/go/thrift/simple_server.go
+++ b/lib/go/thrift/simple_server.go
@@ -272,7 +272,12 @@
return nil
}
- ctx := defaultCtx
+ ctx := SetResponseHelper(
+ defaultCtx,
+ TResponseHelper{
+ THeaderResponseHelper: NewTHeaderResponseHelper(outputProtocol),
+ },
+ )
if headerProtocol != nil {
// We need to call ReadFrame here, otherwise we won't
// get any headers on the AddReadTHeaderToContext call.