Add Future interface to BigQuery jobs

bigquery/google/cloud/bigquery/job.py

@@ -14,8 +14,14 @@
 
 """Define API Jobs."""
 
+import collections
+import threading
+import warnings
+
 import six
+from six.moves import http_client
 
+from google.cloud import exceptions
 from google.cloud.exceptions import NotFound
 from google.cloud._helpers import _datetime_from_microseconds
 from google.cloud.bigquery.dataset import Dataset
@@ -27,6 +33,60 @@
 from google.cloud.bigquery._helpers import UDFResourcesProperty
 from google.cloud.bigquery._helpers import _EnumProperty
 from google.cloud.bigquery._helpers import _TypedProperty
+import google.cloud.future.base
+
+_DONE_STATE = 'DONE'
+_STOPPED_REASON = 'stopped'
+
+_ERROR_REASON_TO_EXCEPTION = {
+ 'accessDenied': http_client.FORBIDDEN,
+ 'backendError': http_client.INTERNAL_SERVER_ERROR,
+ 'billingNotEnabled': http_client.FORBIDDEN,
+ 'billingTierLimitExceeded': http_client.BAD_REQUEST,
+ 'blocked': http_client.FORBIDDEN,
+ 'duplicate': http_client.CONFLICT,
+ 'internalError': http_client.INTERNAL_SERVER_ERROR,
+ 'invalid': http_client.BAD_REQUEST,
+ 'invalidQuery': http_client.BAD_REQUEST,
+ 'notFound': http_client.NOT_FOUND,
+ 'notImplemented': http_client.NOT_IMPLEMENTED,
+ 'quotaExceeded': http_client.FORBIDDEN,
+ 'rateLimitExceeded': http_client.FORBIDDEN,
+ 'resourceInUse': http_client.BAD_REQUEST,
+ 'resourcesExceeded': http_client.BAD_REQUEST,
+ 'responseTooLarge': http_client.FORBIDDEN,
+ 'stopped': http_client.OK,
+ 'tableUnavailable': http_client.BAD_REQUEST,
+}
+
+_FakeResponse = collections.namedtuple('_FakeResponse', ['status'])
+
+
+def _error_result_to_exception(error_result):
+ """Maps BigQuery error reasons to an exception.
+
+ The reasons and their matching HTTP status codes are documented on
+ the `troubleshooting errors`_ page.
+
+ .. _troubleshooting errors: https://cloud.google.com/bigquery\
+ /troubleshooting-errors
+
+ :type error_result: Mapping[str, str]
+ :param error_result: The error result from BigQuery.
+
+ :rtype google.cloud.exceptions.GoogleCloudError:
+ :returns: The mapped exception.
+ """
+ reason = error_result.get('reason')
+ status_code = _ERROR_REASON_TO_EXCEPTION.get(
+ reason, http_client.INTERNAL_SERVER_ERROR)
+ # make_exception expects an httplib2 response object.
+ fake_response = _FakeResponse(status=status_code)
+ return exceptions.make_exception(
+ fake_response,
+ error_result.get('message', ''),
+ error_info=error_result,
+ use_json=False)
 
 
 class Compression(_EnumProperty):
@@ -82,16 +142,23 @@ class WriteDisposition(_EnumProperty):
  ALLOWED = (WRITE_APPEND, WRITE_TRUNCATE, WRITE_EMPTY)
 
 
-class _BaseJob(object):
- """Base class for jobs.
+class _AsyncJob(google.cloud.future.base.PollingFuture):
+ """Base class for asynchronous jobs.
+
+ :type name: str
+ :param name: the name of the job
 
  :type client: :class:`google.cloud.bigquery.client.Client`
  :param client: A client which holds credentials and project configuration
  for the dataset (which requires a project).
  """
- def __init__(self, client):
+ def __init__(self, name, client):
+ super(_AsyncJob, self).__init__()
+ self.name = name
  self._client = client
  self._properties = {}
+ self._result_set = False
+ self._completion_lock = threading.Lock()
 
  @property
  def project(self):
@@ -117,21 +184,6 @@ def _require_client(self, client):
  client = self._client
  return client
 
-
-class _AsyncJob(_BaseJob):
- """Base class for asynchronous jobs.
-
- :type name: str
- :param name: the name of the job
-
- :type client: :class:`google.cloud.bigquery.client.Client`
- :param client: A client which holds credentials and project configuration
- for the dataset (which requires a project).
- """
- def __init__(self, name, client):
- super(_AsyncJob, self).__init__(client)
- self.name = name
-
  @property
  def job_type(self):
  """Type of job
@@ -273,6 +325,9 @@ def _set_properties(self, api_response):
  self._properties.clear()
  self._properties.update(cleaned)
 
