feat: add support for getting and setting table IAM policy

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,70 @@ 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")
+
+        options = {}
+        body = {}
+        options["requestedPolicyVersion"] = requested_policy_version
+        body["options"] = options
+
+        path = "%s:getIamPolicy" % (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 = {}
+
+        if updateMask is not None:
+            body["updateMask"] = updateMask
+
+        body["policy"] = policy.to_api_repr()
+
+        path = "%s:setIamPolicy" % (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 = {}
+
+        body["permissions"] = permissions
+
+        path = "%s:testIamPermissions" % (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_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