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.

Author: google-labs-jules[bot]

google/cloud/bigquery/client.py

@@ -1198,6 +1198,7 @@ def update_dataset(
         fields: Sequence[str],
         retry: retries.Retry = DEFAULT_RETRY,
         timeout: TimeoutType = DEFAULT_TIMEOUT,
+        update_mode: Optional[enums.UpdateMode] = None,
     ) -> Dataset:
         """Change some fields of a dataset.
 
@@ -1249,6 +1250,9 @@ def update_dataset(
             headers = None
         path = dataset.path
         span_attributes = {"path": path, "fields": fields}
+        query_params: Dict[str, Any] = {}
+        if update_mode is not None:
+            query_params["updateMode"] = str(update_mode.value)
 
         api_response = self._call_api(
             retry,
@@ -1259,6 +1263,7 @@ def update_dataset(
             data=partial,
             headers=headers,
             timeout=timeout,
+            query_params=query_params if query_params else None,
         )
         return Dataset.from_api_repr(api_response)
 

google/cloud/bigquery/enums.py

@@ -409,6 +409,24 @@ class BigLakeTableFormat(object):
     """Apache Iceberg format."""
 
 
+class UpdateMode(str, enum.Enum):
+    """Specifies the kind of information to update in a dataset."""
+
+    UPDATE_MODE_UNSPECIFIED = "UPDATE_MODE_UNSPECIFIED"
+    """The default value. Default to the 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."""
 

tests/unit/test_client.py

@@ -2114,6 +2114,56 @@ def test_update_dataset(self):
         client.update_dataset(ds, [])
         req = conn.api_request.call_args
         self.assertEqual(req[1]["headers"]["If-Match"], "etag")
+        self.assertIsNone(req[1].get("query_params"))
+
+    def test_update_dataset_w_update_mode(self):
+        from google.cloud.bigquery.dataset import Dataset
+        from google.cloud.bigquery import enums
+
+        PATH = "projects/%s/datasets/%s" % (self.PROJECT, self.DS_ID)
+        DESCRIPTION = "DESCRIPTION"
+        RESOURCE = {
+            "datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},
+            "etag": "etag",
+            "description": DESCRIPTION,
+        }
+        creds = _make_credentials()
+        client = self._make_one(project=self.PROJECT, credentials=creds)
+        ds = Dataset(DatasetReference(self.PROJECT, self.DS_ID))
+        ds.description = DESCRIPTION
+        filter_fields = ["description"]
+
+        # Test each UpdateMode enum value
+        for update_mode_enum in enums.UpdateMode:
+            conn = client._connection = make_connection(RESOURCE)
+            ds2 = client.update_dataset(
+                ds,
+                fields=filter_fields,
+                update_mode=update_mode_enum,
+            )
+            self.assertEqual(ds2.description, ds.description)
+            conn.api_request.assert_called_once_with(
+                method="PATCH",
+                data={"description": DESCRIPTION},
+                path="/" + PATH,
+                timeout=DEFAULT_TIMEOUT,
+                query_params={"updateMode": str(update_mode_enum.value)},
+            )
+
+        # Test when update_mode is not provided
+        conn = client._connection = make_connection(RESOURCE)
+        ds2 = client.update_dataset(
+            ds,
+            fields=filter_fields,
+        )
+        self.assertEqual(ds2.description, ds.description)
+        conn.api_request.assert_called_once_with(
+            method="PATCH",
+            data={"description": DESCRIPTION},
+            path="/" + PATH,
+            timeout=DEFAULT_TIMEOUT,
+            query_params=None,  # Expect None or empty dict when not provided
+        )
 
     def test_update_dataset_w_custom_property(self):
         # The library should handle sending properties to the API that are not