BigQuery: Poll via getQueryResults method.

This modifies the QueryJob's Futures interface implementation to poll
using getQueryResults instead of jobs.get. This was recommended by
BigQuery engineers because getQueryResults does HTTP long-polling for
closer to realtime results.

Author: tswast

bigquery/google/cloud/bigquery/client.py

@@ -162,6 +162,41 @@ def dataset(self, dataset_name, project=None):
         """
         return Dataset(dataset_name, client=self, project=project)
 
+    def get_query_results(self, job_id, project=None, timeout_ms=None):
+        """Get the query results object for a query job.
+
+        :type job_id: str
+        :param job_id: Name of the query job.
+
+        :type project: str
+        :param project:
+            (Optional) project ID for the query job (defaults to the project of
+            the client).
+
+        :type timeout_ms: int
+        :param timeout_ms:
+            (Optional) number of milliseconds the the API call should wait for
+            the query to complete before the request times out.
+
+        :rtype: :class:`google.cloud.bigquery.query.QueryResults`
+        :returns: a new ``QueryResults`` instance
+        """
+
+        extra_params = {'maxResults': 0}
+
+        if project is None:
+            project = self.project
+
+        if timeout_ms is not None:
+            extra_params['timeoutMs'] = timeout_ms
+
+        path = '/projects/{}/queries/{}'.format(project, job_id)
+
+        resource = self._connection.api_request(
+            method='GET', path=path, query_params=extra_params)
+
+        return QueryResults.from_api_repr(resource, self)
+
     def job_from_resource(self, resource):
         """Detect correct job type from resource and instantiate.
 

bigquery/google/cloud/bigquery/dbapi/cursor.py

@@ -154,20 +154,13 @@ def execute(self, operation, parameters=None, job_id=None):
             query_parameters=query_parameters)
         query_job.use_legacy_sql = False
 
+        # Wait for the query to finish.
         try:
-            query_results = query_job.result()
+            query_job = query_job.result()
         except google.cloud.exceptions.GoogleCloudError:
             raise exceptions.DatabaseError(query_job.errors)
 
-        # Force the iterator to run because the query_results doesn't
-        # have the total_rows populated. See:
-        # https://github.com/GoogleCloudPlatform/google-cloud-python/issues/3506
-        query_iterator = query_results.fetch_data()
-        try:
-            six.next(iter(query_iterator))
-        except StopIteration:
-            pass
-
+        query_results = query_job.query_results()
         self._query_data = iter(
             query_results.fetch_data(max_results=self.arraysize))
         self._set_rowcount(query_results)

bigquery/google/cloud/bigquery/job.py

