Skip to content
This repository was archived by the owner on May 8, 2026. It is now read-only.

Commit bc5befe

Browse files
alexoneillalexoneill
committed
feat: Handle StatusRuntimeException in CbtTestProxy
1 parent 4a99a8c commit bc5befe

1 file changed

Lines changed: 122 additions & 54 deletions

File tree

‎test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java‎

Lines changed: 122 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,17 @@ public void mutateRow(
300300
.build());
301301
responseObserver.onCompleted();
302302
return;
303+
} catch (StatusRuntimeException e) {
304+
responseObserver.onNext(
305+
MutateRowResult.newBuilder()
306+
.setStatus(
307+
com.google.rpc.Status.newBuilder()
308+
.setCode(e.getStatus().getCode().value())
309+
.setMessage(e.getStatus().getDescription())
310+
.build())
311+
.build());
312+
responseObserver.onCompleted();
313+
return;
303314
}
304315

305316
responseObserver.onNext(
@@ -354,6 +365,17 @@ public void bulkMutateRows(
354365
.build());
355366
responseObserver.onCompleted();
356367
return;
368+
} catch (StatusRuntimeException e) {
369+
responseObserver.onNext(
370+
MutateRowsResult.newBuilder()
371+
.setStatus(
372+
com.google.rpc.Status.newBuilder()
373+
.setCode(e.getStatus().getCode().value())
374+
.setMessage(e.getStatus().getDescription())
375+
.build())
376+
.build());
377+
responseObserver.onCompleted();
378+
return;
357379
}
358380

