Skip to content

Commit

Permalink
BQ: use a string or list of string for client.extract_table.
Browse files Browse the repository at this point in the history
This will match behavior of copy_table and remove the need for special
formatting of the keyword arguments.
  • Loading branch information
tswast committed Oct 12, 2017
1 parent d7ad25e commit a492ec6
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 12 deletions.
27 changes: 15 additions & 12 deletions bigquery/google/cloud/bigquery/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
51 changes: 51 additions & 0 deletions bigquery/tests/unit/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down

0 comments on commit a492ec6

Please sign in to comment.