THRIFT-5326: Make PrependError more unwrap-able

Client: go

As a follow up to https://github.com/apache/thrift/pull/2298, make the
error returned by PrependError unwrap-able in certain cases.

3 files changed, 50 insertions, 24 deletions

diff --git a/lib/go/thrift/exception.go b/lib/go/thrift/exception.go
index b6885face..483647995 100644
--- a/lib/go/thrift/exception.go
+++ b/lib/go/thrift/exception.go
@@ -38,11 +38,11 @@ func PrependError(prepend string, err error) error {
 		switch te.TExceptionType() {
 		case TExceptionTypeTransport:
 			if t, ok := err.(TTransportException); ok {
-				return NewTTransportException(t.TypeId(), msg)
+				return prependTTransportException(prepend, t)
 			}
 		case TExceptionTypeProtocol:
 			if t, ok := err.(TProtocolException); ok {
-				return NewTProtocolExceptionWithType(t.TypeId(), errors.New(msg))
+				return prependTProtocolException(prepend, t)
 			}
 		case TExceptionTypeApplication:
 			if t, ok := err.(TApplicationException); ok {
@@ -51,7 +51,8 @@ func PrependError(prepend string, err error) error {
 		}
 
 		return wrappedTException{
-			err:            errors.New(msg),
+			err:            err,
+			msg:            msg,
 			tExceptionType: te.TExceptionType(),
 		}
 	}
@@ -87,17 +88,19 @@ func WrapTException(err error) TException {
 
 	return wrappedTException{
 		err:            err,
+		msg:            err.Error(),
 		tExceptionType: TExceptionTypeUnknown,
 	}
 }
 
 type wrappedTException struct {
 	err            error
+	msg            string
 	tExceptionType TExceptionType
 }
 
 func (w wrappedTException) Error() string {
-	return w.err.Error()
+	return w.msg
 }
 
 func (w wrappedTException) TExceptionType() TExceptionType {
diff --git a/lib/go/thrift/protocol_exception.go b/lib/go/thrift/protocol_exception.go
index b088caf13..c86a9ff47 100644
--- a/lib/go/thrift/protocol_exception.go
+++ b/lib/go/thrift/protocol_exception.go
@@ -42,6 +42,7 @@ const (
 type tProtocolException struct {
 	typeId int
 	err    error
+	msg    string
 }
 
 var _ TProtocolException = (*tProtocolException)(nil)
@@ -55,11 +56,11 @@ func (p *tProtocolException) TypeId() int {
 }
 
 func (p *tProtocolException) String() string {
-	return p.err.Error()
+	return p.msg
 }
 
 func (p *tProtocolException) Error() string {
-	return p.err.Error()
+	return p.msg
 }
 
 func (p *tProtocolException) Unwrap() error {
@@ -70,19 +71,16 @@ func NewTProtocolException(err error) TProtocolException {
 	if err == nil {
 		return nil
 	}
+
 	if e, ok := err.(TProtocolException); ok {
 		return e
 	}
+
 	if _, ok := err.(base64.CorruptInputError); ok {
-		return &tProtocolException{
-			typeId: INVALID_DATA,
-			err:    err,
-		}
-	}
-	return &tProtocolException{
-		typeId: UNKNOWN_PROTOCOL_EXCEPTION,
-		err:    err,
+		return NewTProtocolExceptionWithType(INVALID_DATA, err)
 	}
+
+	return NewTProtocolExceptionWithType(UNKNOWN_PROTOCOL_EXCEPTION, err)
 }
 
 func NewTProtocolExceptionWithType(errType int, err error) TProtocolException {
@@ -92,5 +90,14 @@ func NewTProtocolExceptionWithType(errType int, err error) TProtocolException {
 	return &tProtocolException{
 		typeId: errType,
 		err:    err,
+		msg:    err.Error(),
+	}
+}
+
+func prependTProtocolException(prepend string, err TProtocolException) TProtocolException {
+	return &tProtocolException{
+		typeId: err.TypeId(),
+		err:    err,
+		msg:    prepend + err.Error(),
 	}
 }
diff --git a/lib/go/thrift/transport_exception.go b/lib/go/thrift/transport_exception.go
index cf2cc0059..08f7ce229 100644
--- a/lib/go/thrift/transport_exception.go
+++ b/lib/go/thrift/transport_exception.go
@@ -46,6 +46,7 @@ const (
 type tTransportException struct {
 	typeId int
 	err    error
+	msg    string
 }
 
 var _ TTransportException = (*tTransportException)(nil)
@@ -59,7 +60,7 @@ func (p *tTransportException) TypeId() int {
 }
 
 func (p *tTransportException) Error() string {
-	return p.err.Error()
+	return p.msg
 }
 
 func (p *tTransportException) Err() error {
@@ -75,7 +76,11 @@ func (p *tTransportException) Timeout() bool {
 }
 
 func NewTTransportException(t int, e string) TTransportException {
-	return &tTransportException{typeId: t, err: errors.New(e)}
+	return &tTransportException{
+		typeId: t,
+		err:    errors.New(e),
+		msg:    e,
+	}
 }
 
 func NewTTransportExceptionFromError(e error) TTransportException {
@@ -87,20 +92,31 @@ func NewTTransportExceptionFromError(e error) TTransportException {
 		return t
 	}
 
-	switch v := e.(type) {
-	case TTransportException:
-		return v
-	case timeoutable:
-		if v.Timeout() {
-			return &tTransportException{typeId: TIMED_OUT, err: e}
+	newTTransportException := func(typeID int, err error, prefix string) TTransportException {
+		return &tTransportException{
+			typeId: typeID,
+			err:    err,
+			msg:    prefix + err.Error(),
 		}
 	}
 
+	if isTimeoutError(e) {
+		return newTTransportException(TIMED_OUT, e, "")
+	}
+
 	if e == io.EOF {
-		return &tTransportException{typeId: END_OF_FILE, err: e}
+		return newTTransportException(END_OF_FILE, e, "")
 	}
 
-	return &tTransportException{typeId: UNKNOWN_TRANSPORT_EXCEPTION, err: e}
+	return newTTransportException(UNKNOWN_TRANSPORT_EXCEPTION, e, "")
+}
+
+func prependTTransportException(prepend string, e TTransportException) TTransportException {
+	return &tTransportException{
+		typeId: e.TypeId(),
+		err:    e,
+		msg:    prepend + e.Error(),
+	}
 }
 
 // isTimeoutError returns true when err is a timeout error.