Skip to content
This repository was archived by the owner on Mar 6, 2026. It is now read-only.

Commit 08b1e6f

Browse files
tswasttswast
authored
feat: add fields parameter to set_iam_policy for consistency with update methods (#1872)
1 parent e265db6 commit 08b1e6f

4 files changed

Lines changed: 188 additions & 30 deletions

File tree

‎google/cloud/bigquery/client.py‎

Lines changed: 77 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -882,6 +882,35 @@ def get_iam_policy(
882882
retry: retries.Retry = DEFAULT_RETRY,
883883
timeout: TimeoutType = DEFAULT_TIMEOUT,
884884
) -> Policy:
885+
"""Return the access control policy for a table resource.
886+
887+
Args:
888+
table (Union[ \
889+
google.cloud.bigquery.table.Table, \
890+
google.cloud.bigquery.table.TableReference, \
891+
google.cloud.bigquery.table.TableListItem, \
892+
str, \
893+
]):
894+
The table to get the access control policy for.
895+
If a string is passed in, this method attempts to create a
896+
table reference from a string using
897+
:func:`~google.cloud.bigquery.table.TableReference.from_string`.
898+
requested_policy_version (int):
899+
Optional. The maximum policy version that will be used to format the policy.
900+
901+
Only version ``1`` is currently supported.
902+
903+
See: https://cloud.google.com/bigquery/docs/reference/rest/v2/GetPolicyOptions
904+
retry (Optional[google.api_core.retry.Retry]):
905+
How to retry the RPC.
906+
timeout (Optional[float]):
907+
The number of seconds to wait for the underlying HTTP transport
908+
before using ``retry``.
909+
910+
Returns:
911+
google.api_core.iam.Policy:
912+
The access control policy.
913+
"""
885914
table = _table_arg_to_table_ref(table, default_project=self.project)
886915

887916
if requested_policy_version != 1:
@@ -910,16 +939,62 @@ def set_iam_policy(
910939
updateMask: Optional[str] = None,
911940
retry: retries.Retry = DEFAULT_RETRY,
912941
timeout: TimeoutType = DEFAULT_TIMEOUT,
942+
*,
943+
fields: Sequence[str] = (),
913944
) -> Policy:
945+
"""Return the access control policy for a table resource.
946+
947+
Args:
948+
table (Union[ \
949+
google.cloud.bigquery.table.Table, \
950+
google.cloud.bigquery.table.TableReference, \
951+
google.cloud.bigquery.table.TableListItem, \
952+
str, \
953+
]):
954+
The table to get the access control policy for.
955+
If a string is passed in, this method attempts to create a
956+
table reference from a string using
957+
:func:`~google.cloud.bigquery.table.TableReference.from_string`.
958+
policy (google.api_core.iam.Policy):
959+
The access control policy to set.
960+
updateMask (Optional[str]):
961+
Mask as defined by
962+
https://cloud.google.com/bigquery/docs/reference/rest/v2/tables/setIamPolicy#body.request_body.FIELDS.update_mask
963+
964+
Incompatible with ``fields``.
965+
retry (Optional[google.api_core.retry.Retry]):
966+
How to retry the RPC.
967+
timeout (Optional[float]):
968+
The number of seconds to wait for the underlying HTTP transport
969+
before using ``retry``.
970+
fields (Sequence[str]):
971+
Which properties to set on the policy. See:
972+
https://cloud.google.com/bigquery/docs/reference/rest/v2/tables/setIamPolicy#body.request_body.FIELDS.update_mask
973+
974+
Incompatible with ``updateMask``.
975+
976+
Returns:
977+
google.api_core.iam.Policy:
978+
The updated access control policy.
979+
"""
980+
if updateMask is not None and not fields:
981+
update_mask = updateMask
982+
elif updateMask is not None and fields:
983+
raise ValueError("Cannot set both fields and updateMask")
984+
elif fields:
985+
update_mask = ",".join(fields)
986+
else:
987+
update_mask = None
988+
914989
table = _table_arg_to_table_ref(table, default_project=self.project)
915990

916991
if not isinstance(policy, (Policy)):
917992
raise TypeError("policy must be a Policy")
918993

919994
body = {"policy": policy.to_api_repr()}
920995

921-
if updateMask is not None:
922-
body["updateMask"] = updateMask
996+
if update_mask is not None:
997+
body["updateMask"] = update_mask
923998

924999
path = "{}:setIamPolicy".format(table.path)
9251000
span_attributes = {"path": path}

‎samples/snippets/create_iam_policy_test.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
def test_create_iam_policy(table_id: str):
17+
your_table_id = table_id
18+
19+
# [START bigquery_create_iam_policy]
20+
from google.cloud import bigquery
21+
22+
bqclient = bigquery.Client()
23+
24+
policy = bqclient.get_iam_policy(
25+
your_table_id, # e.g. "project.dataset.table"
26+
)
27+
28+
analyst_email = "example-analyst-group@google.com"
29+
binding = {
30+
"role": "roles/bigquery.dataViewer",
31+
"members": {f"group:{analyst_email}"},
32+
}
33+
policy.bindings.append(binding)
34+
35+
updated_policy = bqclient.set_iam_policy(
36+
your_table_id, # e.g. "project.dataset.table"
37+
policy,
38+
)
39+
40+
for binding in updated_policy.bindings:
41+
print(repr(binding))
42+
# [END bigquery_create_iam_policy]
43+
44+
assert binding in updated_policy.bindings

‎tests/system/test_client.py‎

Lines changed: 0 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,6 @@
3636
from google.api_core.exceptions import InternalServerError
3737
from google.api_core.exceptions import ServiceUnavailable
3838
from google.api_core.exceptions import TooManyRequests
39-
from google.api_core.iam import Policy
4039
from google.cloud import bigquery
4140
from google.cloud.bigquery.dataset import Dataset
4241
from google.cloud.bigquery.dataset import DatasetReference
@@ -1485,33 +1484,6 @@ def test_copy_table(self):
14851484
got_rows = self._fetch_single_page(dest_table)
14861485
self.assertTrue(len(got_rows) > 0)
14871486

1488-
def test_get_set_iam_policy(self):
1489-
from google.cloud.bigquery.iam import BIGQUERY_DATA_VIEWER_ROLE
1490-
1491-
dataset = self.temp_dataset(_make_dataset_id("create_table"))
1492-
table_id = "test_table"
1493-
table_ref = Table(dataset.table(table_id))
1494-
self.assertFalse(_table_exists(table_ref))
1495-
1496-
table = helpers.retry_403(Config.CLIENT.create_table)(table_ref)
1497-
self.to_delete.insert(0, table)
1498-
1499-
self.assertTrue(_table_exists(table))
1500-
1501-
member = "serviceAccount:{}".format(Config.CLIENT.get_service_account_email())
1502-
BINDING = {
1503-
"role": BIGQUERY_DATA_VIEWER_ROLE,
1504-
"members": {member},
1505-
}
1506-
1507-
policy = Config.CLIENT.get_iam_policy(table)
1508-
self.assertIsInstance(policy, Policy)
1509-
self.assertEqual(policy.bindings, [])
1510-
1511-
policy.bindings.append(BINDING)
1512-
returned_policy = Config.CLIENT.set_iam_policy(table, policy)
1513-
self.assertEqual(returned_policy.bindings, policy.bindings)
1514-
15151487
def test_test_iam_permissions(self):
15161488
dataset = self.temp_dataset(_make_dataset_id("create_table"))
15171489
table_id = "test_table"

‎tests/unit/test_client.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1782,6 +1782,60 @@ def test_set_iam_policy(self):
17821782
from google.cloud.bigquery.iam import BIGQUERY_DATA_VIEWER_ROLE
17831783
from google.api_core.iam import Policy
17841784

1785+
PATH = "/projects/%s/datasets/%s/tables/%s:setIamPolicy" % (
1786+
self.PROJECT,
1787+
self.DS_ID,
1788+
self.TABLE_ID,
1789+
)
1790+
ETAG = "foo"
1791+
VERSION = 1
1792+
OWNER1 = "user:phred@example.com"
1793+
OWNER2 = "group:cloud-logs@google.com"
1794+
EDITOR1 = "domain:google.com"
1795+
EDITOR2 = "user:phred@example.com"
1796+
VIEWER1 = "serviceAccount:1234-abcdef@service.example.com"
1797+
VIEWER2 = "user:phred@example.com"
1798+
BINDINGS = [
1799+
{"role": BIGQUERY_DATA_OWNER_ROLE, "members": [OWNER1, OWNER2]},
1800+
{"role": BIGQUERY_DATA_EDITOR_ROLE, "members": [EDITOR1, EDITOR2]},
1801+
{"role": BIGQUERY_DATA_VIEWER_ROLE, "members": [VIEWER1, VIEWER2]},
1802+
]
1803+
FIELDS = ("bindings", "etag")
1804+
RETURNED = {"etag": ETAG, "version": VERSION, "bindings": BINDINGS}
1805+
1806+
policy = Policy()
1807+
for binding in BINDINGS:
1808+
policy[binding["role"]] = binding["members"]
1809+
1810+
BODY = {"policy": policy.to_api_repr(), "updateMask": "bindings,etag"}
1811+
1812+
creds = _make_credentials()
1813+
http = object()
1814+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
1815+
conn = client._connection = make_connection(RETURNED)
1816+
1817+
with mock.patch(
1818+
"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
1819+
) as final_attributes:
1820+
returned_policy = client.set_iam_policy(
1821+
self.TABLE_REF, policy, fields=FIELDS, timeout=7.5
1822+
)
1823+
1824+
final_attributes.assert_called_once_with({"path": PATH}, client, None)
1825+
1826+
conn.api_request.assert_called_once_with(
1827+
method="POST", path=PATH, data=BODY, timeout=7.5
1828+
)
1829+
self.assertEqual(returned_policy.etag, ETAG)
1830+
self.assertEqual(returned_policy.version, VERSION)
1831+
self.assertEqual(dict(returned_policy), dict(policy))
1832+
1833+
def test_set_iam_policy_updateMask(self):
1834+
from google.cloud.bigquery.iam import BIGQUERY_DATA_OWNER_ROLE
1835+
from google.cloud.bigquery.iam import BIGQUERY_DATA_EDITOR_ROLE
1836+
from google.cloud.bigquery.iam import BIGQUERY_DATA_VIEWER_ROLE
1837+
from google.api_core.iam import Policy
1838+
17851839
PATH = "/projects/%s/datasets/%s/tables/%s:setIamPolicy" % (
17861840
self.PROJECT,
17871841
self.DS_ID,
@@ -1858,6 +1912,19 @@ def test_set_iam_policy_no_mask(self):
18581912
method="POST", path=PATH, data=BODY, timeout=7.5
18591913
)
18601914

1915+
def test_set_ia_policy_updateMask_and_fields(self):
1916+
from google.api_core.iam import Policy
1917+
1918+
policy = Policy()
1919+
creds = _make_credentials()
1920+
http = object()
1921+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
1922+
1923+
with pytest.raises(ValueError, match="updateMask"):
1924+
client.set_iam_policy(
1925+
self.TABLE_REF, policy, updateMask="bindings", fields=("bindings",)
1926+
)
1927+
18611928
def test_set_iam_policy_invalid_policy(self):
18621929
from google.api_core.iam import Policy
18631930

0 commit comments

Comments
 (0)