Repository: dbt-labs/dbt-core

Issue:
dbt seed should retry on bigquery when it tells us to
## Issue
Sometimes `dbt seed` runs fail on bigquery with an error telling us to retry, but we don't.

### Issue description
This is especially annoying in integration tests, here's an example:
https://dev.azure.com/fishtown-analytics/dbt/_build/results?buildId=302

The failure text is:
```
Runtime Error in seed source (data\source.csv)
  Error encountered during execution. Retrying may solve the problem.
  You are loading data without specifying data format, data will be treated as CSV format by default. If this is not what you mean, please specify data format by --source_format.
```

And, in fact, retrying _will_ solve the problem. This is some weird transient failure inside bigquery. I'm pretty much positive that the second message is caused by the first (BQ gets an empty csv due to the previous error?) and should be ignored. dbt should catch this and retry the seed.

I'm sure this also happens in the "real world" once in a while.

### Results
`dbt seed` failed with an error. I expected it to succeed

### System information
dbt version 0.14.0-ish

The operating system you're running on: Any OS

The python version you're using (probably the output of `python --version`): Any python

### Steps to reproduce
Run dbt seed a few thousand times because you run a lot of tests
Experience a couple failures, feel frustrated about having to re-run the entire suite each time for this.



Discussion:
This is something I have a need for, and plan on implementing in a fork, but would love to contribute back to dbt. 

My rough plan for how to implement this is to wrap calls to the client (I use bigquery) in a call to a helper function that retries in accordance with the timeout, and a new configurable 'retries' parameter set in profiles.yml. It would be nice to delegate retrying logic to another library that already implements nice things like politely exponentially backing-off (the bigquery adapter already depends on google.api_core which has a 'retry' module that does exactly this).

Before I get too far on something that doesn't align well with what dbt wants, it would be helpful to know if it makes sense to implement this per adapter. The adapter feels to me like the most logical place to implement this, as the first point where exceptions from the client can be caught and retried, but this is a feature that all adapters could benefit from. A compelling reason to implement this per adapter might be that exceptions will vary per adapter, so logic that determines whether an exception is retryable will depend on the adapter.

Curious to hear any thoughts on this approach!
Hey @kconvey - I think this feature definitely has merit on BigQuery. I don't really perceive a need/desire for it on other databases though. BQ returns random 500 errors with some regularity and retrying really does make the query succeed. That just... doesn't usually happen on any of Snowflake/Redshift/Postgres.

My thinking is that when BigQuery returns a 500 error code (or an error message that says "Retrying may solve the problem"), dbt can retry the query. I think there could be merit to retrying with some sort of backoff, but really, I'd be equivalently comfortable retrying a single time after something like 10 seconds. I'm happy for the number of retries and the timeout interval to be configurable though.

Overall, I very much agree with your thinking here! I'd just say that we can make this BQ specific initially, but we should implement it in a way that could be extended to other plugins in the future.
Closed by #1963 - this is going out in dbt v0.15.1

Resolving patch:
diff --git a/plugins/bigquery/dbt/adapters/bigquery/connections.py b/plugins/bigquery/dbt/adapters/bigquery/connections.py
--- a/plugins/bigquery/dbt/adapters/bigquery/connections.py
+++ b/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 patch:
diff --git a/test/unit/test_bigquery_adapter.py b/test/unit/test_bigquery_adapter.py
--- a/test/unit/test_bigquery_adapter.py
+++ b/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()