Add retry of additional errors

CHANGELOG.md

@@ -4,6 +4,7 @@
 ### Features
 - Add support for impersonating a service account using `impersonate_service_account` in the BigQuery profile configuration ([#2677](https://github.com/fishtown-analytics/dbt/issues/2677)) ([docs](https://docs.getdbt.com/reference/warehouse-profiles/bigquery-profile#service-account-impersonation))
 - Macros in the current project can override internal dbt macros that are called through `execute_macros`. ([#2301](https://github.com/fishtown-analytics/dbt/issues/2301), [#2686](https://github.com/fishtown-analytics/dbt/pull/2686))
+- Add better retry support when using the BigQuery adapter ([#2694](https://github.com/fishtown-analytics/dbt/pull/2694), follow-up to [#1963](https://github.com/fishtown-analytics/dbt/pull/1963))
 
 
 ### Breaking changes
@@ -16,6 +17,7 @@
 
 Contributors:
 - [@bbhoss](https://github.com/bbhoss) ([#2677](https://github.com/fishtown-analytics/dbt/pull/2677))
+- [@kconvey](https://github.com/kconvey) ([#2694](https://github.com/fishtown-analytics/dbt/pull/2694))
 
 ## dbt 0.18.0b2 (July 30, 2020)
 

plugins/bigquery/dbt/adapters/bigquery/connections.py

@@ -1,5 +1,6 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
+from requests.exceptions import ConnectionError
 from typing import Optional, Any, Dict
 
 import google.auth
@@ -25,6 +26,18 @@
 
 BQ_QUERY_JOB_SPLIT = '-----Query Job SQL Follows-----'
 
+REOPENABLE_ERRORS = (
+ ConnectionResetError,
+ ConnectionError,
+)
+
+RETRYABLE_ERRORS = (
+ google.cloud.exceptions.ServerError,
+ google.cloud.exceptions.BadRequest,
+ ConnectionResetError,
+ ConnectionError,
+)
+
 
 class Priority(StrEnum):
  Interactive = 'interactive'
@@ -390,12 +403,20 @@ def _query_and_results(self, client, sql, conn, job_params, timeout=None):
 
  def _retry_and_handle(self, msg, conn, fn):
  """retry a function call within the context of exception_handler."""
+ def reopen_conn_on_error(error):
+ if isinstance(error, REOPENABLE_ERRORS):
+ logger.warning('Reopening connection after {!r}', error)
+ self.close(conn)
+ self.open(conn)
+ return
+
  with self.exception_handler(msg):
  return retry.retry_target(
  target=fn,
  predicate=_ErrorCounter(self.get_retries(conn)).count_error,
  sleep_generator=self._retry_generator(),
- deadline=None)
+ deadline=None,
+ on_error=reopen_conn_on_error)
 
  def _retry_generator(self):
  """Generates retry intervals that exponentially back off."""
@@ -425,5 +446,7 @@ def count_error(self, error):
 
 
 def _is_retryable(error):
- """Return true for 500 level (retryable) errors."""
- return isinstance(error, google.cloud.exceptions.ServerError)
+ """Return true for errors that are unlikely to occur again if retried."""
+ if isinstance(error, RETRYABLE_ERRORS):
+ return True
+ return False

test/unit/test_bigquery_adapter.py

@@ -3,6 +3,7 @@
 import re
 import unittest
 from contextlib import contextmanager
+from requests.exceptions import ConnectionError
 from unittest.mock import patch, MagicMock, Mock
 
 import hologram
@@ -427,9 +428,10 @@ def setUp(self):
 
  self.connections.get_thread_connection = lambda: self.mock_connection
 
- def test_retry_and_handle(self):
+ @patch(
+ 'dbt.adapters.bigquery.connections._is_retryable', return_value=True)
+ def test_retry_and_handle(self, is_retryable):
  self.connections.DEFAULT_MAXIMUM_DELAY = 2.0
- dbt.adapters.bigquery.connections._is_retryable = lambda x: True
 
  @contextmanager
  def dummy_handler(msg):
@@ -453,14 +455,42 @@ def raiseDummyException():
  raiseDummyException)
  self.assertEqual(DummyException.count, 9)
 
+ @patch(
+ 'dbt.adapters.bigquery.connections._is_retryable', return_value=True)
+ def test_retry_connection_reset(self, is_retryable):
+ self.connections.open = MagicMock()
+ self.connections.close = MagicMock()
+ self.connections.DEFAULT_MAXIMUM_DELAY = 2.0
+
+ @contextmanager
+ def dummy_handler(msg):
+ yield
+
+ self.connections.exception_handler = dummy_handler
+
+ def raiseConnectionResetError():
+ raise ConnectionResetError("Connection broke")
+
+ mock_conn = Mock(credentials=Mock(retries=1))
+ with self.assertRaises(ConnectionResetError):
+ self.connections._retry_and_handle(
+ "some sql", mock_conn,
+ raiseConnectionResetError)
+ self.connections.close.assert_called_once_with(mock_conn)
+ self.connections.open.assert_called_once_with(mock_conn)
+
  def test_is_retryable(self):
  _is_retryable = dbt.adapters.bigquery.connections._is_retryable
  exceptions = dbt.adapters.bigquery.impl.google.cloud.exceptions
  internal_server_error = exceptions.InternalServerError('code broke')
  bad_request_error = exceptions.BadRequest('code broke')
+ connection_error = ConnectionError('code broke')
+ client_error = exceptions.ClientError('bad code')
 
  self.assertTrue(_is_retryable(internal_server_error))
- self.assertFalse(_is_retryable(bad_request_error))
+ self.assertTrue(_is_retryable(bad_request_error))
+ self.assertTrue(_is_retryable(connection_error))
+ self.assertFalse(_is_retryable(client_error))
 
  def test_drop_dataset(self):
  mock_table = Mock()