bigquery: add Client.update_dataset

bigquery/google/cloud/bigquery/client.py

@@ -190,9 +190,7 @@ def create_dataset(self, dataset):
         path = '/projects/%s/datasets' % (dataset.project,)
         api_response = self._connection.api_request(
             method='POST', path=path, data=dataset._build_resource())
-        ds = Dataset(dataset.dataset_id, project=dataset.project, client=self)
-        ds._set_properties(api_response)
-        return ds
+        return Dataset.from_api_repr(api_response, self)
 
     def get_dataset(self, dataset_ref):
         """Fetch the dataset referenced by ``dataset_ref``
@@ -222,6 +220,52 @@ def get_table(self, table_ref):
             method='GET', path=table_ref.path)
         return Table.from_api_repr(api_response, self)
 
+    def update_dataset(self, dataset, fields):
+        """Change some fields of a dataset.
+
+        Use ``fields`` to specify which fields to update. At least one field
+        must be provided. If a field is listed in ``fields`` and is ``None`` in
+        ``dataset``, it will be deleted.
+
+        If ``dataset.etag`` is not ``None``, the update will only
+        succeed if the dataset on the server has the same ETag. Thus
+        reading a dataset with ``get_dataset``, changing its fields,
+        and then passing it ``update_dataset`` will ensure that the changes
+        will only be saved if no modifications to the dataset occurred
+        since the read.
+
+        :type dataset: :class:`google.cloud.bigquery.dataset.Dataset`
+        :param dataset: the dataset to update.
+
+        :type fields: sequence of string
+        :param fields: the fields of ``dataset`` to change, spelled as the
+                       Dataset properties (e.g. "friendly_name").
+
+        :rtype: :class:`google.cloud.bigquery.dataset.Dataset`
+        :returns: the modified ``Dataset`` instance
+        :raises: ValueError for fields that cannot be updated.
+
+        """
+        if dataset.project is None:
+            dataset._project = self.project
+        path = '/projects/%s/datasets/%s' % (dataset.project,
+                                             dataset.dataset_id)
+        partial = {}
+        for f in fields:
+            if not hasattr(dataset, f):
+                raise ValueError('No Dataset field %s' % f)
+            # snake case to camel case
+            words = f.split('_')
+            api_field = words[0] + ''.join(map(str.capitalize, words[1:]))
+            partial[api_field] = getattr(dataset, f)
+        if dataset.etag is not None:
+            headers = {'If-Match': dataset.etag}
+        else:
+            headers = None
+        api_response = self._connection.api_request(
+            method='PATCH', path=path, data=partial, headers=headers)
+        return Dataset.from_api_repr(api_response, self)
+
     def _get_query_results(self, job_id, project=None, timeout_ms=None):
         """Get the query results object for a query job.
 

bigquery/google/cloud/bigquery/dataset.py

@@ -488,61 +488,6 @@ def _build_resource(self):
 
         return resource
 
