From 331dc40418bd6cfd71a788e6a12df520254e0112 Mon Sep 17 00:00:00 2001
From: Sri Harsha CH <sriharshach@google.com>
Date: Mon, 17 Jul 2023 13:33:20 +0000
Subject: [PATCH] feat(spanner): wrap error into inline begin transaction
 failure to unmask original errors

---
 spanner/errors.go      | 40 ++++++++++++++++++++++++++++++++++++++++
 spanner/session.go     |  2 +-
 spanner/transaction.go | 10 +++++-----
 3 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/spanner/errors.go b/spanner/errors.go
index f4451812765a..c1007d7e9197 100644
--- a/spanner/errors.go
+++ b/spanner/errors.go
@@ -73,6 +73,22 @@ func (*TransactionOutcomeUnknownError) Error() string { return transactionOutcom
 // Unwrap returns the wrapped error (if any).
 func (e *TransactionOutcomeUnknownError) Unwrap() error { return e.err }
 
+// InlineBeginTransactionFailedError is wrapped in a Spanner error when the error
+// occurred during an inline begin transaction.
+type InlineBeginTransactionFailedError struct {
+	// err is the wrapped error that caused this InlineBeginTransactionFailedError
+	// error. The wrapped error can be read with the Unwrap method.
+	err error
+}
+
+const inlineBeginTransactionFailedMsg = "failed inline begin transaction"
+
+// Error implements error.Error.
+func (*InlineBeginTransactionFailedError) Error() string { return inlineBeginTransactionFailedMsg }
+
+// Unwrap returns the wrapped error (if any).
+func (e *InlineBeginTransactionFailedError) Unwrap() error { return e.err }
+
 // Error implements error.Error.
 func (e *Error) Error() string {
 	if e == nil {
@@ -181,6 +197,30 @@ func toSpannerErrorWithCommitInfo(err error, errorDuringCommit bool) error {
 	}
 }
 
+// toSpannerErrorDuringInlineBegin general Go error to *spanner.Error with additional information
+// if the error occurred during an inline begin transaction.
+//
+// If err is already a *spanner.Error, err is returned unmodified.
+func toSpannerErrorDuringInlineBegin(err error) error {
+	if err == nil {
+		return nil
+	}
+	var se *Error
+	if errorAs(err, &se) {
+		return se
+	}
+	statusErr := status.Convert(err)
+	code, desc := statusErr.Code(), statusErr.Message()
+	desc = fmt.Sprintf("%s, %s", desc, inlineBeginTransactionFailedMsg)
+	wrapped := &InlineBeginTransactionFailedError{err: err}
+	return &Error{
+		code,
+		toAPIError(wrapped),
+		desc,
+		"",
+	}
+}
+
 // ErrCode extracts the canonical error code from a Go error.
 func ErrCode(err error) codes.Code {
 	s, ok := status.FromError(err)
diff --git a/spanner/session.go b/spanner/session.go
index 2fc4809ff34a..6613edbbe46b 100644
--- a/spanner/session.go
+++ b/spanner/session.go
@@ -1496,7 +1496,7 @@ func isFailedInlineBeginTransaction(err error) bool {
 	if err == nil {
 		return false
 	}
-	return ErrCode(err) == codes.Internal && strings.Contains(err.Error(), errInlineBeginTransactionFailed().Error())
+	return strings.Contains(err.Error(), errInlineBeginTransactionFailed().Error())
 }
 
 // isClientClosing returns true if the given error is a
diff --git a/spanner/transaction.go b/spanner/transaction.go
index 85de18327f89..1b8dc414320b 100644
--- a/spanner/transaction.go
+++ b/spanner/transaction.go
@@ -264,7 +264,7 @@ func (t *txReadOnly) ReadWithOptions(ctx context.Context, table string, keys Key
 			if err != nil {
 				if _, ok := t.getTransactionSelector().GetSelector().(*sppb.TransactionSelector_Begin); ok {
 					t.setTransactionID(nil)
-					return client, errInlineBeginTransactionFailed()
+					return client, toSpannerErrorDuringInlineBegin(err)
 				}
 				return client, err
 			}
@@ -301,7 +301,7 @@ func errMultipleRowsFound(table string, key Key, index string) error {
 
 // errInlineBeginTransactionFailed returns error for read-write transaction to explicitly begin the transaction
 func errInlineBeginTransactionFailed() error {
-	return spannerErrorf(codes.Internal, "failed inline begin transaction")
+	return spannerErrorf(codes.Internal, inlineBeginTransactionFailedMsg)
 }
 
 // ReadRow reads a single row from the database.
@@ -500,7 +500,7 @@ func (t *txReadOnly) query(ctx context.Context, statement Statement, options Que
 			if err != nil {
 				if _, ok := req.Transaction.GetSelector().(*sppb.TransactionSelector_Begin); ok {
 					t.setTransactionID(nil)
-					return client, errInlineBeginTransactionFailed()
+					return client, toSpannerErrorDuringInlineBegin(err)
 				}
 				return client, err
 			}
@@ -1071,7 +1071,7 @@ func (t *ReadWriteTransaction) update(ctx context.Context, stmt Statement, opts
 	if err != nil {
 		if hasInlineBeginTransaction {
 			t.setTransactionID(nil)
-			return 0, errInlineBeginTransactionFailed()
+			return 0, toSpannerErrorDuringInlineBegin(err)
 		}
 		return 0, ToSpannerError(err)
 	}
@@ -1166,7 +1166,7 @@ func (t *ReadWriteTransaction) batchUpdateWithOptions(ctx context.Context, stmts
 	if err != nil {
 		if hasInlineBeginTransaction {
 			t.setTransactionID(nil)
-			return nil, errInlineBeginTransactionFailed()
+			return nil, toSpannerErrorDuringInlineBegin(err)
 		}
 		return nil, ToSpannerError(err)
 	}