diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 408e7e49c..5521e2e1e 100644 --- a/google/cloud/bigquery/client.py +++ b/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,7 +939,53 @@ 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)): @@ -918,8 +993,8 @@ def set_iam_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} diff --git a/samples/snippets/create_iam_policy_test.py b/samples/snippets/create_iam_policy_test.py new file mode 100644 index 000000000..c41ced2cd --- /dev/null +++ b/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 = "example-analyst-group@google.com" + 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 diff --git a/tests/system/test_client.py b/tests/system/test_client.py index 04740de8a..414239323 100644 --- a/tests/system/test_client.py +++ b/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" diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index d20712a8a..60dcab85e 100644 --- a/tests/unit/test_client.py +++ b/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:phred@example.com" + OWNER2 = "group:cloud-logs@google.com" + EDITOR1 = "domain:google.com" + EDITOR2 = "user:phred@example.com" + VIEWER1 = "serviceAccount:1234-abcdef@service.example.com" + VIEWER2 = "user:phred@example.com" + 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