Skip to content
Permalink
Browse files
feat: add support for getting and setting table IAM policy (#144)
  • Loading branch information
steffnay committed Jul 30, 2020
1 parent 7a6f719 commit f59fc9a482d9f9ae63e2b2bfc80b9a3481d09bde
Showing with 355 additions and 0 deletions.
  1. +58 −0 google/cloud/bigquery/client.py
  2. +38 −0 google/cloud/bigquery/iam.py
  3. +49 −0 tests/system.py
  4. +210 −0 tests/unit/test_client.py
@@ -46,6 +46,7 @@

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

def get_iam_policy(
self, table, requested_policy_version=1, retry=DEFAULT_RETRY, timeout=None,
):
if not isinstance(table, (Table, TableReference)):
raise TypeError("table must be a Table or TableReference")

if requested_policy_version != 1:
raise ValueError("only IAM policy version 1 is supported")

body = {"options": {"requestedPolicyVersion": 1}}

path = "{}:getIamPolicy".format(table.path)

response = self._call_api(
retry, method="POST", path=path, data=body, timeout=timeout,
)

return Policy.from_api_repr(response)

def set_iam_policy(
self, table, policy, updateMask=None, retry=DEFAULT_RETRY, timeout=None,
):
if not isinstance(table, (Table, TableReference)):
raise TypeError("table must be a Table or TableReference")

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

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

response = self._call_api(
retry, method="POST", path=path, data=body, timeout=timeout,
)

return Policy.from_api_repr(response)

def test_iam_permissions(
self, table, permissions, retry=DEFAULT_RETRY, timeout=None,
):
if not isinstance(table, (Table, TableReference)):
raise TypeError("table must be a Table or TableReference")

body = {"permissions": permissions}

path = "{}:testIamPermissions".format(table.path)

response = self._call_api(
retry, method="POST", path=path, data=body, timeout=timeout,
)

return response

def get_model(self, model_ref, retry=DEFAULT_RETRY, timeout=None):
"""[Beta] Fetch the model referenced by ``model_ref``.
@@ -0,0 +1,38 @@
# Copyright 2020 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
#
# http://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.
"""BigQuery API IAM policy definitions
For all allowed roles and permissions, see:
https://cloud.google.com/bigquery/docs/access-control
"""

# BigQuery-specific IAM roles available for tables and views

BIGQUERY_DATA_EDITOR_ROLE = "roles/bigquery.dataEditor"
"""When applied to a table or view, this role provides permissions to
read and update data and metadata for the table or view."""

BIGQUERY_DATA_OWNER_ROLE = "roles/bigquery.dataOwner"
"""When applied to a table or view, this role provides permissions to
read and update data and metadata for the table or view, share the
table/view, and delete the table/view."""

BIGQUERY_DATA_VIEWER_ROLE = "roles/bigquery.dataViewer"
"""When applied to a table or view, this role provides permissions to
read data and metadata from the table or view."""

BIGQUERY_METADATA_VIEWER_ROLE = "roles/bigquery.metadataViewer"
"""When applied to a table or view, this role provides persmissions to
read metadata from the table or view."""
@@ -71,6 +71,7 @@
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 import bigquery_v2
from google.cloud.bigquery.dataset import Dataset
@@ -1407,6 +1408,54 @@ 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 = 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"
table_ref = Table(dataset.table(table_id))
self.assertFalse(_table_exists(table_ref))

table = retry_403(Config.CLIENT.create_table)(table_ref)
self.to_delete.insert(0, table)

self.assertTrue(_table_exists(table))

# Test some default permissions.
permissions = [
"bigquery.tables.get",
"bigquery.tables.getData",
"bigquery.tables.update",
]

response = Config.CLIENT.test_iam_permissions(table, [permissions])
self.assertEqual(set(response["permissions"]), set(permissions))

def test_job_cancel(self):
DATASET_ID = _make_dataset_id("job_cancel")
JOB_ID_PREFIX = "fetch_" + DATASET_ID
@@ -1748,6 +1748,216 @@ def test_get_table_sets_user_agent(self):
)
self.assertIn("my-application/1.2.3", expected_user_agent)

def test_get_iam_policy(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/{}/datasets/{}/tables/{}:getIamPolicy".format(
self.PROJECT, self.DS_ID, self.TABLE_ID,
)
BODY = {"options": {"requestedPolicyVersion": 1}}
ETAG = "CARDI"
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"
RETURNED = {
"resourceId": PATH,
"etag": ETAG,
"version": VERSION,
"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]},
],
}
EXPECTED = {
binding["role"]: set(binding["members"]) for binding in RETURNED["bindings"]
}

creds = _make_credentials()
http = object()
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
conn = client._connection = make_connection(RETURNED)

policy = client.get_iam_policy(self.TABLE_REF, timeout=7.5)

conn.api_request.assert_called_once_with(
method="POST", path=PATH, data=BODY, timeout=7.5
)

self.assertIsInstance(policy, Policy)
self.assertEqual(policy.etag, RETURNED["etag"])
self.assertEqual(policy.version, RETURNED["version"])
self.assertEqual(dict(policy), EXPECTED)

def test_get_iam_policy_w_invalid_table(self):
creds = _make_credentials()
http = object()
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)

table_resource_string = "projects/{}/datasets/{}/tables/{}".format(
self.PROJECT, self.DS_ID, self.TABLE_ID,
)

with self.assertRaises(TypeError):
client.get_iam_policy(table_resource_string)

def test_get_iam_policy_w_invalid_version(self):
creds = _make_credentials()
http = object()
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)

with self.assertRaises(ValueError):
client.get_iam_policy(self.TABLE_REF, requested_policy_version=2)

def test_set_iam_policy(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,
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]},
]
MASK = "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": MASK}

creds = _make_credentials()
http = object()
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
conn = client._connection = make_connection(RETURNED)

returned_policy = client.set_iam_policy(
self.TABLE_REF, policy, updateMask=MASK, timeout=7.5
)

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_no_mask(self):
from google.api_core.iam import Policy

PATH = "/projects/%s/datasets/%s/tables/%s:setIamPolicy" % (
self.PROJECT,
self.DS_ID,
self.TABLE_ID,
)
RETURNED = {"etag": "foo", "version": 1, "bindings": []}

policy = Policy()
BODY = {"policy": policy.to_api_repr()}

creds = _make_credentials()
http = object()
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
conn = client._connection = make_connection(RETURNED)

client.set_iam_policy(self.TABLE_REF, policy, timeout=7.5)

conn.api_request.assert_called_once_with(
method="POST", path=PATH, data=BODY, timeout=7.5
)

def test_set_iam_policy_invalid_policy(self):
from google.api_core.iam import Policy

policy = Policy()
invalid_policy_repr = policy.to_api_repr()

creds = _make_credentials()
http = object()
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)

with self.assertRaises(TypeError):
client.set_iam_policy(self.TABLE_REF, invalid_policy_repr)

def test_set_iam_policy_w_invalid_table(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)

table_resource_string = "projects/%s/datasets/%s/tables/%s" % (
self.PROJECT,
self.DS_ID,
self.TABLE_ID,
)

with self.assertRaises(TypeError):
client.set_iam_policy(table_resource_string, policy)

def test_test_iam_permissions(self):
PATH = "/projects/%s/datasets/%s/tables/%s:testIamPermissions" % (
self.PROJECT,
self.DS_ID,
self.TABLE_ID,
)

PERMISSIONS = ["bigquery.tables.get", "bigquery.tables.update"]
BODY = {"permissions": PERMISSIONS}
RETURNED = {"permissions": PERMISSIONS}

creds = _make_credentials()
http = object()
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
conn = client._connection = make_connection(RETURNED)

client.test_iam_permissions(self.TABLE_REF, PERMISSIONS, timeout=7.5)

conn.api_request.assert_called_once_with(
method="POST", path=PATH, data=BODY, timeout=7.5
)

def test_test_iam_permissions_w_invalid_table(self):
creds = _make_credentials()
http = object()
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)

table_resource_string = "projects/%s/datasets/%s/tables/%s" % (
self.PROJECT,
self.DS_ID,
self.TABLE_ID,
)

PERMISSIONS = ["bigquery.tables.get", "bigquery.tables.update"]

with self.assertRaises(TypeError):
client.test_iam_permissions(table_resource_string, PERMISSIONS)

def test_update_dataset_w_invalid_field(self):
from google.cloud.bigquery.dataset import Dataset

0 comments on commit f59fc9a

Please sign in to comment.