359381
responseObserver.onNext(
@@ -388,6 +410,15 @@ public void readRow(ReadRowRequest request, StreamObserver<RowResult> responseOb
388410
client
389411
.dataClient()
390412
.readRow(tableId, request.getRowKey(), FILTERS.fromProto(request.getFilter()));
413+
if (row != null) {
414+
RowResult.Builder resultBuilder = convertRowResult(row);
415+
responseObserver.onNext(
416+
resultBuilder.setStatus(com.google.rpc.Status.getDefaultInstance()).build());
417+
} else {
418+
logger.info(String.format("readRow() did not find row: %s", request.getRowKey()));
419+
responseObserver.onNext(
420+
RowResult.newBuilder().setStatus(com.google.rpc.Status.getDefaultInstance()).build());
421+
}
391422
} catch (ApiException e) {
392423
responseObserver.onNext(
393424
RowResult.newBuilder()
@@ -399,30 +430,29 @@ public void readRow(ReadRowRequest request, StreamObserver<RowResult> responseOb
399430
.build());
400431
responseObserver.onCompleted();
401432
return;
402-
}
403-
404-
if (row != null) {
405-
try {
406-
RowResult.Builder resultBuilder = convertRowResult(row);
407-
responseObserver.onNext(
408-
resultBuilder.setStatus(com.google.rpc.Status.getDefaultInstance()).build());
409-
} catch (RuntimeException e) {
410-
// If client encounters problem, don't return any row result.
411-
responseObserver.onNext(
412-
RowResult.newBuilder()
413-
.setStatus(
414-
com.google.rpc.Status.newBuilder()
415-
.setCode(Code.INTERNAL.getNumber())
416-
.setMessage(e.getMessage())
417-
.build())
418-
.build());
419-
responseObserver.onCompleted();
420-
return;
421-
}
422-
} else {
423-
logger.info(String.format("readRow() did not find row: %s", request.getRowKey()));
433+
} catch (StatusRuntimeException e) {
424434
responseObserver.onNext(
425-
RowResult.newBuilder().setStatus(com.google.rpc.Status.getDefaultInstance()).build());
435+
RowResult.newBuilder()
436+
.setStatus(
437+
com.google.rpc.Status.newBuilder()
438+
.setCode(e.getStatus().getCode().value())
439+
.setMessage(e.getStatus().getDescription())
440+
.build())
441+
.build());
442+
responseObserver.onCompleted();
443+
return;
444+
} catch (RuntimeException e) {
445+
// If client encounters problem, don't return any row result.
446+
responseObserver.onNext(
447+
RowResult.newBuilder()
448+
.setStatus(
449+
com.google.rpc.Status.newBuilder()
450+
.setCode(Code.INTERNAL.getNumber())
451+
.setMessage(e.getMessage())
452+
.build())
453+
.build());
454+
responseObserver.onCompleted();
455+
return;
426456
}
427457
responseObserver.onCompleted();
428458
}
@@ -441,6 +471,10 @@ public void readRows(ReadRowsRequest request, StreamObserver<RowsResult> respons
441471
Query query = Query.fromProto(request.getRequest());
442472
try {
443473
rows = client.dataClient().readRows(query);
474+
int cancelAfterRows = request.getCancelAfterRows();
475+
RowsResult.Builder resultBuilder = convertRowsResult(rows, cancelAfterRows);
476+
responseObserver.onNext(
477+
resultBuilder.setStatus(com.google.rpc.Status.getDefaultInstance()).build());
444478
} catch (ApiException e) {
445479
responseObserver.onNext(
446480
RowsResult.newBuilder()
@@ -452,13 +486,17 @@ public void readRows(ReadRowsRequest request, StreamObserver<RowsResult> respons
452486
.build());
453487
responseObserver.onCompleted();
454488
return;
455-
}
456-
457-
int cancelAfterRows = request.getCancelAfterRows();
458-
try {
459-
RowsResult.Builder resultBuilder = convertRowsResult(rows, cancelAfterRows);
489+
} catch (StatusRuntimeException e) {
460490
responseObserver.onNext(
461-
resultBuilder.setStatus(com.google.rpc.Status.getDefaultInstance()).build());
491+
RowsResult.newBuilder()
492+
.setStatus(
493+
com.google.rpc.Status.newBuilder()
494+
.setCode(e.getStatus().getCode().value())
495+
.setMessage(e.getStatus().getDescription())
496+
.build())
497+
.build());
498+
responseObserver.onCompleted();
499+
return;
462500
} catch (RuntimeException e) {
463501
// If client encounters problem, don't return any row result.
464502
responseObserver.onNext(
@@ -578,6 +616,17 @@ public void sampleRowKeys(
578616
.build());
579617
responseObserver.onCompleted();
580618
return;
619+
} catch (StatusRuntimeException e) {
620+
responseObserver.onNext(
621+
SampleRowKeysResult.newBuilder()
622+
.setStatus(
623+
com.google.rpc.Status.newBuilder()
624+
.setCode(e.getStatus().getCode().value())
625+
.setMessage(e.getStatus().getDescription())
626+
.build())
627+
.build());
628+
responseObserver.onCompleted();
629+
return;
581630
}
582631

583632
SampleRowKeysResult.Builder resultBuilder = SampleRowKeysResult.newBuilder();
@@ -618,6 +667,17 @@ public void checkAndMutateRow(
618667
.build());
619668
responseObserver.onCompleted();
620669
return;
670+
} catch (StatusRuntimeException e) {
671+
responseObserver.onNext(
672+
CheckAndMutateRowResult.newBuilder()
673+
.setStatus(
674+
com.google.rpc.Status.newBuilder()
675+
.setCode(e.getStatus().getCode().value())
676+
.setMessage(e.getStatus().getDescription())
677+
.build())
678+
.build());
679+
responseObserver.onCompleted();
680+
return;
621681
}
622682

623683
CheckAndMutateRowResult.Builder resultBuilder = CheckAndMutateRowResult.newBuilder();
@@ -642,6 +702,17 @@ public void readModifyWriteRow(
642702
ReadModifyWriteRow mutation = ReadModifyWriteRow.fromProto(request.getRequest());
643703
try {
644704
row = client.dataClient().readModifyWriteRow(mutation);
705+
if (row != null) {
706+
RowResult.Builder resultBuilder = convertRowResult(row);
707+
responseObserver.onNext(
708+
resultBuilder.setStatus(com.google.rpc.Status.getDefaultInstance()).build());
709+
} else {
710+
logger.info(
711+
String.format(
712+
"readModifyWriteRow() did not find row: %s", request.getRequest().getRowKey()));
713+
responseObserver.onNext(
714+
RowResult.newBuilder().setStatus(com.google.rpc.Status.getDefaultInstance()).build());
715+
}
645716
} catch (ApiException e) {
646717
responseObserver.onNext(
647718
RowResult.newBuilder()
@@ -653,32 +724,29 @@ public void readModifyWriteRow(
653724
.build());
654725
responseObserver.onCompleted();
655726
return;
656-
}
657-
658-
if (row != null) {
659-
try {
660-
RowResult.Builder resultBuilder = convertRowResult(row);
661-
responseObserver.onNext(
662-
resultBuilder.setStatus(com.google.rpc.Status.getDefaultInstance()).build());
663-
} catch (RuntimeException e) {
664-
// If client encounters problem, fail the whole operation.
665-
responseObserver.onNext(
666-
RowResult.newBuilder()
667-
.setStatus(
668-
com.google.rpc.Status.newBuilder()
669-
.setCode(Code.INTERNAL.getNumber())
670-
.setMessage(e.getMessage())
671-
.build())
672-
.build());
673-
responseObserver.onCompleted();
674-
return;
675-
}
676-
} else {
677-
logger.info(
678-
String.format(
679-
"readModifyWriteRow() did not find row: %s", request.getRequest().getRowKey()));
727+
} catch (StatusRuntimeException e) {
728+
responseObserver.onNext(
729+
RowResult.newBuilder()
730+
.setStatus(
731+
com.google.rpc.Status.newBuilder()
732+
.setCode(e.getStatus().getCode().value())
733+
.setMessage(e.getStatus().getDescription())
734+
.build())
735+
.build());
736+
responseObserver.onCompleted();
737+
return;
738+
} catch (RuntimeException e) {
739+
// If client encounters problem, fail the whole operation.
680740
responseObserver.onNext(
681-
RowResult.newBuilder().setStatus(com.google.rpc.Status.getDefaultInstance()).build());
741+
RowResult.newBuilder()
742+
.setStatus(
743+
com.google.rpc.Status.newBuilder()
744+
.setCode(Code.INTERNAL.getNumber())
745+
.setMessage(e.getMessage())
746+
.build())
747+
.build());
748+
responseObserver.onCompleted();
749+
return;
682750
}
683751
responseObserver.onCompleted();
684752
}

0 commit comments

Comments
 (0)