BigQuery: Add project parameter to list_datasets and list_jobs (googleapis#5217)

tswast · web-flow · commit 0267a0c080a2 · 2018-04-20T12:10:02.000-07:00
This allows a Client to list datasets or jobs in projects other than the
default project.
diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py
@@ -186,51 +186,53 @@ def list_projects(self, max_results=None, page_token=None,
             page_token=page_token,
             max_results=max_results)
 
-    def list_datasets(self, include_all=False, filter=None, max_results=None,
-                      page_token=None, retry=DEFAULT_RETRY):
+    def list_datasets(
+            self, project=None, include_all=False, filter=None,
+            max_results=None, page_token=None, retry=DEFAULT_RETRY):
         """List datasets for the project associated with this client.
 
         See
         https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/list
 
-        :type include_all: bool
-        :param include_all: True if results include hidden datasets.
-
-        :type filter: str
-        :param filter: (Optional) an expression for filtering the results by
-                       label. For syntax, see
-                       https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/list#filter.
-
-        :type max_results: int
-        :param max_results: (Optional) maximum number of datasets to return,
-                            if not passed, defaults to a value set by the API.
-
-        :type page_token: str
-        :param page_token:
-            (Optional) Token representing a cursor into the datasets. If
-            not passed, the API will return the first page of datasets.
-            The token marks the beginning of the iterator to be returned
-            and the value of the ``page_token`` can be accessed at
-            ``next_page_token`` of the
-            :class:`~google.api_core.page_iterator.HTTPIterator`.
-
-        :type retry: :class:`google.api_core.retry.Retry`
-        :param retry: (Optional) How to retry the RPC.
+        Args:
+            project (str):
+                Optional. Project ID to use for retreiving datasets. Defaults
+                to the client's project.
+            include_all (bool):
+                Optional. True if results include hidden datasets. Defaults
+                to False.
+            filter (str):
+                Optional. An expression for filtering the results by label.
+                For syntax, see
+                https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/list#filter.
+            max_results (int):
+                Optional. Maximum number of datasets to return.
+            page_token (str):
+                Optional. Token representing a cursor into the datasets. If
+                not passed, the API will return the first page of datasets.
+                The token marks the beginning of the iterator to be returned
+                and the value of the ``page_token`` can be accessed at
+                ``next_page_token`` of the
+                :class:`~google.api_core.page_iterator.HTTPIterator`.
+            retry (google.api_core.retry.Retry):
+                Optional. How to retry the RPC.
 
-        :rtype: :class:`~google.api_core.page_iterator.Iterator`
-        :returns:
-            Iterator of
-            :class:`~google.cloud.bigquery.dataset.DatasetListItem`.
-            associated with the client's project.
+        Returns:
+            google.api_core.page_iterator.Iterator:
+                Iterator of
+                :class:`~google.cloud.bigquery.dataset.DatasetListItem`.
+                associated with the project.
         """
         extra_params = {}
+        if project is None:
+            project = self.project
         if include_all:
             extra_params['all'] = True
         if filter:
             # TODO: consider supporting a dict of label -> value for filter,
             # and converting it into a string here.
             extra_params['filter'] = filter
-        path = '/projects/%s/datasets' % (self.project,)
+        path = '/projects/%s/datasets' % (project,)
         return page_iterator.HTTPIterator(
             client=self,
             api_request=functools.partial(self._call_api, retry),
@@ -658,53 +660,55 @@ def cancel_job(
 
         return self.job_from_resource(resource['job'])
 
-    def list_jobs(self, max_results=None, page_token=None, all_users=None,
-                  state_filter=None, retry=DEFAULT_RETRY):
+    def list_jobs(
+            self, project=None, max_results=None, page_token=None,
+            all_users=None, state_filter=None, retry=DEFAULT_RETRY):
         """List jobs for the project associated with this client.
 
         See
         https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/list
 
-        :type max_results: int
-        :param max_results: maximum number of jobs to return, If not
-                            passed, defaults to a value set by the API.
-
-        :type page_token: str
-        :param page_token:
-             (Optional) Opaque marker for the next "page" of jobs. If not
-             passed, the API will return the first page of jobs. The token
-             marks the beginning of the iterator to be returned and the
-             value of the ``page_token`` can be accessed at
-             ``next_page_token`` of
-             :class:`~google.api_core.page_iterator.HTTPIterator`.
-
-        :type all_users: bool
-        :param all_users: if true, include jobs owned by all users in the
-                          project.
-
-        :type state_filter: str
-        :param state_filter: if passed, include only jobs matching the given
-                             state.  One of
-
-                             * ``"done"``
-                             * ``"pending"``
-                             * ``"running"``
-
-        :type retry: :class:`google.api_core.retry.Retry`
-        :param retry: (Optional) How to retry the RPC.
+        Args:
+            project (str):
+                Optional. Project ID to use for retreiving datasets. Defaults
+                to the client's project.
+            max_results (int):
+                Optional. Maximum number of jobs to return.
+            page_token (str):
+                Optional. Opaque marker for the next "page" of jobs. If not
+                passed, the API will return the first page of jobs. The token
+                marks the beginning of the iterator to be returned and the
+                value of the ``page_token`` can be accessed at
+                ``next_page_token`` of
+                :class:`~google.api_core.page_iterator.HTTPIterator`.
+            all_users (bool):
+                If true, include jobs owned by all users in the project.
+            state_filter (str):
+                Optional. If set, include only jobs matching the given
+                state. One of
+
+                    * ``"done"``
+                    * ``"pending"``
+                    * ``"running"``
+            retry (google.api_core.retry.Retry):
+                Optional. How to retry the RPC.
 
-        :rtype: :class:`~google.api_core.page_iterator.Iterator`
-        :returns: Iterable of job instances.
+        Returns:
+            google.api_core.page_iterator.Iterator:
+                Iterable of job instances.
         """
         extra_params = {'projection': 'full'}
 
+        if project is None:
+            project = self.project
+
         if all_users is not None:
             extra_params['allUsers'] = all_users
 
         if state_filter is not None:
             extra_params['stateFilter'] = state_filter
 
-        path = '/projects/%s/jobs' % (self.project,)
+        path = '/projects/%s/jobs' % (project,)
         return page_iterator.HTTPIterator(
             client=self,
             api_request=functools.partial(self._call_api, retry),
diff --git a/bigquery/tests/system.py b/bigquery/tests/system.py
@@ -222,6 +222,12 @@ def test_list_datasets(self):
                    dataset.project == Config.CLIENT.project]
         self.assertEqual(len(created), len(datasets_to_create))
 
+    def test_list_datasets_w_project(self):
+        # Retrieve datasets from a different project.
+        iterator = Config.CLIENT.list_datasets(project='bigquery-public-data')
+        all_datasets = frozenset([dataset.dataset_id for dataset in iterator])
+        self.assertIn('usa_names', all_datasets)
+
     def test_create_table(self):
         dataset = self.temp_dataset(_make_dataset_id('create_table'))
         table_id = 'test_table'
diff --git a/bigquery/tests/unit/test_client.py b/bigquery/tests/unit/test_client.py
@@ -264,6 +264,18 @@ def test_list_datasets_defaults(self):
         conn.api_request.assert_called_once_with(
             method='GET', path='/%s' % PATH, query_params={})
 
+    def test_list_datasets_w_project(self):
+        creds = _make_credentials()
+        client = self._make_one(self.PROJECT, creds)
+        conn = client._connection = _make_connection({})
+
+        list(client.list_datasets(project='other-project'))
+
+        conn.api_request.assert_called_once_with(
+            method='GET',
+            path='/projects/other-project/datasets',
+            query_params={})
+
     def test_list_datasets_explicit_response_missing_datasets_key(self):
         PATH = 'projects/%s/datasets' % self.PROJECT
         TOKEN = 'TOKEN'
@@ -1715,6 +1727,20 @@ def test_list_jobs_explicit_missing(self):
                 'stateFilter': 'done'
             })
 
+    def test_list_jobs_w_project(self):
+        creds = _make_credentials()
+        client = self._make_one(self.PROJECT, creds)
+        conn = client._connection = _make_connection({})
+
+        list(client.list_jobs(project='other-project'))
+
+        conn.api_request.assert_called_once_with(
+            method='GET',
+            path='/projects/other-project/jobs',
+            query_params={
+                'projection': 'full',
+            })
+
     def test_load_table_from_uri(self):
         from google.cloud.bigquery.job import LoadJob
 