+ # For Future interface
+ self._set_future_result()
+
  @classmethod
  def _get_resource_config(cls, resource):
  """Helper for :meth:`from_api_repr`
@@ -345,7 +400,7 @@ def exists(self, client=None):
  return True
 
  def reload(self, client=None):
- """API call: refresh job properties via a GET request
+ """API call: refresh job properties via a GET request.
 
  See
  https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/get
@@ -371,12 +426,85 @@ def cancel(self, client=None):
  ``NoneType``
  :param client: the client to use. If not passed, falls back to the
  ``client`` stored on the current dataset.
+
+ :rtype: bool
+ :returns: Boolean indicating that the cancel request was sent.
  """
  client = self._require_client(client)
 
  api_response = client._connection.api_request(
  method='POST', path='%s/cancel' % (self.path,))
  self._set_properties(api_response['job'])
+ # The Future interface requires that we return True if the *attempt*
+ # to cancel was successful.
+ return True
+
+ # The following methods implement the PollingFuture interface. Note that
+ # the methods above are from the pre-Future interface and are left for
+ # compatibility. The only "overloaded" method is :meth:`cancel`, which
+ # satisfies both interfaces.
+
+ def _set_future_result(self):
+ """Set the result or exception from the job if it is complete."""
+ # This must be done in a lock to prevent the polling thread
+ # and main thread from both executing the completion logic
+ # at the same time.
+ with self._completion_lock:
+ # If the operation isn't complete or if the result has already been
+ # set, do not call set_result/set_exception again.
+ # Note: self._result_set is set to True in set_result and
+ # set_exception, in case those methods are invoked directly.
+ if self.state != _DONE_STATE or self._result_set:
+ return
+
+ if self.error_result is not None:
+ exception = _error_result_to_exception(self.error_result)
+ self.set_exception(exception)
+ else:
+ self.set_result(self)
+
+ def done(self):
+ """Refresh the job and checks if it is complete.
+
+ :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.reload()
+ return self.state == _DONE_STATE
+
+ def result(self, timeout=None):
+ """Start the job and wait for it to complete and get the result.
+
+ :type timeout: int
+ :param timeout: How long to wait for job to complete before raising
+ a :class:`TimeoutError`.
+
+ :rtype: _AsyncJob
+ :returns: This instance.
+
+ :raises: :class:`~google.cloud.exceptions.GoogleCloudError` if the job
+ failed or :class:`TimeoutError` if the job did not complete in the
+ given timeout.
+ """
+ if self.state is None:
+ self.begin()
+ return super(_AsyncJob, self).result(timeout=timeout)
+
+ def cancelled(self):
+ """Check if the job has been cancelled.
+
+ This always returns False. It's not possible to check if a job was
+ cancelled in the API. This method is here to satisfy the interface
+ for :class:`google.cloud.future.Future`.
+
+ :rtype: bool
+ :returns: False
+ """
+ return (self.error_result is not None
+ and self.error_result.get('reason') == _STOPPED_REASON)
 
 
 class _LoadConfiguration(object):
@@ -1127,11 +1255,44 @@ def from_api_repr(cls, resource, client):
  job._set_properties(resource)
  return job
 
- def results(self):
+ def query_results(self):
  """Construct a QueryResults instance, bound to this job.
 
  :rtype: :class:`~google.cloud.bigquery.query.QueryResults`
  :returns: results instance
  """
  from google.cloud.bigquery.query import QueryResults
  return QueryResults.from_query_job(self)
+
+ def results(self):
+ """DEPRECATED.
+
+ This method is deprecated. Use :meth:`query_results` or :meth:`result`.
+
+ Construct a QueryResults instance, bound to this job.
+
+ :rtype: :class:`~google.cloud.bigquery.query.QueryResults`
+ :returns: The query results.
+ """
+ warnings.warn(
+ 'QueryJob.results() is deprecated. Please use query_results() or '
+ 'result().', DeprecationWarning)
+ return self.query_results()
+
+ def result(self, timeout=None):
+ """Start the job and wait for it to complete and get the result.
+
+ :type timeout: int
+ :param timeout: How long to wait for job to complete before raising
+ a :class:`TimeoutError`.
+
+ :rtype: :class:`~google.cloud.bigquery.query.QueryResults`
+ :returns: The query results.
+
+ :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()

bigquery/tests/system.py

@@ -19,6 +19,7 @@
 import os
 import time
 import unittest
+import uuid
 
 from google.cloud import bigquery
 from google.cloud._helpers import UTC
@@ -1013,6 +1014,15 @@ def test_large_query_w_public_data(self):
  rows = list(iterator)
  self.assertEqual(len(rows), LIMIT)
 
+ def test_async_query_future(self):
+ query_job = Config.CLIENT.run_async_query(
+ str(uuid.uuid4()), 'SELECT 1')
+ query_job.use_legacy_sql = False
+
+ iterator = query_job.result().fetch_data()
+ rows = list(iterator)
+ self.assertEqual(rows, [(1,)])
+
  def test_insert_nested_nested(self):
  # See #2951
  SF = bigquery.SchemaField