Merge branch 'main' into additional-column-unit-tests

.changes/unreleased/Fixes-20241211-144752.yaml

@@ -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"

dbt/adapters/bigquery/retry.py

@@ -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
@@ -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)}"
             )
@@ -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

tests/unit/test_bigquery_connection_manager.py

@@ -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")
@@ -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))