feat: Add ability to specify RetryOptions and BigQueryRetryConfig whe… · googleapis/java-bigquery@1f91ae7
@@ -27,6 +27,7 @@
2727import com.google.api.services.bigquery.model.*;
2828import com.google.api.services.bigquery.model.JobStatistics;
2929import com.google.cloud.Policy;
30+import com.google.cloud.RetryOption;
3031import com.google.cloud.ServiceOptions;
3132import com.google.cloud.Tuple;
3233import com.google.cloud.bigquery.BigQuery.JobOption;
@@ -1594,6 +1595,119 @@ public void testCreateJobFailureShouldRetry() {
15941595verify(bigqueryRpcMock, times(6)).create(jobCapture.capture(), eq(EMPTY_RPC_OPTIONS));
15951596 }
159615971598+@Test
1599+public void testCreateJobWithBigQueryRetryConfigFailureShouldRetry() {
1600+// Validate create job with BigQueryRetryConfig that retries on rate limit error message.
1601+JobOption bigQueryRetryConfigOption =
1602+JobOption.bigQueryRetryConfig(
1603+BigQueryRetryConfig.newBuilder()
1604+ .retryOnMessage(BigQueryErrorMessages.RATE_LIMIT_EXCEEDED_MSG)
1605+ .retryOnMessage(BigQueryErrorMessages.JOB_RATE_LIMIT_EXCEEDED_MSG)
1606+ .retryOnRegEx(BigQueryErrorMessages.RetryRegExPatterns.RATE_LIMIT_EXCEEDED_REGEX)
1607+ .build());
1608+1609+Map<BigQueryRpc.Option, ?> bigQueryRpcOptions = optionMap(bigQueryRetryConfigOption);
1610+when(bigqueryRpcMock.create(jobCapture.capture(), eq(bigQueryRpcOptions)))
1611+ .thenThrow(
1612+new BigQueryException(
1613+400, RATE_LIMIT_ERROR_MSG)) // retrial on based on RATE_LIMIT_EXCEEDED_MSG
1614+ .thenThrow(new BigQueryException(200, RATE_LIMIT_ERROR_MSG))
1615+ .thenReturn(newJobPb());
1616+1617+bigquery = options.getService();
1618+bigquery =
1619+options
1620+ .toBuilder()
1621+ .setRetrySettings(ServiceOptions.getDefaultRetrySettings())
1622+ .build()
1623+ .getService();
1624+1625+ ((BigQueryImpl) bigquery)
1626+ .create(JobInfo.of(QUERY_JOB_CONFIGURATION_FOR_DMLQUERY), bigQueryRetryConfigOption);
1627+verify(bigqueryRpcMock, times(3)).create(jobCapture.capture(), eq(bigQueryRpcOptions));
1628+ }
1629+1630+@Test
1631+public void testCreateJobWithBigQueryRetryConfigFailureShouldNotRetry() {
1632+// Validate create job with BigQueryRetryConfig that does not retry on rate limit error message.
1633+JobOption bigQueryRetryConfigOption =
1634+JobOption.bigQueryRetryConfig(BigQueryRetryConfig.newBuilder().build());
1635+1636+Map<BigQueryRpc.Option, ?> bigQueryRpcOptions = optionMap(bigQueryRetryConfigOption);
1637+when(bigqueryRpcMock.create(jobCapture.capture(), eq(bigQueryRpcOptions)))
1638+ .thenThrow(new BigQueryException(400, RATE_LIMIT_ERROR_MSG));
1639+1640+bigquery = options.getService();
1641+bigquery =
1642+options
1643+ .toBuilder()
1644+ .setRetrySettings(ServiceOptions.getDefaultRetrySettings())
1645+ .build()
1646+ .getService();
1647+1648+try {
1649+ ((BigQueryImpl) bigquery)
1650+ .create(JobInfo.of(QUERY_JOB_CONFIGURATION_FOR_DMLQUERY), bigQueryRetryConfigOption);
1651+fail("JobException expected");
1652+ } catch (BigQueryException e) {
1653+assertNotNull(e.getMessage());
1654+ }
1655+// Verify that getQueryResults is attempted only once and not retried since the error message
1656+// does not match.
1657+verify(bigqueryRpcMock, times(1)).create(jobCapture.capture(), eq(bigQueryRpcOptions));
1658+ }
1659+1660+@Test
1661+public void testCreateJobWithRetryOptionsFailureShouldRetry() {
1662+// Validate create job with RetryOptions.
1663+JobOption retryOptions = JobOption.retryOptions(RetryOption.maxAttempts(4));
1664+Map<BigQueryRpc.Option, ?> bigQueryRpcOptions = optionMap(retryOptions);
1665+when(bigqueryRpcMock.create(jobCapture.capture(), eq(bigQueryRpcOptions)))
1666+ .thenThrow(new BigQueryException(500, "InternalError"))
1667+ .thenThrow(new BigQueryException(502, "Bad Gateway"))
1668+ .thenThrow(new BigQueryException(503, "Service Unavailable"))
1669+ .thenReturn(newJobPb());
1670+1671+bigquery = options.getService();
1672+bigquery =
1673+options
1674+ .toBuilder()
1675+ .setRetrySettings(ServiceOptions.getDefaultRetrySettings())
1676+ .build()
1677+ .getService();
1678+1679+ ((BigQueryImpl) bigquery)
1680+ .create(JobInfo.of(QUERY_JOB_CONFIGURATION_FOR_DMLQUERY), retryOptions);
1681+verify(bigqueryRpcMock, times(4)).create(jobCapture.capture(), eq(bigQueryRpcOptions));
1682+ }
1683+1684+@Test
1685+public void testCreateJobWithRetryOptionsFailureShouldNotRetry() {
1686+// Validate create job with RetryOptions that only attempts once (no retry).
1687+JobOption retryOptions = JobOption.retryOptions(RetryOption.maxAttempts(1));
1688+Map<BigQueryRpc.Option, ?> bigQueryRpcOptions = optionMap(retryOptions);
1689+when(bigqueryRpcMock.create(jobCapture.capture(), eq(bigQueryRpcOptions)))
1690+ .thenThrow(new BigQueryException(500, "InternalError"))
1691+ .thenReturn(newJobPb());
1692+1693+bigquery = options.getService();
1694+bigquery =
1695+options
1696+ .toBuilder()
1697+ .setRetrySettings(ServiceOptions.getDefaultRetrySettings())
1698+ .build()
1699+ .getService();
1700+1701+try {
1702+ ((BigQueryImpl) bigquery)
1703+ .create(JobInfo.of(QUERY_JOB_CONFIGURATION_FOR_DMLQUERY), retryOptions);
1704+fail("JobException expected");
1705+ } catch (BigQueryException e) {
1706+assertNotNull(e.getMessage());
1707+ }
1708+verify(bigqueryRpcMock, times(1)).create(jobCapture.capture(), eq(bigQueryRpcOptions));
1709+ }
1710+15971711@Test
15981712public void testCreateJobWithSelectedFields() {
15991713when(bigqueryRpcMock.create(