Skip to content

Commit c4776e4

Browse files
olavloiteolavloite
authored
fix: safeguard against statements errors when requesting tx (#800)
1 parent 5b864a1 commit c4776e4

File tree

7 files changed

+344
-72
lines changed

7 files changed

+344
-72
lines changed

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,7 @@ TransactionSelector getTransactionSelector() {
213213
}
214214

215215
@Override
216-
public void onTransactionMetadata(Transaction transaction) {
216+
public void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId) {
217217
synchronized (lock) {
218218
if (!transaction.hasReadTimestamp()) {
219219
throw newSpannerException(
@@ -394,6 +394,9 @@ void initTransaction() {
394394
// much more frequently.
395395
private static final int MAX_BUFFERED_CHUNKS = 512;
396396

397+
protected static final String NO_TRANSACTION_RETURNED_MSG =
398+
"The statement did not return a transaction even though one was requested";
399+
397400
AbstractReadContext(Builder<?, ?> builder) {
398401
this.session = builder.session;
399402
this.rpc = builder.rpc;
@@ -632,7 +635,7 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
632635
SpannerRpc.StreamingCall call =
633636
rpc.executeQuery(request.build(), stream.consumer(), session.getOptions());
634637
call.request(prefetchChunks);
635-
stream.setCall(call, request.hasTransaction() && request.getTransaction().hasBegin());
638+
stream.setCall(call, request.getTransaction().hasBegin());
636639
return stream;
637640
}
638641
};
@@ -685,7 +688,7 @@ public void close() {
685688

686689
/** This method is called when a statement returned a new transaction as part of its results. */
687690
@Override
688-
public void onTransactionMetadata(Transaction transaction) {}
691+
public void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId) {}
689692

690693
@Override
691694
public void onError(SpannerException e, boolean withBeginTransaction) {}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,8 @@ interface Listener {
7878
* Called when transaction metadata is seen. This method may be invoked at most once. If the
7979
* method is invoked, it will precede {@link #onError(SpannerException)} or {@link #onDone()}.
8080
*/
81-
void onTransactionMetadata(Transaction transaction) throws SpannerException;
81+
void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId)
82+
throws SpannerException;
8283

8384
/** Called when the read finishes with an error. */
8485
void onError(SpannerException e, boolean withBeginTransaction);
@@ -117,12 +118,12 @@ public boolean next() throws SpannerException {
117118
if (currRow == null) {
118119
ResultSetMetadata metadata = iterator.getMetadata();
119120
if (metadata.hasTransaction()) {
120-
listener.onTransactionMetadata(metadata.getTransaction());
121+
listener.onTransactionMetadata(
122+
metadata.getTransaction(), iterator.isWithBeginTransaction());
121123
} else if (iterator.isWithBeginTransaction()) {
122124
// The query should have returned a transaction.
123125
throw SpannerExceptionFactory.newSpannerException(
124-
ErrorCode.FAILED_PRECONDITION,
125-
"Query requested a transaction to be started, but no transaction was returned");
126+
ErrorCode.FAILED_PRECONDITION, AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
126127
}
127128
currRow = new GrpcStruct(iterator.type(), new ArrayList<>());
128129
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java‎

Lines changed: 51 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -463,13 +463,9 @@ TransactionSelector getTransactionSelector() {
463463
// Aborted error if the call that included the BeginTransaction option fails. The
464464
// Aborted error will cause the entire transaction to be retried, and the retry will use
465465
// a separate BeginTransaction RPC.
466-
if (trackTransactionStarter) {
467-
TransactionSelector.newBuilder()
468-
.setId(tx.get(waitForTransactionTimeoutMillis, TimeUnit.MILLISECONDS))
469-
.build();
470-
} else {
471-
TransactionSelector.newBuilder().setId(tx.get()).build();
472-
}
466+
TransactionSelector.newBuilder()
467+
.setId(tx.get(waitForTransactionTimeoutMillis, TimeUnit.MILLISECONDS))
468+
.build();
473469
}
474470
} catch (ExecutionException e) {
475471
if (e.getCause() instanceof AbortedException) {
@@ -479,11 +475,15 @@ TransactionSelector getTransactionSelector() {
479475
}
480476
throw SpannerExceptionFactory.newSpannerException(e.getCause());
481477
} catch (TimeoutException e) {
478+
// Throw an ABORTED exception to force a retry of the transaction if no transaction
479+
// has been returned by the first statement.
482480
SpannerException se =
483481
SpannerExceptionFactory.newSpannerException(
484-
ErrorCode.DEADLINE_EXCEEDED,
485-
"Timeout while waiting for a transaction to be returned by another statement. "
486-
+ "See the suppressed exception for the stacktrace of the caller that should return a transaction",
482+
ErrorCode.ABORTED,
483+
"Timeout while waiting for a transaction to be returned by another statement."
484+
+ (trackTransactionStarter
485+
? " See the suppressed exception for the stacktrace of the caller that should return a transaction"
486+
: ""),
487487
e);
488488
if (transactionStarter != null) {
489489
se.addSuppressed(transactionStarter);
@@ -498,12 +498,20 @@ TransactionSelector getTransactionSelector() {
498498
}
499499

500500
@Override
501-
public void onTransactionMetadata(Transaction transaction) {
502-
// A transaction has been returned by a statement that was executed. Set the id of the
503-
// transaction on this instance and release the lock to allow other statements to proceed.
504-
if (this.transactionId == null && transaction != null && transaction.getId() != null) {
505-
this.transactionId = transaction.getId();
506-
this.transactionIdFuture.set(transaction.getId());
501+
public void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId) {
502+
Preconditions.checkNotNull(transaction);
503+
if (transaction.getId() != ByteString.EMPTY) {
504+
// A transaction has been returned by a statement that was executed. Set the id of the
505+
// transaction on this instance and release the lock to allow other statements to proceed.
506+
if ((transactionIdFuture == null || !this.transactionIdFuture.isDone())
507+
&& this.transactionId == null) {
508+
this.transactionId = transaction.getId();
509+
this.transactionIdFuture.set(transaction.getId());
510+
}
511+
} else if (shouldIncludeId) {
512+
// The statement should have returned a transaction.
513+
throw SpannerExceptionFactory.newSpannerException(
514+
ErrorCode.FAILED_PRECONDITION, AbstractReadContext.NO_TRANSACTION_RETURNED_MSG);
507515
}
508516
}
509517

@@ -580,17 +588,18 @@ public long executeUpdate(Statement statement, UpdateOption... options) {
580588
com.google.spanner.v1.ResultSet resultSet =
581589
rpc.executeQuery(builder.build(), session.getOptions());
582590
if (resultSet.getMetadata().hasTransaction()) {
583-
onTransactionMetadata(resultSet.getMetadata().getTransaction());
591+
onTransactionMetadata(
592+
resultSet.getMetadata().getTransaction(), builder.getTransaction().hasBegin());
584593
}
585594
if (!resultSet.hasStats()) {
586595
throw new IllegalArgumentException(
587596
"DML response missing stats possibly due to non-DML statement as input");
588597
}
589598
// For standard DML, using the exact row count.
590599
return resultSet.getStats().getRowCountExact();
591-
} catch (SpannerException e) {
592-
onError(e, builder.hasTransaction() && builder.getTransaction().hasBegin());
593-
throw e;
600+
} catch (Throwable t) {
601+
onError(SpannerExceptionFactory.asSpannerException(t), builder.getTransaction().hasBegin());
602+
throw t;
594603
}
595604
}
596605

@@ -621,6 +630,12 @@ public Long apply(ResultSet input) {
621630
ErrorCode.INVALID_ARGUMENT,
622631
"DML response missing stats possibly due to non-DML statement as input");
623632
}
633+
if (builder.getTransaction().hasBegin()
634+
&& !(input.getMetadata().hasTransaction()
635+
&& input.getMetadata().getTransaction().getId() != ByteString.EMPTY)) {
636+
throw SpannerExceptionFactory.newSpannerException(
637+
ErrorCode.FAILED_PRECONDITION, NO_TRANSACTION_RETURNED_MSG);
638+
}
624639
// For standard DML, using the exact row count.
625640
return input.getStats().getRowCountExact();
626641
}
@@ -633,8 +648,8 @@ public Long apply(ResultSet input) {
633648
new ApiFunction<Throwable, Long>() {
634649
@Override
635650
public Long apply(Throwable input) {
636-
SpannerException e = SpannerExceptionFactory.newSpannerException(input);
637-
onError(e, builder.hasTransaction() && builder.getTransaction().hasBegin());
651+
SpannerException e = SpannerExceptionFactory.asSpannerException(input);
652+
onError(e, builder.getTransaction().hasBegin());
638653
throw e;
639654
}
640655
},
@@ -645,9 +660,11 @@ public Long apply(Throwable input) {
645660
public void run() {
646661
try {
647662
if (resultSet.get().getMetadata().hasTransaction()) {
648-
onTransactionMetadata(resultSet.get().getMetadata().getTransaction());
663+
onTransactionMetadata(
664+
resultSet.get().getMetadata().getTransaction(),
665+
builder.getTransaction().hasBegin());
649666
}
650-
} catch (ExecutionException | InterruptedException e) {
667+
} catch (Throwable e) {
651668
// Ignore this error here as it is handled by the future that is returned by the
652669
// executeUpdateAsync method.
653670
}
@@ -670,7 +687,9 @@ public long[] batchUpdate(Iterable<Statement> statements, UpdateOption... option
670687
for (int i = 0; i < response.getResultSetsCount(); ++i) {
671688
results[i] = response.getResultSets(i).getStats().getRowCountExact();
672689
if (response.getResultSets(i).getMetadata().hasTransaction()) {
673-
onTransactionMetadata(response.getResultSets(i).getMetadata().getTransaction());
690+
onTransactionMetadata(
691+
response.getResultSets(i).getMetadata().getTransaction(),
692+
builder.getTransaction().hasBegin());
674693
}
675694
}
676695

@@ -686,8 +705,8 @@ public long[] batchUpdate(Iterable<Statement> statements, UpdateOption... option
686705
results);
687706
}
688707
return results;
689-
} catch (SpannerException e) {
690-
onError(e, builder.hasTransaction() && builder.getTransaction().hasBegin());
708+
} catch (Throwable e) {
709+
onError(SpannerExceptionFactory.asSpannerException(e), builder.getTransaction().hasBegin());
691710
throw e;
692711
}
693712
}
@@ -718,7 +737,9 @@ public long[] apply(ExecuteBatchDmlResponse input) {
718737
for (int i = 0; i < input.getResultSetsCount(); ++i) {
719738
results[i] = input.getResultSets(i).getStats().getRowCountExact();
720739
if (input.getResultSets(i).getMetadata().hasTransaction()) {
721-
onTransactionMetadata(input.getResultSets(i).getMetadata().getTransaction());
740+
onTransactionMetadata(
741+
input.getResultSets(i).getMetadata().getTransaction(),
742+
builder.getTransaction().hasBegin());
722743
}
723744
}
724745
// If one of the DML statements was aborted, we should throw an aborted exception.
@@ -743,10 +764,8 @@ public long[] apply(ExecuteBatchDmlResponse input) {
743764
new ApiFunction<Throwable, long[]>() {
744765
@Override
745766
public long[] apply(Throwable input) {
746-
SpannerException e = SpannerExceptionFactory.newSpannerException(input);
747-
onError(
748-
SpannerExceptionFactory.newSpannerException(e.getCause()),
749-
builder.hasTransaction() && builder.getTransaction().hasBegin());
767+
SpannerException e = SpannerExceptionFactory.asSpannerException(input);
768+
onError(e, builder.getTransaction().hasBegin());
750769
throw e;
751770
}
752771
},

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,8 @@ public class GrpcResultSetTest {
5656

5757
private static class NoOpListener implements AbstractResultSet.Listener {
5858
@Override
59-
public void onTransactionMetadata(Transaction transaction) throws SpannerException {}
59+
public void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId)
60+
throws SpannerException {}
6061

6162
@Override
6263
public void onError(SpannerException e, boolean withBeginTransaction) {}

0 commit comments

Comments
 (0)