Repository: spotify/luigi
File: luigi/contrib/bigquery.py (Python)

Code before:

    def dataset_exists(self, dataset):
        """Returns whether the given dataset exists.
        If regional location is specified for the dataset, that is also checked
        to be compatible with the remote dataset, otherwise an exception is thrown.

           :param dataset:
           :type dataset: BQDataset
        """

        try:
            response = self.client.datasets().get(projectId=dataset.project_id,
                                                  datasetId=dataset.dataset_id).execute()
            if dataset.location is not None:
                fetched_location = response.get('location')
                if dataset.location != fetched_location:
                    raise Exception('''Dataset already exists with regional location {}. Can't use {}.'''.format(
                        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
            raise

        return True

    def table_exists(self, table):
        """Returns whether the given table exists.

           :param table:
           :type table: BQTable
        """
        if not self.dataset_exists(table.dataset):
            return False

        try:
            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()

Review comment (refactor):
How about using retry decorator to simplify the codes?
The implementation below will be helpful.
https://github.com/spotify/luigi/blob/master/luigi/contrib/gcs.py#L72-L81

Code after:
    """A client for Google BigQuery.

    For details of how authentication and the descriptor work, see the
    documentation for the GCS client. The descriptor URL for BigQuery is
    https://www.googleapis.com/discovery/v1/apis/bigquery/v2/rest
    """

    def __init__(self, oauth_credentials=None, descriptor='', http_=None):
        # 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.__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)

    @bq_retry
    def dataset_exists(self, dataset):
        """Returns whether the given dataset exists.
        If regional location is specified for the dataset, that is also checked
        to be compatible with the remote dataset, otherwise an exception is thrown.

           :param dataset:
           :type dataset: BQDataset
        """

        try:
            response = self.client.datasets().get(projectId=dataset.project_id,
                                                  datasetId=dataset.dataset_id).execute()
            if dataset.location is not None:
                fetched_location = response.get('location')
                if dataset.location != fetched_location:
                    raise Exception('''Dataset already exists with regional location {}. Can't use {}.'''.format(
                        fetched_location if fetched_location is not None else 'unspecified',
                        dataset.location))
        except http.HttpError as ex:
            if ex.resp.status == 404:
                return False
            raise

        return True

    @bq_retry