Skip to content

Commit 47e3466

Browse files
tswastlandrito
authored andcommitted
BigQuery - add get_query_results method. (googleapis#3838)
This method calls the getQueryResults API directly and returns a QueryResults object. Note: the response from this API does not include the query, so I modified the constructor to make query optional in this case.
1 parent 1e0a788 commit 47e3466

File tree

4 files changed

+111
-0
lines changed

4 files changed

+111
-0
lines changed

bigquery/google/cloud/bigquery/client.py

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,41 @@ def dataset(self, dataset_name, project=None):
162162
"""
163163
return Dataset(dataset_name, client=self, project=project)
164164

165+
def get_query_results(self, job_id, project=None, timeout_ms=None):
166+
"""Get the query results object for a query job.
167+
168+
:type job_id: str
169+
:param job_id: Name of the query job.
170+
171+
:type project: str
172+
:param project:
173+
(Optional) project ID for the query job (defaults to the project of
174+
the client).
175+
176+
:type timeout_ms: int
177+
:param timeout_ms:
178+
(Optional) number of milliseconds the the API call should wait for
179+
the query to complete before the request times out.
180+
181+
:rtype: :class:`google.cloud.bigquery.query.QueryResults`
182+
:returns: a new ``QueryResults`` instance
183+
"""
184+
185+
extra_params = {'maxResults': 0}
186+
187+
if project is None:
188+
project = self.project
189+
190+
if timeout_ms is not None:
191+
extra_params['timeoutMs'] = timeout_ms
192+
193+
path = '/projects/{}/queries/{}'.format(project, job_id)
194+
195+
resource = self._connection.api_request(
196+
method='GET', path=path, query_params=extra_params)
197+
198+
return QueryResults.from_api_repr(resource, self)
199+
165200
def job_from_resource(self, resource):
166201
"""Detect correct job type from resource and instantiate.
167202

bigquery/google/cloud/bigquery/query.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,12 @@ def __init__(self, query, client, udf_resources=(), query_parameters=()):
7676
self.query_parameters = query_parameters
7777
self._job = None
7878

79+
@classmethod
80+
def from_api_repr(cls, api_response, client):
81+
instance = cls(None, client)
82+
instance._set_properties(api_response)
83+
return instance
84+
7985
@classmethod
8086
def from_query_job(cls, job):
8187
"""Factory: construct from an existing job.

bigquery/tests/system.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -599,6 +599,13 @@ def test_job_cancel(self):
599599
# raise an error, and that the job completed (in the `retry()`
600600
# above).
601601

602+
def test_get_query_results(self):
603+
job_id = 'test-get-query-results-' + str(uuid.uuid4())
604+
query_job = Config.CLIENT.run_async_query(job_id, 'SELECT 1')
605+
query_job.begin()
606+
results = Config.CLIENT.get_query_results(job_id)
607+
self.assertEqual(results.total_rows, 1)
608+
602609
def test_sync_query_w_legacy_sql_types(self):
603610
naive = datetime.datetime(2016, 12, 5, 12, 41, 9)
604611
stamp = '%s %s' % (naive.date().isoformat(), naive.time().isoformat())

bigquery/tests/unit/test_client.py

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,64 @@ def test_ctor(self):
4545
self.assertIs(client._connection.credentials, creds)
4646
self.assertIs(client._connection.http, http)
4747

48+
def test_get_job_miss_w_explicit_project_and_timeout(self):
49+
from google.cloud.exceptions import NotFound
50+
51+
project = 'PROJECT'
52+
creds = _make_credentials()
53+
client = self._make_one(project, creds)
54+
conn = client._connection = _Connection()
55+
56+
with self.assertRaises(NotFound):
57+
client.get_query_results(
58+
'nothere', project='other-project', timeout_ms=500)
59+
60+
self.assertEqual(len(conn._requested), 1)
61+
req = conn._requested[0]
62+
self.assertEqual(req['method'], 'GET')
63+
self.assertEqual(
64+
req['path'], '/projects/other-project/queries/nothere')
65+
self.assertEqual(
66+
req['query_params'], {'maxResults': 0, 'timeoutMs': 500})
67+
68+
def test_get_query_results_hit(self):
69+
project = 'PROJECT'
70+
job_id = 'query_job'
71+
data = {
72+
'kind': 'bigquery#getQueryResultsResponse',
73+
'etag': 'some-tag',
74+
'schema': {
75+
'fields': [
76+
{
77+
'name': 'title',
78+
'type': 'STRING',
79+
'mode': 'NULLABLE'
80+
},
81+
{
82+
'name': 'unique_words',
83+
'type': 'INTEGER',
84+
'mode': 'NULLABLE'
85+
}
86+
]
87+
},
88+
'jobReference': {
89+
'projectId': project,
90+
'jobId': job_id,
91+
},
92+
'totalRows': '10',
93+
'totalBytesProcessed': '2464625',
94+
'jobComplete': True,
95+
'cacheHit': False,
96+
}
97+
98+
creds = _make_credentials()
99+
client = self._make_one(project, creds)
100+
client._connection = _Connection(data)
101+
query_results = client.get_query_results(job_id)
102+
103+
self.assertEqual(query_results.total_rows, 10)
104+
self.assertTrue(query_results.complete)
105+
48106
def test_list_projects_defaults(self):
49107
import six
50108
from google.cloud.bigquery.client import Project
@@ -607,6 +665,11 @@ def __init__(self, *responses):
607665
self._requested = []
608666

609667
def api_request(self, **kw):
668+
from google.cloud.exceptions import NotFound
610669
self._requested.append(kw)
670+
671+
if len(self._responses) == 0:
672+
raise NotFound('miss')
673+
611674
response, self._responses = self._responses[0], self._responses[1:]
612675
return response

0 commit comments

Comments
 (0)