Skip to content

Commit 2e1a5fd

Browse files
authored
[Bug] Use bigquery default retryable exceptions (#1431)
* replace our custom list of retryable exceptions with BigQuery's defaults * remove BadRequest as a retryable error
1 parent e1b6e74 commit 2e1a5fd

File tree

3 files changed

+12
-19
lines changed

3 files changed

+12
-19
lines changed
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
kind: Fixes
2+
body: Fix retry scenarios so that dbt always retries when BigQuery recommends a retry
3+
time: 2024-12-11T14:47:52.36905-05:00
4+
custom:
5+
Author: mikealfare
6+
Issue: "263"

dbt/adapters/bigquery/retry.py

Lines changed: 2 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,8 @@
11
from typing import Callable, Optional
22

3-
from google.api_core.exceptions import Forbidden
43
from google.api_core.future.polling import DEFAULT_POLLING
54
from google.api_core.retry import Retry
6-
from google.cloud.bigquery.retry import DEFAULT_RETRY
7-
from google.cloud.exceptions import BadGateway, BadRequest, ServerError
5+
from google.cloud.bigquery.retry import DEFAULT_RETRY, _job_should_retry
86
from requests.exceptions import ConnectionError
97

108
from dbt.adapters.contracts.connection import Connection, ConnectionState
@@ -83,7 +81,7 @@ def __call__(self, error: Exception) -> bool:
8381
self._error_count += 1
8482

8583
# if the error is retryable, and we haven't breached the threshold, log and continue
86-
if _is_retryable(error) and self._error_count <= self._retries:
84+
if _job_should_retry(error) and self._error_count <= self._retries:
8785
_logger.debug(
8886
f"Retry attempt {self._error_count} of {self._retries} after error: {repr(error)}"
8987
)
@@ -113,16 +111,3 @@ def on_error(error: Exception):
113111
raise FailedToConnectError(str(e))
114112

115113
return on_error
116-
117-
118-
def _is_retryable(error: Exception) -> bool:
119-
"""Return true for errors that are unlikely to occur again if retried."""
120-
if isinstance(
121-
error, (BadGateway, BadRequest, ConnectionError, ConnectionResetError, ServerError)
122-
):
123-
return True
124-
elif isinstance(error, Forbidden) and any(
125-
e["reason"] == "rateLimitExceeded" for e in error.errors
126-
):
127-
return True
128-
return False

tests/unit/test_bigquery_connection_manager.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ def generate_connection_reset_error():
5353
assert new_mock_client is not self.mock_client
5454

5555
def test_is_retryable(self):
56-
_is_retryable = dbt.adapters.bigquery.retry._is_retryable
56+
_is_retryable = google.cloud.bigquery.retry._job_should_retry
5757
exceptions = dbt.adapters.bigquery.impl.google.cloud.exceptions
5858
internal_server_error = exceptions.InternalServerError("code broke")
5959
bad_request_error = exceptions.BadRequest("code broke")
@@ -65,7 +65,9 @@ def test_is_retryable(self):
6565
service_unavailable_error = exceptions.ServiceUnavailable("service is unavailable")
6666

6767
self.assertTrue(_is_retryable(internal_server_error))
68-
self.assertTrue(_is_retryable(bad_request_error))
68+
self.assertFalse(
69+
_is_retryable(bad_request_error)
70+
) # this was removed after initially being included
6971
self.assertTrue(_is_retryable(connection_error))
7072
self.assertFalse(_is_retryable(client_error))
7173
self.assertTrue(_is_retryable(rate_limit_error))

0 commit comments

Comments
 (0)