|
1 | 1 | from typing import Callable, Optional |
2 | 2 |
|
3 | | -from google.api_core.exceptions import Forbidden |
4 | 3 | from google.api_core.future.polling import DEFAULT_POLLING |
5 | 4 | 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 |
8 | 6 | from requests.exceptions import ConnectionError |
9 | 7 |
|
10 | 8 | from dbt.adapters.contracts.connection import Connection, ConnectionState |
@@ -83,7 +81,7 @@ def __call__(self, error: Exception) -> bool: |
83 | 81 | self._error_count += 1 |
84 | 82 |
|
85 | 83 | # 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: |
87 | 85 | _logger.debug( |
88 | 86 | f"Retry attempt {self._error_count} of {self._retries} after error: {repr(error)}" |
89 | 87 | ) |
@@ -113,16 +111,3 @@ def on_error(error: Exception): |
113 | 111 | raise FailedToConnectError(str(e)) |
114 | 112 |
|
115 | 113 | 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 |
0 commit comments