Implement retries in BQ adapter

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

@@ -7,6 +7,7 @@
 import google.oauth2
 import google.cloud.exceptions
 import google.cloud.bigquery
+from google.api_core import retry
 
 import dbt.clients.agate_helper
 import dbt.exceptions
@@ -35,6 +36,7 @@ class BigQueryCredentials(Credentials):
     timeout_seconds: Optional[int] = 300
     location: Optional[str] = None
     priority: Optional[Priority] = None
+    retries: Optional[int] = 1
     _ALIASES = {
         'project': 'database',
         'dataset': 'schema',
@@ -57,6 +59,9 @@ class BigQueryConnectionManager(BaseConnectionManager):
              'https://www.googleapis.com/auth/drive')
 
     QUERY_TIMEOUT = 300
+    RETRIES = 1
+    DEFAULT_INITIAL_DELAY = 1.0  # Seconds
+    DEFAULT_MAXIMUM_DELAY = 1.0  # Seconds
 
     @classmethod
     def handle_error(cls, error, message, sql):
@@ -170,6 +175,11 @@ def get_timeout(cls, conn):
         credentials = conn.credentials
         return credentials.timeout_seconds
 
+    @classmethod
+    def get_retries(cls, conn):
+        credentials = conn.credentials
+        return credentials.retries
+
     @classmethod
     def get_table_from_response(cls, resp):
         column_names = [field.name for field in resp.schema]
@@ -182,21 +192,19 @@ def raw_execute(self, sql, fetch=False):
 
         logger.debug('On {}: {}', conn.name, sql)
 
-        job_config = google.cloud.bigquery.QueryJobConfig()
-        job_config.use_legacy_sql = False
+        job_params = {'use_legacy_sql': False}
 
         priority = conn.credentials.priority
         if priority == Priority.Batch:
-            job_config.priority = google.cloud.bigquery.QueryPriority.BATCH
+            job_params['priority'] = google.cloud.bigquery.QueryPriority.BATCH
         else:
-            job_config.priority = \
-                google.cloud.bigquery.QueryPriority.INTERACTIVE
+            job_params[
+                'priority'] = google.cloud.bigquery.QueryPriority.INTERACTIVE
 
-        query_job = client.query(sql, job_config)
+        def fn(): 
+            self._query_and_results(client, sql, conn, job_params)
 
-        # this blocks until the query has completed
-        with self.exception_handler(sql):
-            iterator = query_job.result()
+        query_job, iterator = self._retry_and_handle(msg=sql, conn=conn, fn=fn)
 
         return query_job, iterator
 
@@ -243,8 +251,9 @@ def create_bigquery_table(self, database, schema, table_name, callback,
         view = google.cloud.bigquery.Table(view_ref)
         callback(view)
 
-        with self.exception_handler(sql):
+        def fn():
             client.create_table(view)
+        self._retry_and_handle(msg=sql, conn=conn, fn=fn)
 
     def create_view(self, database, schema, table_name, sql):
         def callback(table):
@@ -258,15 +267,12 @@ def create_table(self, database, schema, table_name, sql):
         client = conn.handle
 
         table_ref = self.table_ref(database, schema, table_name, conn)
-        job_config = google.cloud.bigquery.QueryJobConfig()
-        job_config.destination = table_ref
-        job_config.write_disposition = 'WRITE_TRUNCATE'
-
-        query_job = client.query(sql, job_config=job_config)
+        job_params = {'destination': table_ref,
+                      'write_disposition': 'WRITE_TRUNCATE'}
 
-        # this waits for the job to complete
-        with self.exception_handler(sql):
-            query_job.result(timeout=self.get_timeout(conn))
+        def fn(): 
+            self._query_and_results(client, sql, conn, job_params)
+        self._retry_and_handle(msg=sql, conn=conn, fn=fn)
 
     def create_date_partitioned_table(self, database, schema, table_name):
         def callback(table):
@@ -295,15 +301,70 @@ def drop_dataset(self, database, schema):
         dataset = self.dataset(database, schema, conn)
         client = conn.handle
 
-        with self.exception_handler('drop dataset'):
+        def fn():
             client.delete_dataset(
-                dataset, delete_contents=True, not_found_ok=True
-            )
+                dataset, delete_contents=True, not_found_ok=True)
+
+        self._retry_and_handle(
+          msg='drop dataset', conn=conn, fn=fn)
 
     def create_dataset(self, database, schema):
         conn = self.get_thread_connection()
         client = conn.handle
         dataset = self.dataset(database, schema, conn)
 
