Merge branch 'main' into start-index-2

Author: chalmerlowe

google/cloud/bigquery/client.py

@@ -90,7 +90,8 @@
 from google.cloud.bigquery.dataset import Dataset
 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 AutoRowIDs, DatasetView, UpdateMode
 from google.cloud.bigquery.format_options import ParquetOptions
 from google.cloud.bigquery.job import (
     CopyJob,
@@ -864,6 +865,7 @@ def get_dataset(
         dataset_ref: Union[DatasetReference, str],
         retry: retries.Retry = DEFAULT_RETRY,
         timeout: TimeoutType = DEFAULT_TIMEOUT,
+        dataset_view: Optional[DatasetView] = None,
     ) -> Dataset:
         """Fetch the dataset referenced by ``dataset_ref``
 
@@ -881,7 +883,21 @@ def get_dataset(
             timeout (Optional[float]):
                 The number of seconds to wait for the underlying HTTP transport
                 before using ``retry``.
-
+            dataset_view (Optional[google.cloud.bigquery.enums.DatasetView]):
+                Specifies the view that determines which dataset information is
+                returned. By default, dataset metadata (e.g. friendlyName, description,
+                labels, etc) and ACL information are returned. This argument can
+                take on the following possible enum values.
+
+                * :attr:`~google.cloud.bigquery.enums.DatasetView.ACL`:
+                    Includes dataset metadata and the ACL.
+                * :attr:`~google.cloud.bigquery.enums.DatasetView.FULL`:
+                    Includes all dataset metadata, including the ACL and table metadata.
+                    This view is not supported by the `datasets.list` API method.
+                * :attr:`~google.cloud.bigquery.enums.DatasetView.METADATA`:
+                    Includes basic dataset metadata, but not the ACL.
+                * :attr:`~google.cloud.bigquery.enums.DatasetView.DATASET_VIEW_UNSPECIFIED`:
+                    The server will decide which view to use. Currently defaults to FULL.
         Returns:
             google.cloud.bigquery.dataset.Dataset:
                 A ``Dataset`` instance.
@@ -891,6 +907,12 @@ def get_dataset(
                 dataset_ref, default_project=self.project
             )
         path = dataset_ref.path
+
+        if dataset_view:
+            query_params = {"datasetView": dataset_view.value}
+        else:
+            query_params = {}
+
         span_attributes = {"path": path}
         api_response = self._call_api(
             retry,
@@ -899,6 +921,7 @@ def get_dataset(
             method="GET",
             path=path,
             timeout=timeout,
+            query_params=query_params,
         )
         return Dataset.from_api_repr(api_response)
 
@@ -1198,6 +1221,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 +1261,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 +1288,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 +1302,7 @@ def update_dataset(
             data=partial,
             headers=headers,
             timeout=timeout,
+            query_params=query_params,
         )
         return Dataset.from_api_repr(api_response)
 

google/cloud/bigquery/enums.py

@@ -80,6 +80,24 @@ class CreateDisposition(object):
     returned in the job result."""
 
 
+class DatasetView(enum.Enum):
+    """DatasetView specifies which dataset information is returned."""
+
+    DATASET_VIEW_UNSPECIFIED = "DATASET_VIEW_UNSPECIFIED"
+    """The default value. Currently maps to the FULL view."""
+
+    METADATA = "METADATA"
+    """View metadata information for the dataset, such as friendlyName,
+    description, labels, etc."""
+
+    ACL = "ACL"
+    """View ACL information for the dataset, which defines dataset access
+    for one or more entities."""
+
+    FULL = "FULL"
+    """View both dataset metadata and ACL information."""
+
+
 class DefaultPandasDTypes(enum.Enum):
     """Default Pandas DataFrem DTypes to convert BigQuery data. These
     Sentinel values are used instead of None to maintain backward compatibility,
@@ -409,6 +427,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."""
 

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, DatasetView
 from google.cloud.bigquery import exceptions
 from google.cloud.bigquery import ParquetOptions
 import google.cloud.bigquery.retry
@@ -752,7 +753,7 @@ def test_get_dataset(self):
         final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None)
 
         conn.api_request.assert_called_once_with(
-            method="GET", path="/%s" % path, timeout=7.5
+            method="GET", path="/%s" % path, timeout=7.5, query_params={}
         )
         self.assertEqual(dataset.dataset_id, self.DS_ID)
 
@@ -818,6 +819,72 @@ def test_get_dataset(self):
 
         self.assertEqual(dataset.dataset_id, self.DS_ID)
 
+    def test_get_dataset_with_dataset_view(self):
+        path = "projects/%s/datasets/%s" % (self.PROJECT, self.DS_ID)
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+        resource = {
+            "id": "%s:%s" % (self.PROJECT, self.DS_ID),
+            "datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},
+        }
+        dataset_ref = DatasetReference(self.PROJECT, self.DS_ID)
+
+        test_cases = [
+            (None, None),
+            (DatasetView.DATASET_VIEW_UNSPECIFIED, "DATASET_VIEW_UNSPECIFIED"),
+            (DatasetView.METADATA, "METADATA"),
+            (DatasetView.ACL, "ACL"),
+            (DatasetView.FULL, "FULL"),
+        ]
+
+        for dataset_view_arg, expected_param_value in test_cases:
+            with self.subTest(
+                dataset_view_arg=dataset_view_arg,
+                expected_param_value=expected_param_value,
+            ):
+                # Re-initialize the connection mock for each sub-test to reset side_effect
+                conn = client._connection = make_connection(resource)
+
+                dataset = client.get_dataset(dataset_ref, dataset_view=dataset_view_arg)
+
+                self.assertEqual(dataset.dataset_id, self.DS_ID)
+
+                if expected_param_value:
+                    expected_query_params = {"datasetView": expected_param_value}
+                else:
+                    expected_query_params = {}
+
+                conn.api_request.assert_called_once_with(
+                    method="GET",
+                    path="/%s" % path,
+                    timeout=DEFAULT_TIMEOUT,
+                    query_params=expected_query_params if expected_query_params else {},
+                )
+
+    def test_get_dataset_with_invalid_dataset_view(self):
+        invalid_view_values = [
+            "INVALID_STRING",
+            123,
+            123.45,
+            object(),
+        ]
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+        resource = {
+            "id": "%s:%s" % (self.PROJECT, self.DS_ID),
+            "datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},
+        }
+        conn = client._connection = make_connection(resource)
+        dataset_ref = DatasetReference(self.PROJECT, self.DS_ID)
+
+        for invalid_view_value in invalid_view_values:
+            with self.subTest(invalid_view_value=invalid_view_value):
+                conn.api_request.reset_mock()  # Reset mock for each sub-test
+                with self.assertRaises(AttributeError):
+                    client.get_dataset(dataset_ref, dataset_view=invalid_view_value)
+
     def test_ensure_bqstorage_client_creating_new_instance(self):
         bigquery_storage = pytest.importorskip("google.cloud.bigquery_storage")
 
@@ -2101,6 +2168,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 +2182,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 +2301,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)

tests/unit/test_create_dataset.py

@@ -372,7 +372,12 @@ def test_create_dataset_alreadyexists_w_exists_ok_true(PROJECT, DS_ID, LOCATION)
                 },
                 timeout=DEFAULT_TIMEOUT,
             ),
-            mock.call(method="GET", path=get_path, timeout=DEFAULT_TIMEOUT),
+            mock.call(
+                method="GET",
+                path=get_path,
+                timeout=DEFAULT_TIMEOUT,
+                query_params={},
+            ),
         ]
     )