Skip to content

Commit de0914d

Browse files
whuffman36whuffman36
authored
feat(bigquery): Add custom ExceptionHandler to BigQueryOptions (#3937)
* feat(bigquery): Add custom ExceptionHandler to BigQueryOptions * add default project id to tests * Add dummy project id to tests * Remove defaultInstance() from test * add option to specify retry algo instead of just abortOn/retryOn * fix formatting * remove custom exception handler from bq options
1 parent 7fce427 commit de0914d

File tree

6 files changed

+148
-40
lines changed

6 files changed

+148
-40
lines changed

‎google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryBaseService.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ protected BigQueryBaseService(ServiceOptions options) {
2626
super(options);
2727
}
2828

29-
public static final ExceptionHandler BIGQUERY_EXCEPTION_HANDLER =
29+
public static final ExceptionHandler DEFAULT_BIGQUERY_EXCEPTION_HANDLER =
3030
ExceptionHandler.newBuilder()
3131
.abortOn(RuntimeException.class)
3232
.retryOn(java.net.ConnectException.class) // retry on Connection Exception

‎google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java‎

Lines changed: 31 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -293,7 +293,7 @@ public com.google.api.services.bigquery.model.Dataset call() throws IOException
293293
}
294294
},
295295
getOptions().getRetrySettings(),
296-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
296+
getOptions().getResultRetryAlgorithm(),
297297
getOptions().getClock(),
298298
EMPTY_RETRY_CONFIG,
299299
getOptions().isOpenTelemetryTracingEnabled(),
@@ -340,7 +340,7 @@ public com.google.api.services.bigquery.model.Table call() throws IOException {
340340
}
341341
},
342342
getOptions().getRetrySettings(),
343-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
343+
getOptions().getResultRetryAlgorithm(),
344344
getOptions().getClock(),
345345
EMPTY_RETRY_CONFIG,
346346
getOptions().isOpenTelemetryTracingEnabled(),
@@ -396,7 +396,7 @@ public com.google.api.services.bigquery.model.Routine call() throws IOException
396396
}
397397
},
398398
getOptions().getRetrySettings(),
399-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
399+
getOptions().getResultRetryAlgorithm(),
400400
getOptions().getClock(),
401401
EMPTY_RETRY_CONFIG,
402402
getOptions().isOpenTelemetryTracingEnabled(),
@@ -490,7 +490,7 @@ public com.google.api.services.bigquery.model.Job call() throws IOException {
490490
? RetryOption.mergeToSettings(
491491
getOptions().getRetrySettings(), getRetryOptions(optionsMap))
492492
: getOptions().getRetrySettings(),
493-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
493+
getOptions().getResultRetryAlgorithm(),
494494
getOptions().getClock(),
495495
getBigQueryRetryConfig(optionsMap) != null
496496
? getBigQueryRetryConfig(optionsMap)
@@ -586,7 +586,7 @@ public com.google.api.services.bigquery.model.Dataset call() throws IOException
586586
}
587587
},
588588
getOptions().getRetrySettings(),
589-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
589+
getOptions().getResultRetryAlgorithm(),
590590
getOptions().getClock(),
591591
EMPTY_RETRY_CONFIG,
592592
getOptions().isOpenTelemetryTracingEnabled(),
@@ -652,7 +652,7 @@ private static Page<Dataset> listDatasets(
652652
}
653653
},
654654
serviceOptions.getRetrySettings(),
655-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
655+
serviceOptions.getResultRetryAlgorithm(),
656656
serviceOptions.getClock(),
657657
EMPTY_RETRY_CONFIG,
658658
serviceOptions.isOpenTelemetryTracingEnabled(),
@@ -704,7 +704,7 @@ public Boolean call() throws IOException {
704704
}
705705
},
706706
getOptions().getRetrySettings(),
707-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
707+
getOptions().getResultRetryAlgorithm(),
708708
getOptions().getClock(),
709709
EMPTY_RETRY_CONFIG,
710710
getOptions().isOpenTelemetryTracingEnabled(),
@@ -755,7 +755,7 @@ public Boolean call() throws IOException {
755755
}
756756
},
757757
getOptions().getRetrySettings(),
758-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
758+
getOptions().getResultRetryAlgorithm(),
759759
getOptions().getClock(),
760760
EMPTY_RETRY_CONFIG,
761761
getOptions().isOpenTelemetryTracingEnabled(),
@@ -801,7 +801,7 @@ public Boolean call() throws IOException {
801801
}
802802
},
803803
getOptions().getRetrySettings(),
804-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
804+
getOptions().getResultRetryAlgorithm(),
805805
getOptions().getClock(),
806806
EMPTY_RETRY_CONFIG,
807807
getOptions().isOpenTelemetryTracingEnabled(),
@@ -847,7 +847,7 @@ public Boolean call() throws IOException {
847847
}
848848
},
849849
getOptions().getRetrySettings(),
850-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
850+
getOptions().getResultRetryAlgorithm(),
851851
getOptions().getClock(),
852852
EMPTY_RETRY_CONFIG,
853853
getOptions().isOpenTelemetryTracingEnabled(),
@@ -891,7 +891,7 @@ public Boolean call() throws IOException {
891891
}
892892
},
893893
getOptions().getRetrySettings(),
894-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
894+
getOptions().getResultRetryAlgorithm(),
895895
getOptions().getClock(),
896896
EMPTY_RETRY_CONFIG,
897897
getOptions().isOpenTelemetryTracingEnabled(),
@@ -932,7 +932,7 @@ public com.google.api.services.bigquery.model.Dataset call() throws IOException
932932
}
933933
},
934934
getOptions().getRetrySettings(),
935-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
935+
getOptions().getResultRetryAlgorithm(),
936936
getOptions().getClock(),
937937
EMPTY_RETRY_CONFIG,
938938
getOptions().isOpenTelemetryTracingEnabled(),
@@ -979,7 +979,7 @@ public com.google.api.services.bigquery.model.Table call() throws IOException {
979979
}
980980
},
981981
getOptions().getRetrySettings(),
982-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
982+
getOptions().getResultRetryAlgorithm(),
983983
getOptions().getClock(),
984984
EMPTY_RETRY_CONFIG,
985985
getOptions().isOpenTelemetryTracingEnabled(),
@@ -1025,7 +1025,7 @@ public com.google.api.services.bigquery.model.Model call() throws IOException {
10251025
}
10261026
},
10271027
getOptions().getRetrySettings(),
1028-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1028+
getOptions().getResultRetryAlgorithm(),
10291029
getOptions().getClock(),
10301030
EMPTY_RETRY_CONFIG,
10311031
getOptions().isOpenTelemetryTracingEnabled(),
@@ -1071,7 +1071,7 @@ public com.google.api.services.bigquery.model.Routine call() throws IOException
10711071
}
10721072
},
10731073
getOptions().getRetrySettings(),
1074-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1074+
getOptions().getResultRetryAlgorithm(),
10751075
getOptions().getClock(),
10761076
EMPTY_RETRY_CONFIG,
10771077
getOptions().isOpenTelemetryTracingEnabled(),
@@ -1125,7 +1125,7 @@ public com.google.api.services.bigquery.model.Table call() throws IOException {
11251125
}
11261126
},
11271127
getOptions().getRetrySettings(),
1128-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1128+
getOptions().getResultRetryAlgorithm(),
11291129
getOptions().getClock(),
11301130
EMPTY_RETRY_CONFIG,
11311131
getOptions().isOpenTelemetryTracingEnabled(),
@@ -1184,7 +1184,7 @@ public com.google.api.services.bigquery.model.Model call() throws IOException {
11841184
}
11851185
},
11861186
getOptions().getRetrySettings(),
1187-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1187+
getOptions().getResultRetryAlgorithm(),
11881188
getOptions().getClock(),
11891189
EMPTY_RETRY_CONFIG,
11901190
getOptions().isOpenTelemetryTracingEnabled(),
@@ -1243,7 +1243,7 @@ public com.google.api.services.bigquery.model.Routine call() throws IOException
12431243
}
12441244
},
12451245
getOptions().getRetrySettings(),
1246-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1246+
getOptions().getResultRetryAlgorithm(),
12471247
getOptions().getClock(),
12481248
EMPTY_RETRY_CONFIG,
12491249
getOptions().isOpenTelemetryTracingEnabled(),
@@ -1461,7 +1461,7 @@ public Tuple<String, Iterable<com.google.api.services.bigquery.model.Table>> cal
14611461
}
14621462
},
14631463
serviceOptions.getRetrySettings(),
1464-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1464+
serviceOptions.getResultRetryAlgorithm(),
14651465
serviceOptions.getClock(),
14661466
EMPTY_RETRY_CONFIG,
14671467
serviceOptions.isOpenTelemetryTracingEnabled(),
@@ -1502,7 +1502,7 @@ public Tuple<String, Iterable<com.google.api.services.bigquery.model.Model>> cal
15021502
}
15031503
},
15041504
serviceOptions.getRetrySettings(),
1505-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1505+
serviceOptions.getResultRetryAlgorithm(),
15061506
serviceOptions.getClock(),
15071507
EMPTY_RETRY_CONFIG,
15081508
serviceOptions.isOpenTelemetryTracingEnabled(),
@@ -1543,7 +1543,7 @@ private static Page<Routine> listRoutines(
15431543
}
15441544
},
15451545
serviceOptions.getRetrySettings(),
1546-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1546+
serviceOptions.getResultRetryAlgorithm(),
15471547
serviceOptions.getClock(),
15481548
EMPTY_RETRY_CONFIG,
15491549
serviceOptions.isOpenTelemetryTracingEnabled(),
@@ -1625,7 +1625,7 @@ public TableDataInsertAllResponse call() throws Exception {
16251625
}
16261626
},
16271627
getOptions().getRetrySettings(),
1628-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1628+
getOptions().getResultRetryAlgorithm(),
16291629
getOptions().getClock(),
16301630
EMPTY_RETRY_CONFIG,
16311631
getOptions().isOpenTelemetryTracingEnabled(),
@@ -1719,7 +1719,7 @@ public TableDataList call() throws IOException {
17191719
}
17201720
},
17211721
serviceOptions.getRetrySettings(),
1722-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1722+
serviceOptions.getResultRetryAlgorithm(),
17231723
serviceOptions.getClock(),
17241724
EMPTY_RETRY_CONFIG,
17251725
serviceOptions.isOpenTelemetryTracingEnabled(),
@@ -1793,7 +1793,7 @@ public com.google.api.services.bigquery.model.Job call() throws IOException {
17931793
}
17941794
},
17951795
getOptions().getRetrySettings(),
1796-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1796+
getOptions().getResultRetryAlgorithm(),
17971797
getOptions().getClock(),
17981798
EMPTY_RETRY_CONFIG,
17991799
getOptions().isOpenTelemetryTracingEnabled(),
@@ -1850,7 +1850,7 @@ public Tuple<String, Iterable<com.google.api.services.bigquery.model.Job>> call(
18501850
}
18511851
},
18521852
serviceOptions.getRetrySettings(),
1853-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1853+
serviceOptions.getResultRetryAlgorithm(),
18541854
serviceOptions.getClock(),
18551855
EMPTY_RETRY_CONFIG,
18561856
serviceOptions.isOpenTelemetryTracingEnabled(),
@@ -1905,7 +1905,7 @@ public Boolean call() throws IOException {
19051905
}
19061906
},
19071907
getOptions().getRetrySettings(),
1908-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1908+
getOptions().getResultRetryAlgorithm(),
19091909
getOptions().getClock(),
19101910
EMPTY_RETRY_CONFIG,
19111911
getOptions().isOpenTelemetryTracingEnabled(),
@@ -1992,7 +1992,7 @@ public com.google.api.services.bigquery.model.QueryResponse call()
19921992
}
19931993
},
19941994
getOptions().getRetrySettings(),
1995-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
1995+
getOptions().getResultRetryAlgorithm(),
19961996
getOptions().getClock(),
19971997
DEFAULT_RETRY_CONFIG,
19981998
getOptions().isOpenTelemetryTracingEnabled(),
@@ -2169,7 +2169,7 @@ public GetQueryResultsResponse call() throws IOException {
21692169
}
21702170
},
21712171
serviceOptions.getRetrySettings(),
2172-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
2172+
serviceOptions.getResultRetryAlgorithm(),
21732173
serviceOptions.getClock(),
21742174
DEFAULT_RETRY_CONFIG,
21752175
serviceOptions.isOpenTelemetryTracingEnabled(),
@@ -2240,7 +2240,7 @@ public com.google.api.services.bigquery.model.Policy call() throws IOException {
22402240
}
22412241
},
22422242
getOptions().getRetrySettings(),
2243-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
2243+
getOptions().getResultRetryAlgorithm(),
22442244
getOptions().getClock(),
22452245
EMPTY_RETRY_CONFIG,
22462246
getOptions().isOpenTelemetryTracingEnabled(),
@@ -2286,7 +2286,7 @@ public com.google.api.services.bigquery.model.Policy call() throws IOException {
22862286
}
22872287
},
22882288
getOptions().getRetrySettings(),
2289-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
2289+
getOptions().getResultRetryAlgorithm(),
22902290
getOptions().getClock(),
22912291
EMPTY_RETRY_CONFIG,
22922292
getOptions().isOpenTelemetryTracingEnabled(),
@@ -2334,7 +2334,7 @@ public com.google.api.services.bigquery.model.TestIamPermissionsResponse call()
23342334
}
23352335
},
23362336
getOptions().getRetrySettings(),
2337-
BigQueryBaseService.BIGQUERY_EXCEPTION_HANDLER,
2337+
getOptions().getResultRetryAlgorithm(),
23382338
getOptions().getClock(),
23392339
EMPTY_RETRY_CONFIG,
23402340
getOptions().isOpenTelemetryTracingEnabled(),