-        with self.exception_handler('create dataset'):
-            client.create_dataset(dataset, exists_ok=True)
+        def fn():
+            client.create_dataset(dataset, exist_ok=True)
+        self._retry_and_handle(msg='create dataset', conn=conn, fn=fn)
+
+    def _query_and_results(self, client, sql, conn, job_params):
+        """Query the client and wait for results."""
+        # Cannot reuse job_config if destination is set and ddl is used
+        job_config = google.cloud.bigquery.QueryJobConfig(**job_params)
+        query_job = client.query(sql, job_config=job_config)
+        iterator = query_job.result(timeout=self.get_timeout(conn))
+
+        return query_job, iterator
+
+    def _retry_and_handle(self, msg, conn, fn):
+        """retry a function call within the context of exception_handler."""
+        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)
+
+    def _retry_generator(self):
+        """Generates retry intervals that exponentially back off."""
+        return retry.exponential_sleep_generator(
+            initial=self.DEFAULT_INITIAL_DELAY,
+            maximum=self.DEFAULT_MAXIMUM_DELAY)
+
+
+class _ErrorCounter(object):
+    """Counts errors seen up to a threshold then raises the next error."""
+
+    def __init__(self, retries):
+        self.retries = retries
+        self.error_count = 0
+
+    def count_error(self, error):
+        if self.retries == 0:
+            return False  # Don't log
+        self.error_count += 1
+        if _is_retryable(error) and self.error_count <= self.retries:
+            logger.warning(
+                'Retry attempt %s of %s after error: %s',
+                self.error_count, self.retries, repr(error))
+            return True
+        else:
+            logger.warning(
+                'Not Retrying after %s previous attempts. Error: %s',
+                self.error_count - 1, repr(error))
+            return False
+
+
+def _is_retryable(error):
+    """Return true for 500 level (retryable) errors."""
+    return isinstance(error, google.cloud.exceptions.ServerError)

plugins/bigquery/dbt/adapters/bigquery/impl.py

@@ -93,12 +93,15 @@ def list_schemas(self, database: str) -> List[str]:
         conn = self.connections.get_thread_connection()
         client = conn.handle
 
-        with self.connections.exception_handler('list dataset'):
+        def query_schemas():
             # this is similar to how we have to deal with listing tables
             all_datasets = client.list_datasets(project=database,
                                                 max_results=10000)
             return [ds.dataset_id for ds in all_datasets]
 
+        return self.connections._retry_and_handle(
+            msg='list dataset', conn=conn, fn=query_schemas)
+
     @available.parse(lambda *a, **k: False)
     def check_schema_exists(self, database: str, schema: str) -> bool:
         conn = self.connections.get_thread_connection()

test/unit/test_bigquery_adapter.py

@@ -1,12 +1,15 @@
 import unittest
-from unittest.mock import patch, MagicMock
+from contextlib import contextmanager
+from unittest.mock import patch, MagicMock, Mock
 
 import hologram
 
 import dbt.flags as flags
 
+from dbt.adapters.bigquery import BigQueryCredentials
 from dbt.adapters.bigquery import BigQueryAdapter
 from dbt.adapters.bigquery import BigQueryRelation
+from dbt.adapters.bigquery.connections import BigQueryConnectionManager
 import dbt.exceptions
 from dbt.logger import GLOBAL_LOGGER as logger  # noqa
 
@@ -287,3 +290,81 @@ def test_invalid_relation(self):
         }
         with self.assertRaises(hologram.ValidationError):
             BigQueryRelation.from_dict(kwargs)
+
+
+class TestBigQueryConnectionManager(unittest.TestCase):
+
+    def setUp(self):
+        credentials = Mock(BigQueryCredentials)
+        profile = Mock(query_comment=None, credentials=credentials)
+        self.connections = BigQueryConnectionManager(profile=profile)
+        self.mock_client = Mock(
+          dbt.adapters.bigquery.impl.google.cloud.bigquery.Client)
+        self.mock_connection = MagicMock()
+
+        self.mock_connection.handle = self.mock_client
+
+        self.connections.get_thread_connection = lambda: self.mock_connection
+
+    def test_retry_and_handle(self):
+        self.connections.DEFAULT_MAXIMUM_DELAY = 2.0
+        dbt.adapters.bigquery.connections._is_retryable = lambda x: True
+
+        @contextmanager
+        def dummy_handler(msg):
+            yield
+
+        self.connections.exception_handler = dummy_handler
+
+        class DummyException(Exception):
+            """Count how many times this exception is raised"""
+            count = 0
+
+            def __init__(self):
+                DummyException.count += 1
+
+        def raiseDummyException():
+            raise DummyException()
+
+#        with self.assertLogs(logger.name) as logs:
+        with self.assertRaises(DummyException):
+            self.connections._retry_and_handle(
+                 "some sql", {'credentials': {'retries': 8}},
+                 raiseDummyException)
+            self.assertEqual(DummyException.count, 9)
+#        self.assertIn(
+#            'WARNING:dbt:Retry attempt 1 of 8 after error: DummyException()',
+#            logs.output)
+
+    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')
+
+        self.assertTrue(_is_retryable(internal_server_error))
+        self.assertFalse(_is_retryable(bad_request_error))
+
+    def test_drop_dataset(self):
+        mock_table = Mock()
+        mock_table.reference = 'table1'
+
+        self.mock_client.list_tables.return_value = [mock_table]
+
+        self.connections.drop_dataset('project', 'dataset')
+
+        self.mock_client.list_tables.assert_not_called()
+        self.mock_client.delete_table.assert_not_called()
+        self.mock_client.delete_dataset.assert_called_once()
+
+    @patch('dbt.adapters.bigquery.impl.google.cloud.bigquery')
+    def test_query_and_results(self, mock_bq):
+        self.connections.get_timeout = lambda x: 100.0
+
+        self.connections._query_and_results(
+          self.mock_client, 'sql', self.mock_connection,
+          {'description': 'blah'})
+
+        mock_bq.QueryJobConfig.assert_called_once()
+        self.mock_client.query.assert_called_once_with(
+          'sql', job_config=mock_bq.QueryJobConfig())