feat: add support for getting and setting table IAM policy (#144)

google/cloud/bigquery/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``.
 

google/cloud/bigquery/iam.py

@@ -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."""

tests/system.py

@@ -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

tests/unit/test_client.py

@@ -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