From b3d9b078dacebfed88db905203116dc5e80cab63 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Thu, 12 Oct 2017 14:06:22 -0700 Subject: [PATCH] BQ: use a string or list of string for client.extract_table. This will match behavior of copy_table and remove the need for special formatting of the keyword arguments. --- bigquery/google/cloud/bigquery/client.py | 34 +++++++++------- bigquery/tests/unit/test_client.py | 51 ++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 15 deletions(-) diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py index c52fbbda1ba8..d0ec7953ae1f 100644 --- a/bigquery/google/cloud/bigquery/client.py +++ b/bigquery/google/cloud/bigquery/client.py @@ -832,7 +832,9 @@ def copy_table(self, sources, destination, job_id=None, job_config=None, job.begin(retry=retry) return job - def extract_table(self, source, *destination_uris, **kwargs): + def extract_table( + self, source, destination_uris, job_config=None, job_id=None, + retry=DEFAULT_RETRY): """Start a job to extract a table into Cloud Storage files. See @@ -841,7 +843,9 @@ def extract_table(self, source, *destination_uris, **kwargs): :type source: :class:`google.cloud.bigquery.table.TableReference` :param source: table to be extracted. - :type destination_uris: sequence of string + :type destination_uris: One of: + str or + sequence of str :param destination_uris: URIs of Cloud Storage file(s) into which table data is to be extracted; in format ``gs:///``. @@ -849,25 +853,25 @@ def extract_table(self, source, *destination_uris, **kwargs): :type kwargs: dict :param kwargs: Additional keyword arguments. - :Keyword Arguments: - * *job_config* - (:class:`google.cloud.bigquery.job.ExtractJobConfig`) -- - (Optional) Extra configuration options for the extract job. - * *job_id* (``str``) -- - Additional content - (Optional) The ID of the job. - * *retry* (:class:`google.api.core.retry.Retry`) - (Optional) How to retry the RPC. + :type job_id: str + :param job_id: (Optional) The ID of the job. + + :type job_config: :class:`google.cloud.bigquery.job.ExtractJobConfig` + :param job_config: (Optional) Extra configuration options for the job. + + :type retry: :class:`google.api.core.retry.Retry` + :param retry: (Optional) How to retry the RPC. :rtype: :class:`google.cloud.bigquery.job.ExtractJob` :returns: a new ``ExtractJob`` instance """ - job_config = kwargs.get('job_config') - job_id = _make_job_id(kwargs.get('job_id')) - retry = kwargs.get('retry', DEFAULT_RETRY) + job_id = _make_job_id(job_id) + + if isinstance(destination_uris, six.string_types): + destination_uris = [destination_uris] job = ExtractJob( - job_id, source, list(destination_uris), client=self, + job_id, source, destination_uris, client=self, job_config=job_config) job.begin(retry=retry) return job diff --git a/bigquery/tests/unit/test_client.py b/bigquery/tests/unit/test_client.py index 6a40a65b390a..243a892fdb86 100644 --- a/bigquery/tests/unit/test_client.py +++ b/bigquery/tests/unit/test_client.py @@ -1703,6 +1703,57 @@ def test_extract_table_generated_job_id(self): self.assertEqual(job.source, source) self.assertEqual(list(job.destination_uris), [DESTINATION]) + def test_extract_table_w_destination_uris(self): + from google.cloud.bigquery.job import ExtractJob + + JOB = 'job_id' + SOURCE = 'source_table' + DESTINATION1 = 'gs://bucket_name/object_one' + DESTINATION2 = 'gs://bucket_name/object_two' + RESOURCE = { + 'jobReference': { + 'projectId': self.PROJECT, + 'jobId': JOB, + }, + 'configuration': { + 'extract': { + 'sourceTable': { + 'projectId': self.PROJECT, + 'datasetId': self.DS_ID, + 'tableId': SOURCE, + }, + 'destinationUris': [ + DESTINATION1, + DESTINATION2, + ], + }, + }, + } + creds = _make_credentials() + http = object() + client = self._make_one(project=self.PROJECT, credentials=creds, + _http=http) + conn = client._connection = _Connection(RESOURCE) + dataset = client.dataset(self.DS_ID) + source = dataset.table(SOURCE) + + job = client.extract_table( + source, [DESTINATION1, DESTINATION2], job_id=JOB) + + # Check that extract_table actually starts the job. + self.assertEqual(len(conn._requested), 1) + req = conn._requested[0] + self.assertEqual(req['method'], 'POST') + self.assertEqual(req['path'], '/projects/PROJECT/jobs') + + # Check the job resource. + self.assertIsInstance(job, ExtractJob) + self.assertIs(job._client, client) + self.assertEqual(job.job_id, JOB) + self.assertEqual(job.source, source) + self.assertEqual( + list(job.destination_uris), [DESTINATION1, DESTINATION2]) + def test_query_defaults(self): from google.cloud.bigquery.job import QueryJob