‎google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryOptions.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package com.google.cloud.bigquery;
1818

1919
import com.google.api.core.BetaApi;
20+
import com.google.api.gax.retrying.ResultRetryAlgorithm;
2021
import com.google.cloud.ServiceDefaults;
2122
import com.google.cloud.ServiceOptions;
2223
import com.google.cloud.ServiceRpc;
@@ -43,6 +44,7 @@ public class BigQueryOptions extends ServiceOptions<BigQuery, BigQueryOptions> {
4344
private JobCreationMode defaultJobCreationMode = JobCreationMode.JOB_CREATION_MODE_UNSPECIFIED;
4445
private boolean enableOpenTelemetryTracing;
4546
private Tracer openTelemetryTracer;
47+
private ResultRetryAlgorithm<?> resultRetryAlgorithm;
4648

4749
public static class DefaultBigQueryFactory implements BigQueryFactory {
4850

@@ -70,6 +72,7 @@ public static class Builder extends ServiceOptions.Builder<BigQuery, BigQueryOpt
7072
private boolean useInt64Timestamps;
7173
private boolean enableOpenTelemetryTracing;
7274
private Tracer openTelemetryTracer;
75+
private ResultRetryAlgorithm<?> resultRetryAlgorithm;
7376

7477
private Builder() {}
7578

@@ -118,6 +121,11 @@ public Builder setOpenTelemetryTracer(Tracer tracer) {
118121
return this;
119122
}
120123

124+
public Builder setResultRetryAlgorithm(ResultRetryAlgorithm<?> resultRetryAlgorithm) {
125+
this.resultRetryAlgorithm = resultRetryAlgorithm;
126+
return this;
127+
}
128+
121129
@Override
122130
public BigQueryOptions build() {
123131
return new BigQueryOptions(this);
@@ -130,6 +138,11 @@ private BigQueryOptions(Builder builder) {
130138
this.useInt64Timestamps = builder.useInt64Timestamps;
131139
this.enableOpenTelemetryTracing = builder.enableOpenTelemetryTracing;
132140
this.openTelemetryTracer = builder.openTelemetryTracer;
141+
if (builder.resultRetryAlgorithm != null) {
142+
this.resultRetryAlgorithm = builder.resultRetryAlgorithm;
143+
} else {
144+
this.resultRetryAlgorithm = BigQueryBaseService.DEFAULT_BIGQUERY_EXCEPTION_HANDLER;
145+
}
133146
}
134147

135148
private static class BigQueryDefaults implements ServiceDefaults<BigQuery, BigQueryOptions> {
@@ -221,6 +234,10 @@ public Tracer getOpenTelemetryTracer() {
221234
return openTelemetryTracer;
222235
}
223236

237+
public ResultRetryAlgorithm<?> getResultRetryAlgorithm() {
238+
return resultRetryAlgorithm;
239+
}
240+
224241
@SuppressWarnings("unchecked")
225242
@Override
226243
public Builder toBuilder() {

0 commit comments

Comments
 (0)