feat: Add UpdateMode to update_dataset (#2204)

chalmerlowe · google-labs-jules[bot] · web-flow · commit eb9c2aff242c · 2025-05-30T17:58:02.000-04:00
* feat: Add UpdateMode to update_dataset

This commit introduces the `UpdateMode` enum and integrates it into the
`update_dataset` method in the BigQuery client.

The `UpdateMode` enum allows you to specify which parts of a dataset
should be updated (metadata, ACL, or full update).

The following changes were made:
- Defined the `UpdateMode` enum in `google/cloud/bigquery/enums.py`
  with values: `UPDATE_MODE_UNSPECIFIED`, `UPDATE_METADATA`,
  `UPDATE_ACL`, and `UPDATE_FULL`.
- Modified the `update_dataset` method in
  `google/cloud/bigquery/client.py` to accept an optional
  `update_mode` parameter. This parameter is added to the query
  parameters if provided.
- Added unit tests in `tests/unit/test_client.py` to verify the
  correct handling of the `update_mode` parameter, including
  testing all enum values and the default case where it's not
  provided.

* updates enums, client, and tests

---------

Co-authored-by: google-labs-jules[bot] &lt;161369871+google-labs-jules[bot]@users.noreply.github.com&gt;
diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py
@@ -91,6 +91,7 @@
 from google.cloud.bigquery.dataset import DatasetListItem
 from google.cloud.bigquery.dataset import DatasetReference
 from google.cloud.bigquery.enums import AutoRowIDs
+from google.cloud.bigquery.enums import UpdateMode
 from google.cloud.bigquery.format_options import ParquetOptions
 from google.cloud.bigquery.job import (
     CopyJob,
@@ -1198,6 +1199,7 @@ def update_dataset(
         fields: Sequence[str],
         retry: retries.Retry = DEFAULT_RETRY,
         timeout: TimeoutType = DEFAULT_TIMEOUT,
+        update_mode: Optional[UpdateMode] = None,
     ) -> Dataset:
         """Change some fields of a dataset.
 
@@ -1237,6 +1239,20 @@ def update_dataset(
             timeout (Optional[float]):
                 The number of seconds to wait for the underlying HTTP transport
                 before using ``retry``.
+            update_mode (Optional[google.cloud.bigquery.enums.UpdateMode]):
+                Specifies the kind of information to update in a dataset.
+                By default, dataset metadata (e.g. friendlyName, description,
+                labels, etc) and ACL information are updated. This argument can
+                take on the following possible enum values.
+
+                * :attr:`~google.cloud.bigquery.enums.UPDATE_MODE_UNSPECIFIED`:
+                    The default value. Behavior defaults to UPDATE_FULL.
+                * :attr:`~google.cloud.bigquery.enums.UpdateMode.UPDATE_METADATA`:
+                    Includes metadata information for the dataset, such as friendlyName, description, labels, etc.
+                * :attr:`~google.cloud.bigquery.enums.UpdateMode.UPDATE_ACL`:
+                    Includes ACL information for the dataset, which defines dataset access for one or more entities.
+                * :attr:`~google.cloud.bigquery.enums.UpdateMode.UPDATE_FULL`:
+                    Includes both dataset metadata and ACL information.
 
         Returns:
             google.cloud.bigquery.dataset.Dataset:
@@ -1250,6 +1266,11 @@ def update_dataset(
         path = dataset.path
         span_attributes = {"path": path, "fields": fields}
 
+        if update_mode:
+            query_params = {"updateMode": update_mode.value}
+        else:
+            query_params = {}
+
         api_response = self._call_api(
             retry,
             span_name="BigQuery.updateDataset",
@@ -1259,6 +1280,7 @@ def update_dataset(
             data=partial,
             headers=headers,
             timeout=timeout,
+            query_params=query_params,
         )
         return Dataset.from_api_repr(api_response)
 
diff --git a/google/cloud/bigquery/enums.py b/google/cloud/bigquery/enums.py
@@ -409,6 +409,24 @@ class BigLakeTableFormat(object):
     """Apache Iceberg format."""
 
 
+class UpdateMode(enum.Enum):
+    """Specifies the kind of information to update in a dataset."""
+
+    UPDATE_MODE_UNSPECIFIED = "UPDATE_MODE_UNSPECIFIED"
+    """The default value. Behavior defaults to UPDATE_FULL."""
+
+    UPDATE_METADATA = "UPDATE_METADATA"
+    """Includes metadata information for the dataset, such as friendlyName,
+    description, labels, etc."""
+
+    UPDATE_ACL = "UPDATE_ACL"
+    """Includes ACL information for the dataset, which defines dataset access
+    for one or more entities."""
+
+    UPDATE_FULL = "UPDATE_FULL"
+    """Includes both dataset metadata and ACL information."""
+
+
 class JobCreationMode(object):
     """Documented values for Job Creation Mode."""
 
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
@@ -60,7 +60,8 @@
 
 from google.cloud.bigquery import job as bqjob
 import google.cloud.bigquery._job_helpers
-from google.cloud.bigquery.dataset import DatasetReference
+from google.cloud.bigquery.dataset import DatasetReference, Dataset
+from google.cloud.bigquery.enums import UpdateMode
 from google.cloud.bigquery import exceptions
 from google.cloud.bigquery import ParquetOptions
 import google.cloud.bigquery.retry
@@ -2101,6 +2102,7 @@ def test_update_dataset(self):
             },
             path="/" + PATH,
             timeout=7.5,
+            query_params={},
         )
         self.assertEqual(ds2.description, ds.description)
         self.assertEqual(ds2.friendly_name, ds.friendly_name)
@@ -2114,6 +2116,94 @@ def test_update_dataset(self):
         client.update_dataset(ds, [])
         req = conn.api_request.call_args
         self.assertEqual(req[1]["headers"]["If-Match"], "etag")
+        self.assertEqual(req[1].get("query_params"), {})
+
+    def test_update_dataset_w_update_mode(self):
+        PATH = f"projects/{self.PROJECT}/datasets/{self.DS_ID}"
+        creds = _make_credentials()
+        client = self._make_one(project=self.PROJECT, credentials=creds)
+
+        DESCRIPTION = "DESCRIPTION"
+        RESOURCE = {
+            "datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},
+            "etag": "etag",
+            "description": DESCRIPTION,
+        }
+        dataset_ref = DatasetReference(self.PROJECT, self.DS_ID)
+        orig_dataset = Dataset(dataset_ref)
+        orig_dataset.description = DESCRIPTION
+        filter_fields = ["description"]
+
+        test_cases = [
+            (None, None),
+            (UpdateMode.UPDATE_MODE_UNSPECIFIED, "UPDATE_MODE_UNSPECIFIED"),
+            (UpdateMode.UPDATE_METADATA, "UPDATE_METADATA"),
+            (UpdateMode.UPDATE_ACL, "UPDATE_ACL"),
+            (UpdateMode.UPDATE_FULL, "UPDATE_FULL"),
+        ]
+
+        for update_mode_arg, expected_param_value in test_cases:
+            with self.subTest(
+                update_mode_arg=update_mode_arg,
+                expected_param_value=expected_param_value,
+            ):
+                conn = client._connection = make_connection(RESOURCE, RESOURCE)
+
+                new_dataset = client.update_dataset(
+                    orig_dataset,
+                    fields=filter_fields,
+                    update_mode=update_mode_arg,
+                )
+                self.assertEqual(orig_dataset.description, new_dataset.description)
+
+                if expected_param_value:
+                    expected_query_params = {"updateMode": expected_param_value}
+                else:
+                    expected_query_params = {}
+
+                conn.api_request.assert_called_once_with(
+                    method="PATCH",
+                    path="/" + PATH,
+                    data={"description": DESCRIPTION},
+                    timeout=DEFAULT_TIMEOUT,
+                    query_params=expected_query_params if expected_query_params else {},
+                )
+
+    def test_update_dataset_w_invalid_update_mode(self):
+        creds = _make_credentials()
+        client = self._make_one(project=self.PROJECT, credentials=creds)
+
+        DESCRIPTION = "DESCRIPTION"
+        resource = {
+            "datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},
+            "etag": "etag",
+        }
+
+        dataset_ref = DatasetReference(self.PROJECT, self.DS_ID)
+        orig_dataset = Dataset(dataset_ref)
+        orig_dataset.description = DESCRIPTION
+        filter_fields = ["description"]  # A non-empty list of fields is required
+
+        # Mock the connection to prevent actual API calls
+        # and to provide a minimal valid response if the call were to proceed.
+        conn = client._connection = make_connection(resource)
+
+        test_cases = [
+            "INVALID_STRING",
+            123,
+            123.45,
+            object(),
+        ]
+
+        for invalid_update_mode in test_cases:
+            with self.subTest(invalid_update_mode=invalid_update_mode):
+                conn.api_request.reset_mock()  # Reset mock for each sub-test
+                with self.assertRaises(AttributeError):
+                    client.update_dataset(
+                        orig_dataset,
+                        fields=filter_fields,
+                        update_mode=invalid_update_mode,
+                    )
 
     def test_update_dataset_w_custom_property(self):
         # The library should handle sending properties to the API that are not
@@ -2145,6 +2235,7 @@ def test_update_dataset_w_custom_property(self):
             data={"newAlphaProperty": "unreleased property"},
             path=path,
             timeout=DEFAULT_TIMEOUT,
+            query_params={},
         )
 
         self.assertEqual(dataset.dataset_id, self.DS_ID)
