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

[Bug] Use bigquery default retryable exceptions #1431

Merged
merged 5 commits into from
Dec 13, 2024
Merged

[Bug] Use bigquery default retryable exceptions #1431

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
89766ae
changelog
mikealfare Dec 11, 2024
5900683
replace our custom list of retryable exceptions with BigQuery's defau…
mikealfare Dec 11, 2024
b96493a
remove unused imports
mikealfare Dec 11, 2024
0d4994c
remove BadRequest as a retryable error
mikealfare Dec 11, 2024
a17b060
code quality checks
mikealfare Dec 11, 2024
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
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

While this is a private function, there is no reasonable way to capture BQ's defaults here. We need to add an exception into list that goes into this function, which would then require vendoring the whole function. Even if we remove BadRequest, we would then need to pull a private attribute (_predicate) off of DEFAULT_JOB_RETRY so that we can implement the retry cap (job_retries).

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:
mikealfare marked this conversation as resolved.
Show resolved Hide resolved
"""Return true for errors that are unlikely to occur again if retried."""
if isinstance(
error, (BadGateway, BadRequest, ConnectionError, ConnectionResetError, ServerError)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only BadRequest is missing from BQ's defaults; the rest can be removed since they would be caught in _job_should_retry.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We determined we can remove BadRequest here. The rationale is that we would expect a bad request to always be a bad request since the issue originates from the user side and not the BQ side.

):
return True
elif isinstance(error, Forbidden) and any(
e["reason"] == "rateLimitExceeded" for e in error.errors
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This condition is caught in the default retryable reasons (for more than just Forbidden).

):
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(
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's worth keeping this, but flipping it, to be clear that this was once supported, and no longer is supported.

_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