Patch BigQueryClient to retry on network error

luigi/contrib/bigquery.py

@@ -120,11 +120,22 @@ class BigQueryClient:
     https://www.googleapis.com/discovery/v1/apis/bigquery/v2/rest
     """
 
-    def __init__(self, oauth_credentials=None, descriptor='', http_=None):
-        authenticate_kwargs = gcp.get_authenticate_kwargs(oauth_credentials, http_)
-
-        if descriptor:
-            self.client = discovery.build_from_document(descriptor, **authenticate_kwargs)
+    def __init__(self, oauth_credentials=None, descriptor='', http_=None, retry_limit=2):
+        # Save initialisation arguments in case we need to re-create client
+        # due to connection timeout
+        self.oauth_credentials = oauth_credentials
+        self.descriptor = descriptor
+        self.http_ = http_
+        self.retry_limit = retry_limit
+        self.retry_count = 0
+
+        self.__initialise_client()
+
+    def __initialise_client(self):
+        authenticate_kwargs = gcp.get_authenticate_kwargs(self.oauth_credentials, self.http_)
+
+        if self.descriptor:
+            self.client = discovery.build_from_document(self.descriptor, **authenticate_kwargs)
         else:
             self.client = discovery.build('bigquery', 'v2', cache_discovery=False, **authenticate_kwargs)
 
@@ -147,6 +158,12 @@ def dataset_exists(self, dataset):
                         fetched_location if fetched_location is not None else 'unspecified',
                         dataset.location))
 
+        except (TimeoutError, BrokenPipeError, IOError) as bq_connection_error:
+            if self.retry_count > self.retry_limit:
+                raise Exception(f"Exceeded max retries for BigQueryClient connection error: {bq_connection_error}") from bq_connection_error
+            self.retry_count += 1
+            self.__initialise_client()
+            self.dataset_exists(dataset)
         except http.HttpError as ex:
             if ex.resp.status == 404:
                 return False
@@ -167,6 +184,12 @@ def table_exists(self, table):
             self.client.tables().get(projectId=table.project_id,
                                      datasetId=table.dataset_id,
                                      tableId=table.table_id).execute()
+        except (TimeoutError, BrokenPipeError, IOError) as bq_connection_error:
+            if self.retry_count > self.retry_limit:
+                raise Exception(f"Exceeded max retries for BigQueryClient connection error: {bq_connection_error}") from bq_connection_error
+            self.retry_count += 1
+            self.__initialise_client()
+            self.table_exists(table)
         except http.HttpError as ex:
             if ex.resp.status == 404:
                 return False