Skip to content

Commit f59fc9a

Browse files
steffnaysteffnay
authored
feat: add support for getting and setting table IAM policy (#144)
1 parent 7a6f719 commit f59fc9a

File tree

4 files changed

+355
-0
lines changed

4 files changed

+355
-0
lines changed

‎google/cloud/bigquery/client.py

+58
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646

4747
import google.api_core.client_options
4848
import google.api_core.exceptions
49+
from google.api_core.iam import Policy
4950
from google.api_core import page_iterator
5051
import google.cloud._helpers
5152
from google.cloud import exceptions
@@ -605,6 +606,63 @@ def get_dataset(self, dataset_ref, retry=DEFAULT_RETRY, timeout=None):
605606
)
606607
return Dataset.from_api_repr(api_response)
607608

609+
def get_iam_policy(
610+
self, table, requested_policy_version=1, retry=DEFAULT_RETRY, timeout=None,
611+
):
612+
if not isinstance(table, (Table, TableReference)):
613+
raise TypeError("table must be a Table or TableReference")
614+
615+
if requested_policy_version != 1:
616+
raise ValueError("only IAM policy version 1 is supported")
617+
618+
body = {"options": {"requestedPolicyVersion": 1}}
619+
620+
path = "{}:getIamPolicy".format(table.path)
621+
622+
response = self._call_api(
623+
retry, method="POST", path=path, data=body, timeout=timeout,
624+
)
625+
626+
return Policy.from_api_repr(response)
627+
628+
def set_iam_policy(
629+
self, table, policy, updateMask=None, retry=DEFAULT_RETRY, timeout=None,
630+
):
631+
if not isinstance(table, (Table, TableReference)):
632+
raise TypeError("table must be a Table or TableReference")
633+
634+
if not isinstance(policy, (Policy)):
635+
raise TypeError("policy must be a Policy")
636+
637+
body = {"policy": policy.to_api_repr()}
638+
639+
if updateMask is not None:
640+
body["updateMask"] = updateMask
641+
642+
path = "{}:setIamPolicy".format(table.path)
643+
644+
response = self._call_api(
645+
retry, method="POST", path=path, data=body, timeout=timeout,
646+
)
647+
648+
return Policy.from_api_repr(response)
649+
650+
def test_iam_permissions(
651+
self, table, permissions, retry=DEFAULT_RETRY, timeout=None,
652+
):
653+
if not isinstance(table, (Table, TableReference)):
654+
raise TypeError("table must be a Table or TableReference")
655+
656+
body = {"permissions": permissions}
657+
658+
path = "{}:testIamPermissions".format(table.path)
659+
660+
response = self._call_api(
661+
retry, method="POST", path=path, data=body, timeout=timeout,
662+
)
663+
664+
return response
665+
608666
def get_model(self, model_ref, retry=DEFAULT_RETRY, timeout=None):
609667
"""[Beta] Fetch the model referenced by ``model_ref``.
610668

‎google/cloud/bigquery/iam.py

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
# Copyright 2020 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+
# http://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+
"""BigQuery API IAM policy definitions
15+
16+
For all allowed roles and permissions, see:
17+
18+
https://cloud.google.com/bigquery/docs/access-control
19+
"""
20+
21+
# BigQuery-specific IAM roles available for tables and views
22+
23+
BIGQUERY_DATA_EDITOR_ROLE = "roles/bigquery.dataEditor"
24+
"""When applied to a table or view, this role provides permissions to
25+
read and update data and metadata for the table or view."""
26+
27+
BIGQUERY_DATA_OWNER_ROLE = "roles/bigquery.dataOwner"
28+
"""When applied to a table or view, this role provides permissions to
29+
read and update data and metadata for the table or view, share the
30+
table/view, and delete the table/view."""
31+
32+
BIGQUERY_DATA_VIEWER_ROLE = "roles/bigquery.dataViewer"
33+
"""When applied to a table or view, this role provides permissions to
34+
read data and metadata from the table or view."""
35+
36+
BIGQUERY_METADATA_VIEWER_ROLE = "roles/bigquery.metadataViewer"
37+
"""When applied to a table or view, this role provides persmissions to
38+
read metadata from the table or view."""

‎tests/system.py

+49
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@
7171
from google.api_core.exceptions import InternalServerError
7272
from google.api_core.exceptions import ServiceUnavailable
7373
from google.api_core.exceptions import TooManyRequests
74+
from google.api_core.iam import Policy
7475
from google.cloud import bigquery
7576
from google.cloud import bigquery_v2
7677
from google.cloud.bigquery.dataset import Dataset
@@ -1407,6 +1408,54 @@ def test_copy_table(self):
14071408
got_rows = self._fetch_single_page(dest_table)
14081409
self.assertTrue(len(got_rows) > 0)
14091410

1411+
def test_get_set_iam_policy(self):
1412+
from google.cloud.bigquery.iam import BIGQUERY_DATA_VIEWER_ROLE
1413+
1414+
dataset = self.temp_dataset(_make_dataset_id("create_table"))
1415+
table_id = "test_table"
1416+
table_ref = Table(dataset.table(table_id))
1417+
self.assertFalse(_table_exists(table_ref))
1418+
1419+
table = retry_403(Config.CLIENT.create_table)(table_ref)
1420+
self.to_delete.insert(0, table)
1421+
1422+
self.assertTrue(_table_exists(table))
1423+
1424+
member = "serviceAccount:{}".format(Config.CLIENT.get_service_account_email())
1425+
BINDING = {
1426+
"role": BIGQUERY_DATA_VIEWER_ROLE,
1427+
"members": {member},
1428+
}
1429+
1430+
policy = Config.CLIENT.get_iam_policy(table)
1431+
self.assertIsInstance(policy, Policy)
1432+
self.assertEqual(policy.bindings, [])
1433+
1434+
policy.bindings.append(BINDING)
1435+
returned_policy = Config.CLIENT.set_iam_policy(table, policy)
1436+
self.assertEqual(returned_policy.bindings, policy.bindings)
1437+
1438+
def test_test_iam_permissions(self):
1439+
dataset = self.temp_dataset(_make_dataset_id("create_table"))
1440+
table_id = "test_table"
1441+
table_ref = Table(dataset.table(table_id))
1442+
self.assertFalse(_table_exists(table_ref))
1443+
1444+
table = retry_403(Config.CLIENT.create_table)(table_ref)
1445+
self.to_delete.insert(0, table)
1446+
1447+
self.assertTrue(_table_exists(table))
1448+
1449+
# Test some default permissions.
1450+
permissions = [
1451+
"bigquery.tables.get",
1452+
"bigquery.tables.getData",
1453+
"bigquery.tables.update",
1454+
]
1455+
1456+
response = Config.CLIENT.test_iam_permissions(table, [permissions])
1457+
self.assertEqual(set(response["permissions"]), set(permissions))
1458+
14101459
def test_job_cancel(self):
14111460
DATASET_ID = _make_dataset_id("job_cancel")
14121461
JOB_ID_PREFIX = "fetch_" + DATASET_ID

‎tests/unit/test_client.py

+210
Original file line numberDiff line numberDiff line change
@@ -1748,6 +1748,216 @@ def test_get_table_sets_user_agent(self):
17481748
)
17491749
self.assertIn("my-application/1.2.3", expected_user_agent)
17501750

1751+
def test_get_iam_policy(self):
1752+
from google.cloud.bigquery.iam import BIGQUERY_DATA_OWNER_ROLE
1753+
from google.cloud.bigquery.iam import BIGQUERY_DATA_EDITOR_ROLE
1754+
from google.cloud.bigquery.iam import BIGQUERY_DATA_VIEWER_ROLE
1755+
from google.api_core.iam import Policy
1756+
1757+
PATH = "/projects/{}/datasets/{}/tables/{}:getIamPolicy".format(
1758+
self.PROJECT, self.DS_ID, self.TABLE_ID,
1759+
)
1760+
BODY = {"options": {"requestedPolicyVersion": 1}}
1761+
ETAG = "CARDI"
1762+
VERSION = 1
1763+
OWNER1 = "user:[email protected]"
1764+
OWNER2 = "group:[email protected]"
1765+
EDITOR1 = "domain:google.com"
1766+
EDITOR2 = "user:[email protected]"
1767+
VIEWER1 = "serviceAccount:[email protected]"
1768+
VIEWER2 = "user:[email protected]"
1769+
RETURNED = {
1770+
"resourceId": PATH,
1771+
"etag": ETAG,
1772+
"version": VERSION,
1773+
"bindings": [
1774+
{"role": BIGQUERY_DATA_OWNER_ROLE, "members": [OWNER1, OWNER2]},
1775+
{"role": BIGQUERY_DATA_EDITOR_ROLE, "members": [EDITOR1, EDITOR2]},
1776+
{"role": BIGQUERY_DATA_VIEWER_ROLE, "members": [VIEWER1, VIEWER2]},
1777+
],
1778+
}
1779+
EXPECTED = {
1780+
binding["role"]: set(binding["members"]) for binding in RETURNED["bindings"]
1781+
}
1782+
1783+
creds = _make_credentials()
1784+
http = object()
1785+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
1786+
conn = client._connection = make_connection(RETURNED)
1787+
1788+
policy = client.get_iam_policy(self.TABLE_REF, timeout=7.5)
1789+
1790+
conn.api_request.assert_called_once_with(
1791+
method="POST", path=PATH, data=BODY, timeout=7.5
1792+
)
1793+
1794+
self.assertIsInstance(policy, Policy)
1795+
self.assertEqual(policy.etag, RETURNED["etag"])
1796+
self.assertEqual(policy.version, RETURNED["version"])
1797+
self.assertEqual(dict(policy), EXPECTED)
1798+
1799+
def test_get_iam_policy_w_invalid_table(self):
1800+
creds = _make_credentials()
1801+
http = object()
1802+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
1803+
1804+
table_resource_string = "projects/{}/datasets/{}/tables/{}".format(
1805+
self.PROJECT, self.DS_ID, self.TABLE_ID,
1806+
)
1807+
1808+
with self.assertRaises(TypeError):
1809+
client.get_iam_policy(table_resource_string)
1810+
1811+
def test_get_iam_policy_w_invalid_version(self):
1812+
creds = _make_credentials()
1813+
http = object()
1814+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
1815+
1816+
with self.assertRaises(ValueError):
1817+
client.get_iam_policy(self.TABLE_REF, requested_policy_version=2)
1818+
1819+
def test_set_iam_policy(self):
1820+
from google.cloud.bigquery.iam import BIGQUERY_DATA_OWNER_ROLE
1821+
from google.cloud.bigquery.iam import BIGQUERY_DATA_EDITOR_ROLE
1822+
from google.cloud.bigquery.iam import BIGQUERY_DATA_VIEWER_ROLE
1823+
from google.api_core.iam import Policy
1824+
1825+
PATH = "/projects/%s/datasets/%s/tables/%s:setIamPolicy" % (
1826+
self.PROJECT,
1827+
self.DS_ID,
1828+
self.TABLE_ID,
1829+
)
1830+
ETAG = "foo"
1831+
VERSION = 1
1832+
OWNER1 = "user:[email protected]"
1833+
OWNER2 = "group:[email protected]"
1834+
EDITOR1 = "domain:google.com"
1835+
EDITOR2 = "user:[email protected]"
1836+
VIEWER1 = "serviceAccount:[email protected]"
1837+
VIEWER2 = "user:[email protected]"
1838+
BINDINGS = [
1839+
{"role": BIGQUERY_DATA_OWNER_ROLE, "members": [OWNER1, OWNER2]},
1840+
{"role": BIGQUERY_DATA_EDITOR_ROLE, "members": [EDITOR1, EDITOR2]},
1841+
{"role": BIGQUERY_DATA_VIEWER_ROLE, "members": [VIEWER1, VIEWER2]},
1842+
]
1843+
MASK = "bindings,etag"
1844+
RETURNED = {"etag": ETAG, "version": VERSION, "bindings": BINDINGS}
1845+
1846+
policy = Policy()
1847+
for binding in BINDINGS:
1848+
policy[binding["role"]] = binding["members"]
1849+
1850+
BODY = {"policy": policy.to_api_repr(), "updateMask": MASK}
1851+
1852+
creds = _make_credentials()
1853+
http = object()
1854+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
1855+
conn = client._connection = make_connection(RETURNED)
1856+
1857+
returned_policy = client.set_iam_policy(
1858+
self.TABLE_REF, policy, updateMask=MASK, timeout=7.5
1859+
)
1860+
1861+
conn.api_request.assert_called_once_with(
1862+
method="POST", path=PATH, data=BODY, timeout=7.5
1863+
)
1864+
self.assertEqual(returned_policy.etag, ETAG)
1865+
self.assertEqual(returned_policy.version, VERSION)
1866+
self.assertEqual(dict(returned_policy), dict(policy))
1867+
1868+
def test_set_iam_policy_no_mask(self):
1869+
from google.api_core.iam import Policy
1870+
1871+
PATH = "/projects/%s/datasets/%s/tables/%s:setIamPolicy" % (
1872+
self.PROJECT,
1873+
self.DS_ID,
1874+
self.TABLE_ID,
1875+
)
1876+
RETURNED = {"etag": "foo", "version": 1, "bindings": []}
1877+
1878+
policy = Policy()
1879+
BODY = {"policy": policy.to_api_repr()}
1880+
1881+
creds = _make_credentials()
1882+
http = object()
1883+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
1884+
conn = client._connection = make_connection(RETURNED)
1885+
1886+
client.set_iam_policy(self.TABLE_REF, policy, timeout=7.5)
1887+
1888+
conn.api_request.assert_called_once_with(
1889+
method="POST", path=PATH, data=BODY, timeout=7.5
1890+
)
1891+
1892+
def test_set_iam_policy_invalid_policy(self):
1893+
from google.api_core.iam import Policy
1894+
1895+
policy = Policy()
1896+
invalid_policy_repr = policy.to_api_repr()
1897+
1898+
creds = _make_credentials()
1899+
http = object()
1900+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
1901+
1902+
with self.assertRaises(TypeError):
1903+
client.set_iam_policy(self.TABLE_REF, invalid_policy_repr)
1904+
1905+
def test_set_iam_policy_w_invalid_table(self):
1906+
from google.api_core.iam import Policy
1907+
1908+
policy = Policy()
1909+
1910+
creds = _make_credentials()
1911+
http = object()
1912+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
1913+
1914+
table_resource_string = "projects/%s/datasets/%s/tables/%s" % (
1915+
self.PROJECT,
1916+
self.DS_ID,
1917+
self.TABLE_ID,
1918+
)
1919+
1920+
with self.assertRaises(TypeError):
1921+
client.set_iam_policy(table_resource_string, policy)
1922+
1923+
def test_test_iam_permissions(self):
1924+
PATH = "/projects/%s/datasets/%s/tables/%s:testIamPermissions" % (
1925+
self.PROJECT,
1926+
self.DS_ID,
1927+
self.TABLE_ID,
1928+
)
1929+
1930+
PERMISSIONS = ["bigquery.tables.get", "bigquery.tables.update"]
1931+
BODY = {"permissions": PERMISSIONS}
1932+
RETURNED = {"permissions": PERMISSIONS}
1933+
1934+
creds = _make_credentials()
1935+
http = object()
1936+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
1937+
conn = client._connection = make_connection(RETURNED)
1938+
1939+
client.test_iam_permissions(self.TABLE_REF, PERMISSIONS, timeout=7.5)
1940+
1941+
conn.api_request.assert_called_once_with(
1942+
method="POST", path=PATH, data=BODY, timeout=7.5
1943+
)
1944+
1945+
def test_test_iam_permissions_w_invalid_table(self):
1946+
creds = _make_credentials()
1947+
http = object()
1948+
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
1949+
1950+
table_resource_string = "projects/%s/datasets/%s/tables/%s" % (
1951+
self.PROJECT,
1952+
self.DS_ID,
1953+
self.TABLE_ID,
1954+
)
1955+
1956+
PERMISSIONS = ["bigquery.tables.get", "bigquery.tables.update"]
1957+
1958+
with self.assertRaises(TypeError):
1959+
client.test_iam_permissions(table_resource_string, PERMISSIONS)
1960+
17511961
def test_update_dataset_w_invalid_field(self):
17521962
from google.cloud.bigquery.dataset import Dataset
17531963

0 commit comments

Comments
 (0)