-    def patch(self, client=None, **kw):
-        """API call:  update individual dataset properties via a PATCH request.
-
-        See
-        https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/patch
-
-        :type client: :class:`~google.cloud.bigquery.client.Client` or
-                      ``NoneType``
-        :param client: the client to use.  If not passed, falls back to the
-                       ``client`` stored on the current dataset.
-
-        :type kw: ``dict``
-        :param kw: properties to be patched.
-
-        :raises: ValueError for invalid value types.
-        """
-        client = self._require_client(client)
-
-        partial = {}
-
-        if 'default_table_expiration_ms' in kw:
-            value = kw['default_table_expiration_ms']
-            if not isinstance(value, six.integer_types) and value is not None:
-                raise ValueError("Pass an integer, or None")
-            partial['defaultTableExpirationMs'] = value
-
-        if 'description' in kw:
-            partial['description'] = kw['description']
-
-        if 'friendly_name' in kw:
-            partial['friendlyName'] = kw['friendly_name']
-
-        if 'location' in kw:
-            partial['location'] = kw['location']
-
-        api_response = client._connection.api_request(
-            method='PATCH', path=self.path, data=partial)
-        self._set_properties(api_response)
-
-    def update(self, client=None):
-        """API call:  update dataset properties via a PUT request.
-
-        See
-        https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/update
-
-        :type client: :class:`~google.cloud.bigquery.client.Client` or
-                      ``NoneType``
-        :param client: the client to use.  If not passed, falls back to the
-                       ``client`` stored on the current dataset.
-        """
-        client = self._require_client(client)
-        api_response = client._connection.api_request(
-            method='PUT', path=self.path, data=self._build_resource())
-        self._set_properties(api_response)
-
     def delete(self, client=None):
         """API call:  delete the dataset via a DELETE request.
 

bigquery/tests/system.py

@@ -132,35 +132,22 @@ def test_get_dataset(self):
         self.assertEqual(got.friendly_name, 'Friendly')
         self.assertEqual(got.description, 'Description')
 
-    def test_patch_dataset(self):
+    def test_update_dataset(self):
         dataset = retry_403(Config.CLIENT.create_dataset)(
             Dataset(_make_dataset_id('patch_dataset')))
         self.to_delete.append(dataset)
 
         self.assertTrue(_dataset_exists(dataset))
         self.assertIsNone(dataset.friendly_name)
         self.assertIsNone(dataset.description)
-        dataset.patch(friendly_name='Friendly', description='Description')
-        self.assertEqual(dataset.friendly_name, 'Friendly')
-        self.assertEqual(dataset.description, 'Description')
-
-    def test_update_dataset(self):
-        dataset = retry_403(Config.CLIENT.create_dataset)(
-            Dataset(_make_dataset_id('update_dataset')))
-        self.to_delete.append(dataset)
-
-        self.assertTrue(_dataset_exists(dataset))
-        after = [entry for entry in dataset.access_entries
-                 if entry.entity_id != 'projectWriters']
-        dataset.access_entries = after
-
-        retry_403(dataset.update)()
-
-        self.assertEqual(len(dataset.access_entries), len(after))
-        for found, expected in zip(dataset.access_entries, after):
-            self.assertEqual(found.role, expected.role)
-            self.assertEqual(found.entity_type, expected.entity_type)
-            self.assertEqual(found.entity_id, expected.entity_id)
+        dataset.friendly_name = 'Friendly'
+        dataset.description = 'Description'
+        ds2 = Config.CLIENT.update_dataset(dataset,
+                                           ['friendly_name', 'description'])
+        self.assertEqual(ds2.friendly_name, 'Friendly')
+        self.assertEqual(ds2.description, 'Description')
+
+    # TODO(jba): test that read-modify-write with ETag works.
 
     def test_list_datasets(self):
         datasets_to_create = [

bigquery/tests/unit/test_client.py

@@ -336,6 +336,7 @@ def test_create_dataset_w_attrs(self):
         PATH = 'projects/%s/datasets' % PROJECT
         DESCRIPTION = 'DESC'
         FRIENDLY_NAME = 'FN'
+        LOCATION = 'US'
         USER_EMAIL = 'phred@example.com'
         VIEW = {
             'projectId': 'my-proj',
@@ -349,6 +350,8 @@ def test_create_dataset_w_attrs(self):
             'id': "%s:%s" % (PROJECT, DS_ID),
             'description': DESCRIPTION,
             'friendlyName': FRIENDLY_NAME,
+            'location': LOCATION,
+            'defaultTableExpirationMs': 3600,
             'access': [
                 {'role': 'OWNER', 'userByEmail': USER_EMAIL},
                 {'view': VIEW}],
@@ -361,6 +364,8 @@ def test_create_dataset_w_attrs(self):
         ds_arg = Dataset(DS_ID, project=PROJECT, access_entries=entries)
         ds_arg.description = DESCRIPTION
         ds_arg.friendly_name = FRIENDLY_NAME
+        ds_arg.default_table_expiration_ms = 3600
+        ds_arg.location = LOCATION
         ds = client.create_dataset(ds_arg)
         self.assertEqual(len(conn._requested), 1)
         req = conn._requested[0]
@@ -371,6 +376,8 @@ def test_create_dataset_w_attrs(self):
                 {'projectId': PROJECT, 'datasetId': DS_ID},
             'description': DESCRIPTION,
             'friendlyName': FRIENDLY_NAME,
+            'location': LOCATION,
+            'defaultTableExpirationMs': 3600,
             'access': [
                 {'role': 'OWNER', 'userByEmail': USER_EMAIL},
                 {'view': VIEW}],
@@ -382,6 +389,8 @@ def test_create_dataset_w_attrs(self):
         self.assertEqual(ds.full_dataset_id, RESOURCE['id'])
         self.assertEqual(ds.description, DESCRIPTION)
         self.assertEqual(ds.friendly_name, FRIENDLY_NAME)
+        self.assertEqual(ds.location, LOCATION)
+        self.assertEqual(ds.default_table_expiration_ms, 3600)
 
     def test_get_table(self):
         project = 'PROJECT'
@@ -411,6 +420,66 @@ def test_get_table(self):
         self.assertEqual(req['path'], '/%s' % path)
         self.assertEqual(table.table_id, table_id)
 
+    def test_update_dataset_w_invalid_field(self):
+        from google.cloud.bigquery.dataset import Dataset
+
+        PROJECT = 'PROJECT'
+        DS_ID = 'DATASET_ID'
+        creds = _make_credentials()
+        client = self._make_one(project=PROJECT, credentials=creds)
+        with self.assertRaises(ValueError):
+            client.update_dataset(Dataset(DS_ID), ["foo"])
+
+    def test_update_dataset(self):
+        from google.cloud.bigquery.dataset import Dataset
+
+        PROJECT = 'PROJECT'
+        DS_ID = 'DATASET_ID'
+        PATH = 'projects/%s/datasets/%s' % (PROJECT, DS_ID)
+        DESCRIPTION = 'DESCRIPTION'
+        FRIENDLY_NAME = 'TITLE'
+        LOCATION = 'loc'
+        EXP = 17
+        RESOURCE = {
+            'datasetReference':
+                {'projectId': PROJECT, 'datasetId': DS_ID},
+            'etag': "etag",
+            'description': DESCRIPTION,
+            'friendlyName': FRIENDLY_NAME,
+            'location': LOCATION,
+            'defaultTableExpirationMs': EXP,
+        }
+        creds = _make_credentials()
+        client = self._make_one(project=PROJECT, credentials=creds)
+        conn = client._connection = _Connection(RESOURCE, RESOURCE)
+        ds = Dataset(DS_ID, project=PROJECT)
+        ds.description = DESCRIPTION
+        ds.friendly_name = FRIENDLY_NAME
+        ds.location = LOCATION
+        ds.default_table_expiration_ms = EXP
+        ds2 = client.update_dataset(
+            ds, ['description', 'friendly_name', 'location'])
+        self.assertEqual(len(conn._requested), 1)
+        req = conn._requested[0]
+        self.assertEqual(req['method'], 'PATCH')
+        SENT = {
+            'description': DESCRIPTION,
+            'friendlyName': FRIENDLY_NAME,
+            'location': LOCATION,
+        }
+        self.assertEqual(req['data'], SENT)
+        self.assertEqual(req['path'], '/' + PATH)
+        self.assertIsNone(req['headers'])
+        self.assertEqual(ds2.description, ds.description)
+        self.assertEqual(ds2.friendly_name, ds.friendly_name)
+        self.assertEqual(ds2.location, ds.location)
+
+        # ETag becomes If-Match header.
+        ds._properties['etag'] = 'etag'
+        client.update_dataset(ds, [])
+        req = conn._requested[1]
+        self.assertEqual(req['headers']['If-Match'], 'etag')
+
     def test_job_from_resource_unknown_type(self):
         PROJECT = 'PROJECT'
         creds = _make_credentials()