Merge branch 'master' into renovate/pyarrow-2.x

Author: tswast

google/cloud/bigquery/client.py

@@ -1534,7 +1534,7 @@ def _get_query_results(
                 A new ``_QueryResults`` instance.
         """
 
-        extra_params = {"maxResults": 0}
+        extra_params = {}
 
         if project is None:
             project = self.project
@@ -3187,6 +3187,7 @@ def _list_rows_from_query_results(
         page_size=None,
         retry=DEFAULT_RETRY,
         timeout=None,
+        first_page_response=None,
     ):
         """List the rows of a completed query.
         See
@@ -3247,6 +3248,7 @@ def _list_rows_from_query_results(
             table=destination,
             extra_params=params,
             total_rows=total_rows,
+            first_page_response=first_page_response,
         )
         return row_iterator
 

google/cloud/bigquery/dataset.py

@@ -79,38 +79,47 @@ class AccessEntry(object):
     """Represents grant of an access role to an entity.
 
     An entry must have exactly one of the allowed :attr:`ENTITY_TYPES`. If
-    anything but ``view`` is set, a ``role`` is also required. ``role`` is
-    omitted for a ``view``, because ``view`` s are always read-only.
+    anything but ``view`` or ``routine`` are set, a ``role`` is also required.
+    ``role`` is omitted for ``view`` and ``routine``, because they are always
+    read-only.
 
     See https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets.
 
     Args:
         role (str):
             Role granted to the entity. The following string values are
             supported: `'READER'`, `'WRITER'`, `'OWNER'`. It may also be
-            :data:`None` if the ``entity_type`` is ``view``.
+            :data:`None` if the ``entity_type`` is ``view`` or ``routine``.
 
         entity_type (str):
             Type of entity being granted the role. One of :attr:`ENTITY_TYPES`.
 
         entity_id (Union[str, Dict[str, str]]):
-            If the ``entity_type`` is not 'view', the ``entity_id`` is the
-            ``str`` ID of the entity being granted the role. If the
-            ``entity_type`` is 'view', the ``entity_id`` is a ``dict``
-            representing the view from a different dataset to grant access to
-            in the following format::
+            If the ``entity_type`` is not 'view' or 'routine', the ``entity_id``
+            is the ``str`` ID of the entity being granted the role. If the
+            ``entity_type`` is 'view' or 'routine', the ``entity_id`` is a ``dict``
+            representing the view  or routine from a different dataset to grant
+            access to in the following format for views::
 
                 {
                     'projectId': string,
                     'datasetId': string,
                     'tableId': string
                 }
 
+            For routines::
+
+                {
+                    'projectId': string,
+                    'datasetId': string,
+                    'routineId': string
+                }
+
     Raises:
         ValueError:
             If the ``entity_type`` is not among :attr:`ENTITY_TYPES`, or if a
-            ``view`` has ``role`` set, or a non ``view`` **does not** have a
-            ``role`` set.
+            ``view`` or a ``routine`` has ``role`` set, or a non ``view`` and
+            non ``routine`` **does not** have a ``role`` set.
 
     Examples:
         >>> entry = AccessEntry('OWNER', 'userByEmail', '[email protected]')
@@ -124,7 +133,15 @@ class AccessEntry(object):
     """
 
     ENTITY_TYPES = frozenset(
-        ["userByEmail", "groupByEmail", "domain", "specialGroup", "view", "iamMember"]
+        [
+            "userByEmail",
+            "groupByEmail",
+            "domain",
+            "specialGroup",
+            "view",
+            "iamMember",
+            "routine",
+        ]
     )
     """Allowed entity types."""
 
@@ -135,10 +152,11 @@ def __init__(self, role, entity_type, entity_id):
                 ", ".join(self.ENTITY_TYPES),
             )
             raise ValueError(message)
-        if entity_type == "view":
+        if entity_type in ("view", "routine"):
             if role is not None:
                 raise ValueError(
-                    "Role must be None for a view. Received " "role: %r" % (role,)
+                    "Role must be None for a %r. Received "
+                    "role: %r" % (entity_type, role)
                 )
         else:
             if role is None:
@@ -409,7 +427,7 @@ def access_entries(self):
         entries.
 
         ``role`` augments the entity type and must be present **unless** the
-        entity type is ``view``.
+        entity type is ``view`` or ``routine``.
 
         Raises:
             TypeError: If 'value' is not a sequence

google/cloud/bigquery/job/query.py

