BQ: use a string or list of string for client.extract_table.

tswast · tswast · commit a492ec612f6b · 2017-10-12T14:06:22.000-07:00
This will match behavior of copy_table and remove the need for special
formatting of the keyword arguments.
diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py
@@ -772,7 +772,8 @@ def copy_table(self, sources, destination, job_id=None, job_config=None):
         job.begin()
         return job
 
-    def extract_table(self, source, *destination_uris, **kwargs):
+    def extract_table(
+            self, source, destination_uris, job_config=None, job_id=None):
         """Start a job to extract a table into Cloud Storage files.
 
         See
@@ -781,30 +782,32 @@ 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://<bucket_name>/<object_name_or_glob>``.
 
         :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.
+        :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.
 
         :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'))
+        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()
         return job
diff --git a/bigquery/tests/unit/test_client.py b/bigquery/tests/unit/test_client.py
@@ -1668,6 +1668,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
 
