feat: add fields parameter to set_iam_policy for consistency with update methods (#1872)

Author: tswast

google/cloud/bigquery/client.py

@@ -882,6 +882,35 @@ def get_iam_policy(
         retry: retries.Retry = DEFAULT_RETRY,
         timeout: TimeoutType = DEFAULT_TIMEOUT,
     ) -> Policy:
+        """Return the access control policy for a table resource.
+
+        Args:
+            table (Union[ \
+                google.cloud.bigquery.table.Table, \
+                google.cloud.bigquery.table.TableReference, \
+                google.cloud.bigquery.table.TableListItem, \
+                str, \
+            ]):
+                The table to get the access control policy for.
+                If a string is passed in, this method attempts to create a
+                table reference from a string using
+                :func:`~google.cloud.bigquery.table.TableReference.from_string`.
+            requested_policy_version (int):
+                Optional. The maximum policy version that will be used to format the policy.
+
+                Only version ``1`` is currently supported.
+
+                See: https://cloud.google.com/bigquery/docs/reference/rest/v2/GetPolicyOptions
+            retry (Optional[google.api_core.retry.Retry]):
+                How to retry the RPC.
+            timeout (Optional[float]):
+                The number of seconds to wait for the underlying HTTP transport
+                before using ``retry``.
+
+        Returns:
+            google.api_core.iam.Policy:
+                The access control policy.
+        """
         table = _table_arg_to_table_ref(table, default_project=self.project)
 
         if requested_policy_version != 1:
@@ -910,16 +939,62 @@ def set_iam_policy(
         updateMask: Optional[str] = None,
         retry: retries.Retry = DEFAULT_RETRY,
         timeout: TimeoutType = DEFAULT_TIMEOUT,
+        *,
+        fields: Sequence[str] = (),
     ) -> Policy:
+        """Return the access control policy for a table resource.
+
+        Args:
+            table (Union[ \
+                google.cloud.bigquery.table.Table, \
+                google.cloud.bigquery.table.TableReference, \
+                google.cloud.bigquery.table.TableListItem, \
+                str, \
+            ]):
+                The table to get the access control policy for.
+                If a string is passed in, this method attempts to create a
+                table reference from a string using
+                :func:`~google.cloud.bigquery.table.TableReference.from_string`.
+            policy (google.api_core.iam.Policy):
+                The access control policy to set.
+            updateMask (Optional[str]):
+                Mask as defined by
+                https://cloud.google.com/bigquery/docs/reference/rest/v2/tables/setIamPolicy#body.request_body.FIELDS.update_mask
+
+                Incompatible with ``fields``.
+            retry (Optional[google.api_core.retry.Retry]):
+                How to retry the RPC.
+            timeout (Optional[float]):
+                The number of seconds to wait for the underlying HTTP transport
+                before using ``retry``.
+            fields (Sequence[str]):
+                Which properties to set on the policy. See:
+                https://cloud.google.com/bigquery/docs/reference/rest/v2/tables/setIamPolicy#body.request_body.FIELDS.update_mask
+
+                Incompatible with ``updateMask``.
+
+        Returns:
+            google.api_core.iam.Policy:
+                The updated access control policy.
+        """
+        if updateMask is not None and not fields:
+            update_mask = updateMask
+        elif updateMask is not None and fields:
+            raise ValueError("Cannot set both fields and updateMask")
+        elif fields:
+            update_mask = ",".join(fields)
+        else:
+            update_mask = None
+
         table = _table_arg_to_table_ref(table, default_project=self.project)
 
         if not isinstance(policy, (Policy)):
             raise TypeError("policy must be a Policy")
 
         body = {"policy": policy.to_api_repr()}
 
-        if updateMask is not None:
-            body["updateMask"] = updateMask
+        if update_mask is not None:
+            body["updateMask"] = update_mask
 
         path = "{}:setIamPolicy".format(table.path)
         span_attributes = {"path": path}

samples/snippets/create_iam_policy_test.py

@@ -0,0 +1,44 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def test_create_iam_policy(table_id: str):
+    your_table_id = table_id
+
+    # [START bigquery_create_iam_policy]
+    from google.cloud import bigquery
+
+    bqclient = bigquery.Client()
+
+    policy = bqclient.get_iam_policy(
+        your_table_id,  # e.g. "project.dataset.table"
+    )
+
+    analyst_email = "[email protected]"
+    binding = {
+        "role": "roles/bigquery.dataViewer",
+        "members": {f"group:{analyst_email}"},
+    }
+    policy.bindings.append(binding)
+
+    updated_policy = bqclient.set_iam_policy(
+        your_table_id,  # e.g. "project.dataset.table"
+        policy,
+    )
+
+    for binding in updated_policy.bindings:
+        print(repr(binding))
+    # [END bigquery_create_iam_policy]
+
+    assert binding in updated_policy.bindings

tests/system/test_client.py

@@ -36,7 +36,6 @@
 from google.api_core.exceptions import InternalServerError
 from google.api_core.exceptions import ServiceUnavailable
 from google.api_core.exceptions import TooManyRequests
-from google.api_core.iam import Policy
 from google.cloud import bigquery
 from google.cloud.bigquery.dataset import Dataset
 from google.cloud.bigquery.dataset import DatasetReference
@@ -1485,33 +1484,6 @@ def test_copy_table(self):
         got_rows = self._fetch_single_page(dest_table)
         self.assertTrue(len(got_rows) > 0)
 
-    def test_get_set_iam_policy(self):
-        from google.cloud.bigquery.iam import BIGQUERY_DATA_VIEWER_ROLE
-
-        dataset = self.temp_dataset(_make_dataset_id("create_table"))
-        table_id = "test_table"
-        table_ref = Table(dataset.table(table_id))
-        self.assertFalse(_table_exists(table_ref))
-
-        table = helpers.retry_403(Config.CLIENT.create_table)(table_ref)
-        self.to_delete.insert(0, table)
-
-        self.assertTrue(_table_exists(table))
-
-        member = "serviceAccount:{}".format(Config.CLIENT.get_service_account_email())
-        BINDING = {
-            "role": BIGQUERY_DATA_VIEWER_ROLE,
-            "members": {member},
-        }
-
-        policy = Config.CLIENT.get_iam_policy(table)
-        self.assertIsInstance(policy, Policy)
-        self.assertEqual(policy.bindings, [])
-
-        policy.bindings.append(BINDING)
-        returned_policy = Config.CLIENT.set_iam_policy(table, policy)
-        self.assertEqual(returned_policy.bindings, policy.bindings)
-
     def test_test_iam_permissions(self):
         dataset = self.temp_dataset(_make_dataset_id("create_table"))
         table_id = "test_table"

tests/unit/test_client.py

@@ -1782,6 +1782,60 @@ def test_set_iam_policy(self):
         from google.cloud.bigquery.iam import BIGQUERY_DATA_VIEWER_ROLE
         from google.api_core.iam import Policy
 
+        PATH = "/projects/%s/datasets/%s/tables/%s:setIamPolicy" % (
+            self.PROJECT,
+            self.DS_ID,
+            self.TABLE_ID,
+        )
+        ETAG = "foo"
+        VERSION = 1
+        OWNER1 = "user:[email protected]"
+        OWNER2 = "group:[email protected]"
+        EDITOR1 = "domain:google.com"
+        EDITOR2 = "user:[email protected]"
+        VIEWER1 = "serviceAccount:[email protected]"
+        VIEWER2 = "user:[email protected]"
+        BINDINGS = [
+            {"role": BIGQUERY_DATA_OWNER_ROLE, "members": [OWNER1, OWNER2]},
+            {"role": BIGQUERY_DATA_EDITOR_ROLE, "members": [EDITOR1, EDITOR2]},
+            {"role": BIGQUERY_DATA_VIEWER_ROLE, "members": [VIEWER1, VIEWER2]},
+        ]
+        FIELDS = ("bindings", "etag")
+        RETURNED = {"etag": ETAG, "version": VERSION, "bindings": BINDINGS}
+
+        policy = Policy()
+        for binding in BINDINGS:
+            policy[binding["role"]] = binding["members"]
+
+        BODY = {"policy": policy.to_api_repr(), "updateMask": "bindings,etag"}
+
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+        conn = client._connection = make_connection(RETURNED)
+
+        with mock.patch(
+            "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
+        ) as final_attributes:
+            returned_policy = client.set_iam_policy(
+                self.TABLE_REF, policy, fields=FIELDS, timeout=7.5
+            )
+
+        final_attributes.assert_called_once_with({"path": PATH}, client, None)
+
+        conn.api_request.assert_called_once_with(
+            method="POST", path=PATH, data=BODY, timeout=7.5
+        )
+        self.assertEqual(returned_policy.etag, ETAG)
+        self.assertEqual(returned_policy.version, VERSION)
+        self.assertEqual(dict(returned_policy), dict(policy))
+
+    def test_set_iam_policy_updateMask(self):
+        from google.cloud.bigquery.iam import BIGQUERY_DATA_OWNER_ROLE
+        from google.cloud.bigquery.iam import BIGQUERY_DATA_EDITOR_ROLE
+        from google.cloud.bigquery.iam import BIGQUERY_DATA_VIEWER_ROLE
+        from google.api_core.iam import Policy
+
         PATH = "/projects/%s/datasets/%s/tables/%s:setIamPolicy" % (
             self.PROJECT,
             self.DS_ID,
@@ -1858,6 +1912,19 @@ def test_set_iam_policy_no_mask(self):
             method="POST", path=PATH, data=BODY, timeout=7.5
         )
 
+    def test_set_ia_policy_updateMask_and_fields(self):
+        from google.api_core.iam import Policy
+
+        policy = Policy()
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+
+        with pytest.raises(ValueError, match="updateMask"):
+            client.set_iam_policy(
+                self.TABLE_REF, policy, updateMask="bindings", fields=("bindings",)
+            )
+
     def test_set_iam_policy_invalid_policy(self):
         from google.api_core.iam import Policy