@@ -990,48 +990,22 @@ def done(self, retry=DEFAULT_RETRY, timeout=None, reload=True):
         Returns:
             bool: True if the job is complete, False otherwise.
         """
-        is_done = (
-            # Only consider a QueryJob complete when we know we have the final
-            # query results available.
-            self._query_results is not None
-            and self._query_results.complete
-            and self.state == _DONE_STATE
-        )
         # Do not refresh if the state is already done, as the job will not
         # change once complete.
+        is_done = self.state == _DONE_STATE
         if not reload or is_done:
             return is_done
 
-        # Since the API to getQueryResults can hang up to the timeout value
-        # (default of 10 seconds), set the timeout parameter to ensure that
-        # the timeout from the futures API is respected. See:
-        # https://github.com/GoogleCloudPlatform/google-cloud-python/issues/4135
-        timeout_ms = None
-        if self._done_timeout is not None:
-            # Subtract a buffer for context switching, network latency, etc.
-            api_timeout = self._done_timeout - _TIMEOUT_BUFFER_SECS
-            api_timeout = max(min(api_timeout, 10), 0)
-            self._done_timeout -= api_timeout
-            self._done_timeout = max(0, self._done_timeout)
-            timeout_ms = int(api_timeout * 1000)
+        self._reload_query_results(retry=retry, timeout=timeout)
 
         # If an explicit timeout is not given, fall back to the transport timeout
         # stored in _blocking_poll() in the process of polling for job completion.
         transport_timeout = timeout if timeout is not None else self._transport_timeout
 
-        self._query_results = self._client._get_query_results(
-            self.job_id,
-            retry,
-            project=self.project,
-            timeout_ms=timeout_ms,
-            location=self.location,
-            timeout=transport_timeout,
-        )
-
         # 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 and self.state != _DONE_STATE:
+        if self._query_results.complete:
             self.reload(retry=retry, timeout=transport_timeout)
 
         return self.state == _DONE_STATE
@@ -1098,6 +1072,45 @@ def _begin(self, client=None, retry=DEFAULT_RETRY, timeout=None):
             exc.query_job = self
             raise
 
+    def _reload_query_results(self, retry=DEFAULT_RETRY, timeout=None):
+        """Refresh the cached query results.
+
+        Args:
+            retry (Optional[google.api_core.retry.Retry]):
+                How to retry the call that retrieves query results.
+            timeout (Optional[float]):
+                The number of seconds to wait for the underlying HTTP transport
+                before using ``retry``.
+        """
+        if self._query_results and self._query_results.complete:
+            return
+
+        # Since the API to getQueryResults can hang up to the timeout value
+        # (default of 10 seconds), set the timeout parameter to ensure that
+        # the timeout from the futures API is respected. See:
+        # https://github.com/GoogleCloudPlatform/google-cloud-python/issues/4135
+        timeout_ms = None
+        if self._done_timeout is not None:
+            # Subtract a buffer for context switching, network latency, etc.
+            api_timeout = self._done_timeout - _TIMEOUT_BUFFER_SECS
+            api_timeout = max(min(api_timeout, 10), 0)
+            self._done_timeout -= api_timeout
+            self._done_timeout = max(0, self._done_timeout)
+            timeout_ms = int(api_timeout * 1000)
+
+        # If an explicit timeout is not given, fall back to the transport timeout
+        # stored in _blocking_poll() in the process of polling for job completion.
+        transport_timeout = timeout if timeout is not None else self._transport_timeout
+
+        self._query_results = self._client._get_query_results(
+            self.job_id,
+            retry,
+            project=self.project,
+            timeout_ms=timeout_ms,
+            location=self.location,
+            timeout=transport_timeout,
+        )
+
     def result(
         self,
         page_size=None,
@@ -1144,6 +1157,11 @@ def result(
         """
         try:
             super(QueryJob, self).result(retry=retry, timeout=timeout)
+
+            # Since the job could already be "done" (e.g. got a finished job
+            # via client.get_job), the superclass call to done() might not
+            # set the self._query_results cache.
+            self._reload_query_results(retry=retry, timeout=timeout)
         except exceptions.GoogleAPICallError as exc:
             exc.message += self._format_for_exception(self.query, self.job_id)
             exc.query_job = self
@@ -1158,10 +1176,14 @@ def result(
         if self._query_results.total_rows is None:
             return _EmptyRowIterator()
 
+        first_page_response = None
+        if max_results is None and page_size is None and start_index is None:
+            first_page_response = self._query_results._properties
+
         rows = self._client._list_rows_from_query_results(
-            self._query_results.job_id,
+            self.job_id,
             self.location,
-            self._query_results.project,
+            self.project,
             self._query_results.schema,
             total_rows=self._query_results.total_rows,
             destination=self.destination,
@@ -1170,6 +1192,7 @@ def result(
             start_index=start_index,
             retry=retry,
             timeout=timeout,
+            first_page_response=first_page_response,
         )
         rows._preserve_order = _contains_order_by(self.query)
         return rows

google/cloud/bigquery/table.py

@@ -1308,7 +1308,9 @@ class RowIterator(HTTPIterator):
             A subset of columns to select from this table.
         total_rows (Optional[int]):
             Total number of rows in the table.
-
+        first_page_response (Optional[dict]):
+            API response for the first page of results. These are returned when
+            the first page is requested.
     """
 
     def __init__(
@@ -1324,6 +1326,7 @@ def __init__(
         table=None,
         selected_fields=None,
         total_rows=None,
+        first_page_response=None,
     ):
         super(RowIterator, self).__init__(
             client,
@@ -1346,6 +1349,7 @@ def __init__(
         self._selected_fields = selected_fields
         self._table = table
         self._total_rows = total_rows
+        self._first_page_response = first_page_response
 
     def _get_next_page_response(self):
         """Requests the next page from the path provided.
@@ -1354,6 +1358,11 @@ def _get_next_page_response(self):
             Dict[str, object]:
                 The parsed JSON response of the next page's contents.
         """
+        if self._first_page_response:
+            response = self._first_page_response
+            self._first_page_response = None
+            return response
+
         params = self._get_query_params()
         if self._page_size is not None:
             if self.page_number and "startIndex" in params: