Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Backport 1.9.latest] [Bug] Use bigquery default retryable exceptions #1439

Merged
merged 1 commit into from
Jan 9, 2025
+12 −19
Merged

[Backport 1.9.latest] [Bug] Use bigquery default retryable exceptions #1439

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions .changes/unreleased/Fixes-20241211-144752.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
kind: Fixes
body: Fix retry scenarios so that dbt always retries when BigQuery recommends a retry
time: 2024-12-11T14:47:52.36905-05:00
custom:
Author: mikealfare
Issue: "263"
19 changes: 2 additions & 17 deletions dbt/adapters/bigquery/retry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,8 @@
from typing import Callable, Optional

from google.api_core.exceptions import Forbidden
from google.api_core.future.polling import DEFAULT_POLLING
from google.api_core.retry import Retry
from google.cloud.bigquery.retry import DEFAULT_RETRY
from google.cloud.exceptions import BadGateway, BadRequest, ServerError
from google.cloud.bigquery.retry import DEFAULT_RETRY, _job_should_retry
from requests.exceptions import ConnectionError

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

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

return on_error


def _is_retryable(error: Exception) -> bool:
"""Return true for errors that are unlikely to occur again if retried."""
if isinstance(
error, (BadGateway, BadRequest, ConnectionError, ConnectionResetError, ServerError)
):
return True
elif isinstance(error, Forbidden) and any(
e["reason"] == "rateLimitExceeded" for e in error.errors
):
return True
return False
6 changes: 4 additions & 2 deletions tests/unit/test_bigquery_connection_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ def generate_connection_reset_error():
assert new_mock_client is not self.mock_client

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

self.assertTrue(_is_retryable(internal_server_error))
self.assertTrue(_is_retryable(bad_request_error))
self.assertFalse(
_is_retryable(bad_request_error)
) # this was removed after initially being included
self.assertTrue(_is_retryable(connection_error))
self.assertFalse(_is_retryable(client_error))
self.assertTrue(_is_retryable(rate_limit_error))
Expand Down
Loading