@@ -1085,6 +1085,7 @@ def __init__(self, name, query, client,
         self.udf_resources = udf_resources
         self.query_parameters = query_parameters
         self._configuration = _AsyncQueryConfiguration()
+        self._query_results = None
 
     allow_large_results = _TypedProperty('allow_large_results', bool)
     """See
@@ -1284,23 +1285,25 @@ def query_results(self):
         :rtype: :class:`~google.cloud.bigquery.query.QueryResults`
         :returns: results instance
         """
-        from google.cloud.bigquery.query import QueryResults
-        return QueryResults.from_query_job(self)
+        if not self._query_results:
+            self._query_results = self._client.get_query_results(self.name)
+        return self._query_results
 
-    def result(self, timeout=None):
-        """Start the job and wait for it to complete and get the result.
+    def done(self):
+        """Refresh the job and checks if it is complete.
 
-        :type timeout: int
-        :param timeout: How long to wait for job to complete before raising
-            a :class:`TimeoutError`.
+        :rtype: bool
+        :returns: True if the job is complete, False otherwise.
+        """
+        # Do not refresh is the state is already done, as the job will not
+        # change once complete.
+        if self.state != _DONE_STATE:
+            self._query_results = self._client.get_query_results(self.name)
 
-        :rtype: :class:`~google.cloud.bigquery.query.QueryResults`
-        :returns: The query results.
+            # Only reload the job once we know the query is complete.
+            # This will ensure that fields such as the destination table are
+            # correctly populated.
+            if self._query_results.complete:
+                self.reload()
 
-        :raises: :class:`~google.cloud.exceptions.GoogleCloudError` if the job
-            failed or  :class:`TimeoutError` if the job did not complete in the
-            given timeout.
-        """
-        super(QueryJob, self).result(timeout=timeout)
-        # Return a QueryResults instance instead of returning the job.
-        return self.query_results()
+        return self.state == _DONE_STATE

bigquery/google/cloud/bigquery/query.py

@@ -76,6 +76,12 @@ def __init__(self, query, client, udf_resources=(), query_parameters=()):
         self.query_parameters = query_parameters
         self._job = None
 
+    @classmethod
+    def from_api_repr(cls, api_response, client):
+        instance = cls(None, client)
+        instance._set_properties(api_response)
+        return instance
+
     @classmethod
     def from_query_job(cls, job):
         """Factory: construct from an existing job.

bigquery/tests/system.py

@@ -599,6 +599,13 @@ def test_job_cancel(self):
         # raise an error, and that the job completed (in the `retry()`
         # above).
 
+    def test_get_query_results(self):
+        job_id = 'test-get-query-results-' + str(uuid.uuid4())
+        query_job = Config.CLIENT.run_async_query(job_id, 'SELECT 1')
+        query_job.begin()
+        results = Config.CLIENT.get_query_results(job_id)
+        self.assertEqual(results.total_rows, 1)
+
     def test_sync_query_w_legacy_sql_types(self):
         naive = datetime.datetime(2016, 12, 5, 12, 41, 9)
         stamp = '%s %s' % (naive.date().isoformat(), naive.time().isoformat())
@@ -1086,7 +1093,8 @@ def test_async_query_future(self):
             str(uuid.uuid4()), 'SELECT 1')
         query_job.use_legacy_sql = False
 
-        iterator = query_job.result(timeout=JOB_TIMEOUT).fetch_data()
+        query_job = query_job.result(timeout=JOB_TIMEOUT)
+        iterator = query_job.query_results().fetch_data()
         rows = list(iterator)
         self.assertEqual(rows, [(1,)])
 

bigquery/tests/unit/test_client.py

@@ -45,6 +45,64 @@ def test_ctor(self):
         self.assertIs(client._connection.credentials, creds)
         self.assertIs(client._connection.http, http)
 
+    def test_get_job_miss_w_explicit_project_and_timeout(self):
+        from google.cloud.exceptions import NotFound
+
+        project = 'PROJECT'
+        creds = _make_credentials()
+        client = self._make_one(project, creds)
+        conn = client._connection = _Connection()
+
+        with self.assertRaises(NotFound):
+            client.get_query_results(
+                'nothere', project='other-project', timeout_ms=500)
+
+        self.assertEqual(len(conn._requested), 1)
+        req = conn._requested[0]
+        self.assertEqual(req['method'], 'GET')
+        self.assertEqual(
+            req['path'], '/projects/other-project/queries/nothere')
+        self.assertEqual(
+            req['query_params'], {'maxResults': 0, 'timeoutMs': 500})
+
+    def test_get_query_results_hit(self):
+        project = 'PROJECT'
+        job_id = 'query_job'
+        data = {
+            'kind': 'bigquery#getQueryResultsResponse',
+            'etag': 'some-tag',
+            'schema': {
+                'fields': [
+                    {
+                        'name': 'title',
+                        'type': 'STRING',
+                        'mode': 'NULLABLE'
+                    },
+                    {
+                        'name': 'unique_words',
+                        'type': 'INTEGER',
+                        'mode': 'NULLABLE'
+                    }
+                ]
+            },
+            'jobReference': {
+                'projectId': project,
+                'jobId': job_id,
+            },
+            'totalRows': '10',
+            'totalBytesProcessed': '2464625',
+            'jobComplete': True,
+            'cacheHit': False,
+        }
+
+        creds = _make_credentials()
+        client = self._make_one(project, creds)
+        client._connection = _Connection(data)
+        query_results = client.get_query_results(job_id)
+
+        self.assertEqual(query_results.total_rows, 10)
+        self.assertTrue(query_results.complete)
+
     def test_list_projects_defaults(self):
         import six
         from google.cloud.bigquery.client import Project
@@ -607,6 +665,11 @@ def __init__(self, *responses):
         self._requested = []
 
     def api_request(self, **kw):
+        from google.cloud.exceptions import NotFound
         self._requested.append(kw)
+
+        if len(self._responses) == 0:
+            raise NotFound('miss')
+
         response, self._responses = self._responses[0], self._responses[1:]
         return response

bigquery/tests/unit/test_dbapi_cursor.py

@@ -42,7 +42,8 @@ def _mock_job(
         mock_job = mock.create_autospec(job.QueryJob)
         mock_job.error_result = None
         mock_job.state = 'DONE'
-        mock_job.result.return_value = self._mock_results(
+        mock_job.result.return_value = mock_job
+        mock_job.query_results.return_value = self._mock_results(
             rows=rows, schema=schema,
             num_dml_affected_rows=num_dml_affected_rows)
         return mock_job

bigquery/tests/unit/test_job.py

@@ -171,6 +171,7 @@ def _makeResource(self, started=False, ended=False):
         }
 
         if ended:
+            resource['status'] = {'state': 'DONE'}
             resource['statistics']['load']['inputFiles'] = self.INPUT_FILES
             resource['statistics']['load']['inputFileBytes'] = self.INPUT_BYTES
             resource['statistics']['load']['outputBytes'] = self.OUTPUT_BYTES
@@ -310,6 +311,37 @@ def test_ctor_w_schema(self):
                              schema=[full_name, age])
         self.assertEqual(job.schema, [full_name, age])
 
+    def test_done(self):
+        client = _Client(self.PROJECT)
+        resource = self._makeResource(ended=True)
+        job = self._get_target_class().from_api_repr(resource, client)
+        self.assertTrue(job.done())
+
+    def test_result(self):
+        client = _Client(self.PROJECT)
+        resource = self._makeResource(ended=True)
+        job = self._get_target_class().from_api_repr(resource, client)
+
+        result = job.result()
+
+        self.assertIs(result, job)
+
+    def test_result_invokes_begins(self):
+        begun_resource = self._makeResource()
+        done_resource = copy.deepcopy(begun_resource)
+        done_resource['status'] = {'state': 'DONE'}
+        connection = _Connection(begun_resource, done_resource)
+        client = _Client(self.PROJECT, connection=connection)
+        table = _Table()
+        job = self._make_one(self.JOB_NAME, table, [self.SOURCE1], client)
+
+        job.result()
+
+        self.assertEqual(len(connection._requested), 2)
+        begin_request,  reload_request = connection._requested
+        self.assertEqual(begin_request['method'], 'POST')
+        self.assertEqual(reload_request['method'], 'GET')
+
     def test_schema_setter_non_list(self):
         client = _Client(self.PROJECT)
         table = _Table()
@@ -1421,6 +1453,10 @@ def _makeResource(self, started=False, ended=False):
             started, ended)
         config = resource['configuration']['query']
         config['query'] = self.QUERY
+
+        if ended:
+            resource['status'] = {'state': 'DONE'}
+
         return resource
 
     def _verifyBooleanResourceProperties(self, job, config):
@@ -1640,40 +1676,60 @@ def test_cancelled(self):
 
         self.assertTrue(job.cancelled())
 
+    def test_done(self):
+        client = _Client(self.PROJECT)
+        resource = self._makeResource(ended=True)
+        job = self._get_target_class().from_api_repr(resource, client)
+        self.assertTrue(job.done())
+
     def test_query_results(self):
         from google.cloud.bigquery.query import QueryResults
 
-        client = _Client(self.PROJECT)
+        query_resource = {'jobComplete': True}
+        connection = _Connection(query_resource)
+        client = _Client(self.PROJECT, connection=connection)
         job = self._make_one(self.JOB_NAME, self.QUERY, client)
         results = job.query_results()
         self.assertIsInstance(results, QueryResults)
-        self.assertIs(results._job, job)
 
-    def test_result(self):
+    def test_query_results_w_cached_value(self):
         from google.cloud.bigquery.query import QueryResults
 
         client = _Client(self.PROJECT)
         job = self._make_one(self.JOB_NAME, self.QUERY, client)
-        job._properties['status'] = {'state': 'DONE'}
+        query_results = QueryResults(None, client)
+        job._query_results = query_results
+
+        results = job.query_results()
+
+        self.assertIs(results, query_results)
+
+    def test_result(self):
+        client = _Client(self.PROJECT)
+        resource = self._makeResource(ended=True)
+        job = self._get_target_class().from_api_repr(resource, client)
 
         result = job.result()
 
-        self.assertIsInstance(result, QueryResults)
-        self.assertIs(result._job, job)
+        self.assertIs(result, job)
 
     def test_result_invokes_begins(self):
         begun_resource = self._makeResource()
+        incomplete_resource = {'jobComplete': False}
+        query_resource = {'jobComplete': True}
         done_resource = copy.deepcopy(begun_resource)
         done_resource['status'] = {'state': 'DONE'}
-        connection = _Connection(begun_resource, done_resource)
+        connection = _Connection(
+            begun_resource, incomplete_resource, query_resource, done_resource)
         client = _Client(self.PROJECT, connection=connection)
         job = self._make_one(self.JOB_NAME, self.QUERY, client)
 
         job.result()
 
-        self.assertEqual(len(connection._requested), 2)
-        begin_request, reload_request = connection._requested
+        self.assertEqual(len(connection._requested), 4)
+        begin_request, _, query_request, reload_request = connection._requested
         self.assertEqual(begin_request['method'], 'POST')
+        self.assertEqual(query_request['method'], 'GET')
         self.assertEqual(reload_request['method'], 'GET')
 
     def test_result_error(self):
@@ -2088,6 +2144,12 @@ def dataset(self, name):
 
         return Dataset(name, client=self)
 
+    def get_query_results(self, job_id):
+        from google.cloud.bigquery.query import QueryResults
+
+        resource = self._connection.api_request(method='GET')
+        return QueryResults.from_api_repr(resource, self)
+
 
 class _Table(object):