From 813c649c1455fbfa60b6cd6275ef906990db159a Mon Sep 17 00:00:00 2001
From: Dan Rammer <daniel@union.ai>
Date: Tue, 14 Sep 2021 12:40:08 -0500
Subject: [PATCH] Update EventError with Is and Unwrap functions (#221)

* update EventError with Is and Unwrap functions for use in errors API

Signed-off-by: Daniel Rammer <daniel@union.ai>

* changed EventError Is* functionality to check entire error stack

Signed-off-by: Daniel Rammer <daniel@union.ai>

* added nil tests to EventError.Is function

Signed-off-by: Daniel Rammer <daniel@union.ai>
---
 flyteidl/clients/go/events/errors/errors.go | 50 ++++++++++-----------
 1 file changed, 24 insertions(+), 26 deletions(-)

diff --git a/flyteidl/clients/go/events/errors/errors.go b/flyteidl/clients/go/events/errors/errors.go
index cc970c45da..3f71b9c084 100644
--- a/flyteidl/clients/go/events/errors/errors.go
+++ b/flyteidl/clients/go/events/errors/errors.go
@@ -2,6 +2,7 @@ package errors
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/admin"
@@ -31,6 +32,24 @@ func (r EventError) Error() string {
 	return fmt.Sprintf("%s: %s, caused by [%s]", r.Code, r.Message, r.Cause.Error())
 }
 
+func (r *EventError) Is(target error) bool {
+	t, ok := target.(*EventError)
+	if !ok {
+		return false
+	}
+	if r == nil && t == nil {
+		return true
+	}
+	if r == nil || t == nil {
+		return false
+	}
+	return r.Code == t.Code && (r.Cause == t.Cause || t.Cause == nil) && (r.Message == t.Message || t.Message == "")
+}
+
+func (r *EventError) Unwrap() error {
+	return r.Cause
+}
+
 func WrapError(err error) error {
 	// check if error is gRPC, and convert into our own custom error
 	statusErr, ok := status.FromError(err)
@@ -78,46 +97,25 @@ func wrapf(code ErrorCode, cause error, msg string) error {
 
 // Checks if the error is of type EventError and the ErrorCode is of type AlreadyExists
 func IsAlreadyExists(err error) bool {
-	e, ok := err.(*EventError)
-	if ok {
-		return e.Code == AlreadyExists
-	}
-	return false
+	return errors.Is(err, &EventError{Code: AlreadyExists})
 }
 
 // Checks if the error is of type EventError and the ErrorCode is of type InvalidArgument
 func IsInvalidArguments(err error) bool {
-	e, ok := err.(*EventError)
-	if ok {
-		return e.Code == InvalidArgument
-	}
-	return false
+	return errors.Is(err, &EventError{Code: InvalidArgument})
 }
 
 // Checks if the error is of type EventError and the ErrorCode is of type ExecutionNotFound
 func IsNotFound(err error) bool {
-	e, ok := err.(*EventError)
-	if ok {
-		return e.Code == ExecutionNotFound
-	}
-	return false
+	return errors.Is(err, &EventError{Code: ExecutionNotFound})
 }
 
 // Checks if the error is of type EventError and the ErrorCode is of type ResourceExhausted
 func IsResourceExhausted(err error) bool {
-	e, ok := err.(*EventError)
-	if ok {
-		return e.Code == ResourceExhausted
-	}
-	return false
+	return errors.Is(err, &EventError{Code: ResourceExhausted})
 }
 
 // Checks if the error is of type EventError and the ErrorCode is of type EventAlreadyInTerminalStateError
 func IsEventAlreadyInTerminalStateError(err error) bool {
-	// TODO: don't rely on the specific type here as it could be wrapped in another object.
-	e, ok := err.(*EventError)
-	if ok {
-		return e.Code == EventAlreadyInTerminalStateError
-	}
-	return false
+	return errors.Is(err, &EventError{Code: EventAlreadyInTerminalStateError})
 }